From 28262c20df6f2945c1b1cd7b27b0faef59cb0438 Mon Sep 17 00:00:00 2001 From: Alex Nails Date: Wed, 2 Sep 2026 19:46:08 -0700 Subject: [PATCH] [CI][RFC] Replace black-jupyter with ruff-format (#37210) Co-authored-by: Alison Shao --- .../scripts/triage_kernel_helpers.py | 10 +- .../mechanical_refactor_proof_generator.py | 4 +- .../mechanical_refactor_reproduction_utils.py | 62 ++--- .../tests/proof_generator/test_infer_moves.py | 15 +- .../tests/reproduction_cli/cli_testlib.py | 8 +- .../reproduction_utils/test_add_imports.py | 15 +- .../test_extract_symbols_to_new_module.py | 6 +- .../reproduction_utils/test_move_assign.py | 16 +- .../reproduction_utils/test_move_symbol.py | 8 +- .../scripts/incident_artifact_tool.py | 7 +- .pre-commit-config.yaml | 5 +- 3rdparty/amd/tuning/benchmark_moe_rocm.py | 6 +- benchmark/asr/bench_sglang.py | 2 +- .../bench_kda_flashinfer_mtp.py | 4 +- benchmark/gsm8k/bench_sglang.py | 6 +- benchmark/hf3fs/bench_client.py | 8 +- benchmark/hf3fs/bench_storage.py | 8 +- benchmark/hf3fs/bench_zerocopy.py | 4 +- benchmark/hicache/bench_mix.py | 6 +- benchmark/hicache/bench_serving.py | 6 +- benchmark/hicache/nextqa.py | 2 +- .../kernels/bench_fused_rmsnorm_fp8_quant.py | 6 +- benchmark/kernels/deepep/deepep_utils.py | 6 +- benchmark/kernels/deepep/tuning_deepep.py | 36 +-- .../deepseek/benchmark_deepgemm_fp8_gemm.py | 3 +- .../benchmark_deepgemm_fp8_gemm_blackwell.py | 2 +- .../benchmark_deepgemm_fp8_group_gemm.py | 2 +- .../deepseek/benchmark_q8kv8_kv_gather.py | 3 +- .../elementwise/benchmark_concat_mla.py | 4 +- .../tuning_fused_moe_triton_sep.py | 4 +- .../kernels/lora_csgmv/tune_lora_csgmv.py | 16 +- benchmark/lean_kernel_sweep.py | 8 +- benchmark/mmlu/bench_sglang.py | 2 +- benchmark/ocr/bench_sglang.py | 2 +- benchmark/ocr/generate_report.py | 8 +- benchmark/prefill_only/util.py | 7 +- benchmark/scheduler/bench_token_storage.py | 10 +- docs/demo/deepseek_v4_flash.ipynb | 4 +- .../usage/llava_video/srt_example_llava_v.py | 1 - .../usage/readme_examples.py | 2 +- .../nsys_profile_tools/gputrc2graph.py | 8 +- .../engine/fastapi_engine_inference.py | 2 +- .../engine/offline_batch_inference_qwen_1m.py | 7 +- examples/runtime/qwen3_vl_reranker.py | 6 +- .../usage/modelopt_quantize_and_export.py | 2 +- .../rac_collect_traces.py | 2 +- .../rac_serve_and_eval.py | 2 +- .../test_two_router_convergence.py | 24 +- .../e2e/chat_completions/test_validation.py | 12 +- .../e2e/k8s_integration/test_discovery.py | 20 +- .../sgl-router/tests/e2e/test_chat_smoke.py | 12 +- .../tests/e2e/test_tokenize_smoke.py | 12 +- .../tests/scripts/generate_parity_fixtures.py | 2 +- .../sglang/benchmark/dspark_sps_profiler.py | 3 +- python/sglang/benchmark/offline_throughput.py | 10 +- python/sglang/benchmark/one_batch_server.py | 16 +- python/sglang/benchmark/serving.py | 52 ++--- python/sglang/cli/killall.py | 3 +- python/sglang/cli/serve.py | 2 +- python/sglang/cli/utils.py | 3 +- .../bench_amd_deterministic_allreduce.py | 22 +- .../aot/benchmark/bench_cutlass_mla.py | 6 +- .../kernels/aot/benchmark/bench_fp4_gemm.py | 12 +- .../kernels/aot/benchmark/bench_int8_gemm.py | 4 +- .../aot/python/sgl_kernel/attention.py | 18 +- .../aot/python/sgl_kernel/elementwise.py | 12 +- .../aot/python/sgl_kernel/flash_mla.py | 12 +- .../kernels/aot/python/sgl_kernel/metal.py | 4 +- .../kernels/aot/python/sgl_kernel/musa.py | 12 +- .../aot/python/sgl_kernel/scalar_type.py | 30 +-- .../aot/python/sgl_kernel/test_utils.py | 6 +- .../sgl_kernel/testing/rotary_embedding.py | 12 +- .../kernels/aot/python/sgl_kernel/top_k.py | 18 +- .../speculative/test_speculative_sampling.py | 18 +- .../aot/tests/test_custom_allreduce.py | 6 +- .../aot/tests/test_flash_attn_sparse.py | 14 +- .../sglang/kernels/aot/tests/test_flashmla.py | 4 +- .../aot/tests/test_fused_qk_norm_rope.py | 6 +- .../aot/tests/test_moe_topk_sigmoid.py | 48 ++-- .../aot/tests/test_moe_topk_softmax.py | 48 ++-- python/sglang/kernels/jit/__main__.py | 6 +- python/sglang/kernels/jit/benchmark/marker.py | 21 +- .../kernels/jit/utils/compile/loader.py | 2 +- .../kda_kernels/causal_conv3d_cat_pad_jit.py | 2 +- .../kda_kernels/qwen3x_nvfp4_gemm_sm120.py | 6 +- .../kda_kernels/residual_gate_add_jit.py | 4 +- .../kernels/ops/activation/activation.py | 6 +- .../sglang/kernels/ops/activation/softcap.py | 6 +- .../attention/cutedsl_fp8_paged_mqa_logits.py | 12 +- .../ops/attention/cutedsl_gdn_mtp_ring.py | 114 ++++----- .../kernels/ops/attention/cutedsl_kda.py | 12 +- .../kernels/ops/attention/decode_attention.py | 1 - .../kernels/ops/attention/deepseek_v4_rope.py | 18 +- .../ops/attention/dsa/dequant_k_cache.py | 6 +- .../ops/attention/dsa/index_buf_accessor.py | 6 +- .../ops/attention/dsa/tilelang_kernel.py | 79 +++---- .../ops/attention/dsa/triton_kernel.py | 6 +- .../ops/attention/dsa/triton_sparse_mla.py | 8 +- .../ops/attention/dsv4/index_buf_accessor.py | 4 +- .../ops/attention/dsv4/metadata_kernel.py | 18 +- .../ops/attention/fa4_sm120/flash_fwd.py | 10 +- .../ops/attention/fa4_sm120/scheduler.py | 6 +- .../sglang/kernels/ops/attention/fla/chunk.py | 13 +- .../ops/attention/fla/chunk_delta_h.py | 13 +- .../kernels/ops/attention/fla/chunk_fwd.py | 14 +- .../kernels/ops/attention/fla/chunk_intra.py | 28 ++- .../fla/chunk_intra_token_parallel.py | 7 +- .../kernels/ops/attention/fla/chunk_o.py | 14 +- .../kernels/ops/attention/fla/cumsum.py | 46 ++-- .../ops/attention/fla/fused_norm_gate.py | 6 +- .../ops/attention/fla/fused_recurrent.py | 16 +- .../fla/gdn_replayssm_spec_decode.py | 6 +- .../sglang/kernels/ops/attention/fla/kda.py | 12 +- .../kernels/ops/attention/fla/l2norm.py | 2 - .../ops/attention/fla/layernorm_gated.py | 7 +- .../sglang/kernels/ops/attention/fla/utils.py | 6 +- .../kernels/ops/attention/fla/wy_fast.py | 14 +- .../kernels/ops/attention/flash_mla_sm120.py | 4 +- .../linear/kda_nvidia_prefill/chunk_fwd.py | 3 +- .../fuse_k4_only_persistent.py | 3 - .../linear/kda_ptx_prefill/__init__.py | 18 +- .../ops/attention/linear/lightning_attn.py | 1 - .../kernels/ops/attention/linear/seg_la.py | 1 - .../sglang/kernels/ops/attention/metadata.py | 12 +- .../ops/attention/minimax_qknorm_rope.py | 6 +- .../minimax_sparse/decode/topk_sparse.py | 6 +- .../prefill/flash_with_topk_idx.py | 6 +- .../ops/attention/mla_kv_pack_quantize_fp8.py | 18 +- .../ops/attention/rocm_mla_decode_rope.py | 6 +- .../sglang/kernels/ops/attention/score_mod.py | 6 +- .../sparse_mla_q8kv8_prefill_sm90.py | 10 +- .../ops/communication/inkling_all_reduce.py | 6 +- python/sglang/kernels/ops/communication/mp.py | 3 +- .../attention/sana_wm_gdn_chunkwise_triton.py | 12 +- .../diffusion/attention/sana_wm_gdn_triton.py | 6 +- .../ops/diffusion/layout/usp_relayout_jit.py | 2 +- .../diffusion/modulate/scale_shift_triton.py | 6 +- .../norm/fused_residual_norm_flydsl.py | 12 +- .../kernels/ops/elementwise/elementwise.py | 12 +- .../ops/gemm/cutedsl_dsv3_fused_a_gemm.py | 6 +- ...fer_pr4266_dense_bf16_gemm_sm100_splitk.py | 4 +- .../sglang/kernels/ops/grammar/bitmask_ops.py | 6 +- .../kernels/ops/kimi_k3/sp_collective.py | 5 +- .../sglang/kernels/ops/kvcache/cache_ops.py | 48 ++-- python/sglang/kernels/ops/kvcache/hisparse.py | 6 +- .../kernels/ops/kvcache/kv_read_table.py | 8 +- .../sglang/kernels/ops/kvcache/rope_cache.py | 54 ++--- .../ops/kvcache/trtllm_mha_page_table.py | 12 +- .../sglang/kernels/ops/layernorm/__init__.py | 3 +- python/sglang/kernels/ops/layernorm/mhc.py | 12 +- python/sglang/kernels/ops/lplb/cuda_solver.py | 12 +- .../sglang/kernels/ops/lplb/shmem_budget.py | 12 +- .../kernels/ops/mamba/causal_conv1d_triton.py | 13 +- .../kernels/ops/mamba/triton_ops/mamba_ssm.py | 24 +- .../ops/mamba/triton_ops/ssd_chunk_scan.py | 10 +- .../ops/mamba/triton_ops/ssd_chunk_state.py | 2 - .../ops/mamba/triton_ops/ssd_combined.py | 12 +- .../ops/mamba/triton_ops/ssd_state_passing.py | 12 +- .../sglang/kernels/ops/moe/ep_moe_kernels.py | 18 +- .../ops/moe/fused_moe_triton_kernels.py | 54 ++--- .../ops/moe/inkling_gate_topk_renorm.py | 6 +- python/sglang/kernels/ops/moe/inkling_moe.py | 18 +- .../sglang/kernels/ops/moe/moe_fused_gate.py | 28 ++- .../sglang/kernels/ops/moe/pack_topk_ids.py | 13 +- .../sglang/kernels/ops/moe/rocm_moe_utils.py | 12 +- .../ops/moe/shuffle_rows_with_scales.py | 12 +- .../ops/moe/sigmoid_gate_topk_renorm.py | 18 +- .../kimi_k2_moe_fused_gate.py | 12 +- .../kernels/ops/quantization/fp8_kernel.py | 60 ++--- .../kernels/ops/quantization/fp8_quantize.py | 2 +- .../kernels/ops/quantization/int8_kernel.py | 18 +- .../ops/quantization/mxfp8_amd_gfx95.py | 6 +- .../ops/quantization/mxfp8_interleave_sf.py | 6 +- .../kernels/ops/sampling/murmur_hash.py | 12 +- .../ops/speculative/dspark/dspark_accept.py | 18 +- .../dspark/dspark_verify_window.py | 6 +- python/sglang/lang/backend/openai.py | 18 +- .../sglang/lang/backend/runtime_endpoint.py | 4 - python/sglang/lang/choices.py | 4 - python/sglang/lang/interpreter.py | 6 +- python/sglang/lang/ir.py | 1 - .../scripts/bench_diffusion_denoise.py | 3 +- .../test/test_flux_pipeline.py | 12 +- .../test/test_h3_request.py | 39 ++-- .../test/test_qwen_image_edit_pipeline.py | 12 +- .../test/test_qwen_image_pipeline.py | 12 +- .../test/test_zimage_pipeline.py | 12 +- .../benchmarks/bench_serving.py | 2 +- .../configs/models/adapter/ltx_2_connector.py | 1 - .../multimodal_gen/configs/models/base.py | 6 +- .../configs/models/dits/flux.py | 1 - .../configs/models/dits/krea2.py | 6 +- .../configs/models/dits/mova_audio.py | 6 +- .../configs/models/dits/mova_video.py | 6 +- .../configs/models/dits/wanvideo.py | 4 +- .../configs/pipeline_configs/base.py | 6 +- .../configs/pipeline_configs/longlive2.py | 1 - .../configs/pipeline_configs/sana.py | 1 - .../qwen_image_rollout_pipeline_mixin.py | 1 - .../zimage_rollout_pipeline_mixin.py | 1 - .../multimodal_gen/configs/sample/cosmos3.py | 3 +- .../configs/sample/minimax_h3.py | 4 +- .../configs/sample/qwenimage.py | 4 +- .../configs/sample/sampling_params.py | 15 +- .../multimodal_gen/configs/sample/wan.py | 8 +- python/sglang/multimodal_gen/registry.py | 46 ++-- .../runtime/disaggregation/scheduler_mixin.py | 2 +- .../disaggregation/transport/buffer.py | 3 +- .../base_device_communicator.py | 12 +- .../device_communicators/cpu_communicator.py | 8 +- .../device_communicators/cuda_communicator.py | 1 - .../device_communicators/pynccl.py | 7 +- .../runtime/distributed/group_coordinator.py | 63 +++-- .../runtime/distributed/parallel_groups.py | 6 +- .../runtime/distributed/parallel_state.py | 31 ++- .../runtime/distributed/utils.py | 4 +- .../runtime/entrypoints/cli/utils.py | 1 - .../entrypoints/post_training/rollout_api.py | 6 +- .../multimodal_gen/runtime/launch_server.py | 2 +- .../layers/attention/STA_configuration.py | 8 +- .../layers/attention/backends/aiter_sage.py | 2 - .../layers/attention/backends/ascend_fa.py | 2 - .../attention/backends/attention_backend.py | 2 - .../attention/backends/block_sparse_attn.py | 2 - .../backends/cube_sparse_attn/backend.py | 3 - .../layers/attention/backends/flash_attn.py | 13 +- .../layers/attention/backends/flash_attn_2.py | 1 - .../layers/attention/backends/laser_attn.py | 2 - .../attention/backends/rain_fusion_attn.py | 2 - .../layers/attention/backends/sage_attn.py | 2 - .../runtime/layers/attention/backends/sdpa.py | 4 +- .../attention/backends/sliding_tile_attn.py | 3 - .../layers/attention/backends/sol_attn.py | 1 - .../attention/backends/sparse_linear_attn.py | 10 +- .../backends/sparse_video_gen_2_attn.py | 9 +- .../attention/backends/video_sparse_attn.py | 4 +- .../backends/video_sparse_attn_h3.py | 2 +- .../layers/attention/backends/vmoba.py | 7 +- .../layers/attention/backends/xpu_backend.py | 1 - .../runtime/layers/attention/layer.py | 24 +- .../runtime/layers/attention/selector.py | 3 +- .../runtime/layers/attention/turbo_layer.py | 6 +- .../layers/kvcache/causal_attention_cache.py | 12 +- .../runtime/layers/layernorm.py | 1 - .../runtime/layers/quantization/__init__.py | 2 +- .../configs/kitchen_w4a4_config.py | 3 +- .../configs/kitchen_w4a8_config.py | 3 +- .../runtime/layers/quantization/fp8.py | 6 +- .../layers/quantization/modelopt_quant.py | 2 +- .../runtime/layers/quantization/modelslim.py | 1 - .../quantization/modelslim_mxfp4_scheme.py | 1 - .../quantization/modelslim_mxfp8_scheme.py | 1 - .../runtime/layers/quantization/mxfp4.py | 3 +- .../layers/rotary_embedding/factory.py | 6 +- .../runtime/layers/rotary_embedding/mrope.py | 12 +- .../runtime/layers/rotary_embedding/utils.py | 2 +- .../multimodal_gen/runtime/layers/usp.py | 42 ++-- .../multimodal_gen/runtime/layers/utils.py | 13 +- .../runtime/layers/visual_embedding.py | 6 +- .../component_loaders/component_loader.py | 6 +- .../loader/component_loaders/pe_loader.py | 1 - .../component_loaders/scheduler_loader.py | 6 +- .../sound_tokenizer_loader.py | 6 +- .../component_loaders/text_encoder_loader.py | 5 +- .../loader/component_loaders/vae_loader.py | 12 +- .../component_loaders/vl_encoder_loader.py | 1 - .../component_loaders/vocoder_loader.py | 6 +- .../runtime/loader/transformer_load_utils.py | 3 +- .../runtime/loader/weight_utils.py | 1 - .../runtime/managers/gpu_worker.py | 4 +- .../memory_managers/component_residency.py | 6 +- .../memory_managers/layerwise_offload.py | 6 +- .../runtime/models/bridges/mova_dual_tower.py | 10 +- .../runtime/models/dits/causal_wanvideo.py | 14 +- .../runtime/models/dits/ernie_image.py | 6 +- .../runtime/models/dits/glm_image.py | 12 +- .../runtime/models/dits/hunyuanvideo.py | 16 +- .../runtime/models/dits/krea2.py | 6 +- .../runtime/models/dits/lingbot_video_moe.py | 6 +- .../runtime/models/dits/longcat_image.py | 12 +- .../runtime/models/dits/longlive2.py | 14 +- .../runtime/models/dits/minimax_h3.py | 15 +- .../runtime/models/dits/qwen_image.py | 6 +- .../runtime/models/dits/sana.py | 1 - .../runtime/models/dits/sana_wm_components.py | 13 +- .../runtime/models/dits/wanvideo.py | 42 ++-- .../runtime/models/dits/zimage.py | 18 +- .../runtime/models/encoders/clip.py | 3 - .../runtime/models/encoders/hunyuan3d.py | 2 +- .../runtime/models/encoders/llama.py | 7 +- .../runtime/models/encoders/qwen2_5vl.py | 4 +- .../runtime/models/encoders/qwen_vl_rope.py | 3 +- .../runtime/models/encoders/t5.py | 9 - .../multimodal_gen/runtime/models/registry.py | 3 +- .../scheduling_flow_unipc_multistep.py | 16 +- .../scheduling_minimax_h3_euler_ancestral.py | 3 +- .../runtime/models/vaes/common.py | 1 - .../runtime/models/vaes/hunyuan3d_vae.py | 1 - .../runtime/models/vaes/hunyuanvae.py | 8 - .../vaes/minimax_h3_video_vae/processor.py | 1 - .../vaes/minimax_h3_video_vae/vae_vit.py | 3 +- .../runtime/models/vaes/wanvae.py | 1 - .../pipelines/comfyui_flux_pipeline.py | 20 +- .../runtime/pipelines/ernie_image.py | 7 +- .../runtime/pipelines/hunyuan_pipeline.py | 1 - .../pipelines_core/composed_pipeline_base.py | 6 +- .../pipelines_core/lora/format_adapter.py | 6 +- .../pipelines_core/lora/peft_adapter.py | 4 +- .../runtime/pipelines_core/stages/base.py | 3 +- .../runtime/pipelines_core/stages/decoding.py | 6 +- .../pipelines_core/stages/denoising.py | 6 +- .../pipelines_core/stages/denoising_dmd.py | 6 +- .../pipelines_core/stages/input_validation.py | 10 +- .../stages/latent_preparation.py | 10 +- .../stages/model_specific_stages/cosmos3.py | 9 +- .../model_specific_stages/ernie_image_pe.py | 1 - .../model_specific_stages/ltx_2/denoising.py | 4 +- .../ltx_2/latent_preparation_av.py | 8 +- .../ltx_2/text_connector.py | 3 +- .../minimax_h3/condition_noise.py | 3 +- .../minimax_h3/release_metadata.py | 2 +- .../minimax_h3/request_validation.py | 2 +- .../minimax_h3/stages/latent_preparation.py | 6 +- .../minimax_h3/stages/timestep_preparation.py | 3 +- .../minimax_h3/video_adapter.py | 3 +- .../progressive_resolution/denoising.py | 2 +- .../pipelines_core/stages/text_encoding.py | 14 +- .../runtime/platforms/__init__.py | 2 +- .../runtime/platforms/interface.py | 3 +- .../post_training/rollout_denoising_mixin.py | 1 - .../post_training/scheduler_rl_mixin.py | 18 +- .../runtime/server_args/server_args.py | 2 +- .../runtime/utils/distributed.py | 6 +- .../runtime/utils/logging_utils.py | 12 +- .../multimodal_gen/runtime/weights/source.py | 3 +- .../sglang/multimodal_gen/test/run_suite.py | 11 +- .../test/server/common/slack.py | 6 +- .../test/server/test_server_common.py | 83 +++---- .../component_accuracy/engine.py | 6 +- .../test/single_test_file/test_ar_models.py | 4 +- .../test_update_weights_from_disk.py | 42 ++-- .../sglang/multimodal_gen/test/test_utils.py | 48 ++-- .../unit/test_attention_backend_override.py | 2 +- .../test_component_quantization_admission.py | 10 +- .../test/unit/test_consistency_metrics.py | 6 +- .../test/unit/test_cosmos3_rollout.py | 2 - .../test_diffusion_bcg_tp_graph_capture.py | 110 +++++---- .../test/unit/test_host_memory_budget.py | 6 +- .../test/unit/test_image_encoder_loader.py | 126 +++++----- .../test/unit/test_layerwise_offload.py | 18 +- .../test/unit/test_lora_format_adapter.py | 6 +- .../test_minimax_h3_vae_parallel_modes.py | 8 +- .../test/unit/test_ministral3_generation.py | 5 +- .../test/unit/test_modelopt_fp4_backend.py | 5 +- .../test/unit/test_nvtx_pytorch_hooks.py | 4 +- .../test/unit/test_pi05_runtime_helpers.py | 6 +- .../test/unit/test_regional_torch_compile.py | 5 +- .../test/unit/test_rollout_api.py | 1 - .../test/unit/test_scheduler_rollout_unit.py | 6 +- .../test/unit/test_server_args.py | 1 - .../test/unit/test_text_encoder_loader.py | 37 +-- .../test/unit/test_usp_ipc_a2a_guard.py | 7 +- .../test/unit/test_usp_ring_replicated.py | 25 +- python/sglang/multimodal_gen/utils.py | 17 +- .../sglang/srt/arg_groups/attention_hook.py | 6 +- .../sglang/srt/arg_groups/cuda_graph_hook.py | 65 ++++-- .../sglang/srt/arg_groups/deepseek_v4_hook.py | 24 +- python/sglang/srt/arg_groups/hisparse_hook.py | 6 +- python/sglang/srt/arg_groups/lora_hook.py | 24 +- python/sglang/srt/arg_groups/mamba_hook.py | 18 +- python/sglang/srt/arg_groups/memory_hook.py | 6 +- python/sglang/srt/arg_groups/model_hook.py | 28 +-- .../arg_groups/model_overrides/deepseek_v2.py | 12 +- .../arg_groups/model_overrides/deepseek_v4.py | 4 +- .../srt/arg_groups/model_overrides/gpt_oss.py | 1 - python/sglang/srt/arg_groups/moe_hook.py | 26 ++- python/sglang/srt/arg_groups/overrides.py | 28 +-- python/sglang/srt/arg_groups/parallel_hook.py | 112 ++++----- .../srt/arg_groups/pd_disaggregation_hook.py | 6 +- python/sglang/srt/arg_groups/serving_hook.py | 11 +- .../sglang/srt/arg_groups/speculative_hook.py | 1 - .../sglang/srt/arg_groups/validation_hook.py | 104 +++++---- .../batch_invariant_ops.py | 12 +- python/sglang/srt/batch_overlap/operations.py | 6 +- .../srt/batch_overlap/two_batch_overlap.py | 24 +- python/sglang/srt/beam_search/beam_group.py | 6 +- python/sglang/srt/compilation/backend.py | 14 +- .../srt/compilation/compiler_interface.py | 4 +- .../srt/compilation/cuda_piecewise_backend.py | 6 +- .../srt/compilation/fix_functionalization.py | 6 +- .../sglang/srt/compilation/inductor_pass.py | 3 - python/sglang/srt/configs/bailing_hybrid.py | 7 +- python/sglang/srt/configs/deepseek_ocr.py | 7 +- python/sglang/srt/configs/deepseekvl2.py | 24 +- python/sglang/srt/configs/inkling.py | 12 +- python/sglang/srt/configs/model_config.py | 11 +- python/sglang/srt/configs/muse_glimmer.py | 2 - python/sglang/srt/configs/nano_nemotron_vl.py | 6 +- python/sglang/srt/configs/olmo3.py | 1 - python/sglang/srt/configs/qwen3_omni.py | 5 - python/sglang/srt/configs/update_config.py | 7 +- python/sglang/srt/configs/zaya.py | 7 +- python/sglang/srt/connector/base_connector.py | 1 - python/sglang/srt/connector/redis.py | 1 - .../sglang/srt/connector/remote_instance.py | 19 +- python/sglang/srt/connector/s3.py | 1 - .../sglang/srt/connector/serde/safe_serde.py | 2 - python/sglang/srt/connector/serde/serde.py | 2 - .../srt/constrained/base_grammar_backend.py | 1 - .../srt/constrained/llguidance_backend.py | 2 - .../srt/constrained/xgrammar_backend.py | 1 - .../token_aligner/smart/aux_plugins.py | 18 +- .../aligner/token_aligner/smart/executor.py | 6 +- .../comparator/output_formatter.py | 5 +- .../sglang/srt/debug_utils/cuda_coredump.py | 4 +- .../sglang/srt/debug_utils/dump_comparator.py | 2 +- python/sglang/srt/debug_utils/dump_loader.py | 6 +- python/sglang/srt/debug_utils/dumper.py | 13 +- .../debug_utils/tensor_dump_forward_hook.py | 6 +- .../disaggregation/ascend/transfer_engine.py | 1 - .../sglang/srt/disaggregation/common/conn.py | 9 +- .../disaggregation/common/staging_buffer.py | 4 +- .../disaggregation/common/staging_handler.py | 8 +- .../sglang/srt/disaggregation/common/utils.py | 4 +- python/sglang/srt/disaggregation/decode.py | 18 +- .../decode_schedule_batch_mixin.py | 13 +- .../srt/disaggregation/encoder/receiver.py | 9 +- .../srt/disaggregation/encoder/runtime.py | 6 +- .../srt/disaggregation/encoder/server.py | 7 +- .../srt/disaggregation/mooncake/conn.py | 11 +- python/sglang/srt/disaggregation/mori/conn.py | 1 - python/sglang/srt/disaggregation/nixl/conn.py | 45 ++-- python/sglang/srt/disaggregation/prefill.py | 6 +- .../custom_all_reduce_utils.py | 10 +- .../custom_all_reduce_v2.py | 6 +- .../device_communicators/hpu_communicator.py | 1 - .../device_communicators/npu_communicator.py | 1 - .../device_communicators/pymscclpp.py | 6 +- .../device_communicators/pynccl.py | 7 +- .../device_communicators/pynccl_allocator.py | 18 +- .../device_communicators/quick_all_reduce.py | 7 +- .../device_communicators/shm_broadcast.py | 2 - .../triton_symm_mem_ag.py | 18 +- .../device_communicators/xpu_communicator.py | 1 - .../sglang/srt/distributed/parallel_state.py | 110 ++++----- python/sglang/srt/distributed/utils.py | 4 +- python/sglang/srt/dllm/mixin/req.py | 6 +- python/sglang/srt/dllm/mixin/scheduler.py | 11 +- .../srt/entrypoints/anthropic/protocol.py | 11 +- .../srt/entrypoints/anthropic/serving.py | 3 +- python/sglang/srt/entrypoints/context.py | 4 - python/sglang/srt/entrypoints/grpc_bridge.py | 8 +- python/sglang/srt/entrypoints/http_server.py | 24 +- .../srt/entrypoints/openai/encoding_dsv32.py | 4 +- .../srt/entrypoints/openai/encoding_dsv4.py | 22 +- .../entrypoints/openai/serving_completions.py | 8 +- .../entrypoints/openai/serving_responses.py | 37 ++- .../srt/entrypoints/openai/tool_server.py | 3 - python/sglang/srt/entrypoints/ssl_utils.py | 2 +- python/sglang/srt/entrypoints/tool.py | 3 - python/sglang/srt/eplb/eplb_manager.py | 4 +- .../sglang/srt/eplb/eplb_simulator/reader.py | 6 +- python/sglang/srt/eplb/expert_distribution.py | 22 +- python/sglang/srt/eplb/expert_location.py | 3 +- .../srt/eplb/expert_location_dispatch.py | 2 +- .../srt/function_call/base_format_detector.py | 6 +- .../srt/function_call/deepseekv31_detector.py | 6 +- .../srt/function_call/deepseekv3_detector.py | 6 +- .../srt/function_call/glm47_moe_detector.py | 6 +- .../srt/function_call/glm4_moe_detector.py | 18 +- .../srt/function_call/internlm_detector.py | 4 +- .../srt/function_call/kimik2_detector.py | 6 +- .../function_call/kimik3_structural_tag.py | 4 +- python/sglang/srt/function_call/minimax_m2.py | 6 +- .../cpu/quantization/awq_kernels.py | 6 +- .../cpu/quantization/gptq_kernels.py | 6 +- .../gpu/quantization/gptq_kernels.py | 1 - python/sglang/srt/hardware_backend/mlx/aot.py | 3 +- .../srt/hardware_backend/mlx/model_runner.py | 9 +- .../mlx/models/muse_glimmer_mlx.py | 3 +- .../hardware_backend/mlx/moe/fused_swiglu.py | 12 +- .../musa/attention/flashattention_backend.py | 1 - .../npu/attention/ascend_backend.py | 2 - .../npu/attention/ascend_dsv4_backend.py | 22 +- .../npu/attention/ascend_gdn_backend.py | 1 - .../ascend_hybrid_linear_attn_backend.py | 10 +- .../npu/attention/ascend_kda_backend.py | 4 +- .../npu/attention/mla_preprocess.py | 6 +- .../npu/dsv4/dsv4_memory_pool.py | 6 +- .../hardware_backend/npu/memory_pool_npu.py | 2 - .../npu/quantization/gptq_kernels.py | 12 +- .../npu/quantization/linear_method_npu.py | 4 - .../xpu/graph_runner/xpu_graph_runner.py | 30 +-- .../xpu/kernels/fla/chunk_delta_h.py | 13 +- .../xpu/kernels/fla/chunk_fwd.py | 14 +- .../fla/fused_sigmoid_gating_recurrent.py | 5 +- python/sglang/srt/kv_canary/api.py | 6 +- python/sglang/srt/kv_canary/perturb/utils.py | 3 +- .../srt/kv_canary/runner/canary_manager.py | 6 +- .../srt/kv_canary/sweep_plan_builder.py | 6 +- python/sglang/srt/layers/activation.py | 4 +- python/sglang/srt/layers/amx_utils.py | 6 +- .../srt/layers/attention/aiter_backend.py | 1 - .../layers/attention/attention_registry.py | 5 +- .../layers/attention/deepseek_v4_backend.py | 30 +-- .../deepseek_v4_backend_hip_radix.py | 30 +-- .../srt/layers/attention/dsa/dsa_indexer.py | 42 ++-- .../layers/attention/dsa/dsa_topk_backend.py | 10 +- .../layers/attention/dsa/kpool_fp8_index.py | 18 +- .../sglang/srt/layers/attention/dsa/utils.py | 6 +- .../srt/layers/attention/dsa_backend.py | 30 ++- .../srt/layers/attention/dsv4/compress_hip.py | 6 +- .../srt/layers/attention/dsv4/indexer.py | 6 +- .../srt/layers/attention/dsv4/metadata.py | 12 +- .../attention/dsv4/sparse_prefill_utils.py | 12 +- .../attention/flashattention_backend.py | 33 ++- .../layers/attention/flashinfer_backend.py | 42 ++-- .../attention/flashinfer_mla_backend.py | 6 +- .../srt/layers/attention/flashmla_backend.py | 6 +- .../attention/hybrid_linear_attn_backend.py | 34 +-- .../layers/attention/linear/gdn_backend.py | 6 +- .../linear/kernels/gdn_flashinfer.py | 7 +- .../srt/layers/attention/mamba/mamba.py | 28 +-- .../layers/attention/mamba/mamba2_metadata.py | 1 - .../attention/mamba/mixer2_rms_norm_gated.py | 6 +- .../srt/layers/attention/minicpm/backend.py | 9 +- .../attention/minimax_sparse_backend.py | 5 +- .../attention/minimax_sparse_ops/msa.py | 12 +- .../naive/flash_with_topk_idx.py | 6 +- .../tests/test_flash_with_topk_idx.py | 24 +- .../tests/test_msa_fp8_parity.py | 2 +- .../tests/test_sparse_gqa.py | 18 +- .../srt/layers/attention/tbo_backend.py | 6 +- .../srt/layers/attention/triton_backend.py | 6 +- .../layers/attention/trtllm_mha_backend.py | 8 +- .../layers/attention/trtllm_mla_backend.py | 12 +- python/sglang/srt/layers/attention/vision.py | 9 +- .../srt/layers/attention/wave_backend.py | 7 +- .../srt/layers/attention/xpu_backend.py | 34 ++- python/sglang/srt/layers/communicator.py | 6 +- .../sglang/srt/layers/communicator_dsa_cp.py | 6 +- python/sglang/srt/layers/cp/base.py | 3 +- python/sglang/srt/layers/cp/zigzag.py | 6 +- .../layers/deep_gemm_wrapper/entrypoint.py | 1 - python/sglang/srt/layers/dp_attention.py | 13 +- python/sglang/srt/layers/linear.py | 30 +-- python/sglang/srt/layers/logits_processor.py | 2 +- python/sglang/srt/layers/logprob_processor.py | 12 +- python/sglang/srt/layers/moe/cutlass_moe.py | 12 +- python/sglang/srt/layers/moe/ep_moe/layer.py | 24 +- .../srt/layers/moe/flashinfer_cutedsl_moe.py | 81 ++++--- .../moe/fused_moe_triton/fused_marlin_moe.py | 18 +- .../fused_moe_triton/triton_kernels_moe.py | 24 +- python/sglang/srt/layers/moe/hash_topk.py | 6 +- python/sglang/srt/layers/moe/kt_ep_wrapper.py | 6 +- python/sglang/srt/layers/moe/mega_moe_sm90.py | 3 +- .../sglang/srt/layers/moe/moe_runner/aiter.py | 6 +- .../sglang/srt/layers/moe/moe_runner/base.py | 24 +- .../srt/layers/moe/moe_runner/deep_gemm.py | 18 +- .../moe/moe_runner/flashinfer_cutedsl.py | 18 +- .../moe/moe_runner/flashinfer_cutlass.py | 30 +-- .../moe/moe_runner/flashinfer_trtllm.py | 24 +- .../srt/layers/moe/moe_runner/hpc_ops.py | 12 +- .../srt/layers/moe/moe_runner/triton.py | 3 - .../layers/moe/moe_runner/triton_kernels.py | 18 +- .../moe/moe_runner/triton_utils/fused_moe.py | 18 +- .../srt/layers/moe/token_dispatcher/base.py | 8 - .../srt/layers/moe/token_dispatcher/moriep.py | 2 - .../layers/moe/token_dispatcher/standard.py | 5 +- python/sglang/srt/layers/moe/topk.py | 39 ++-- python/sglang/srt/layers/moe/utils.py | 4 - python/sglang/srt/layers/n_gram_embedding.py | 13 +- python/sglang/srt/layers/parameter.py | 12 +- .../sglang/srt/layers/quantization/awq/awq.py | 1 - .../srt/layers/quantization/base_config.py | 6 +- .../srt/layers/quantization/bitsandbytes.py | 6 +- .../compressed_tensors_w4a4_mxint4_moe.py | 24 +- .../compressed_tensors_w4a4_nvfp4_moe.py | 7 +- .../compressed_tensors_w4a8_fp8_moe.py | 10 +- .../compressed_tensors_w4a8_int8_moe.py | 1 - .../compressed_tensors_w8a8_fp8_moe.py | 7 +- .../schemes/compressed_tensors_w8a8_int8.py | 2 - .../compressed_tensors_w8a8_int8_moe.py | 1 - .../schemes/compressed_tensors_wNa16_moe.py | 14 +- python/sglang/srt/layers/quantization/fp8.py | 31 ++- .../srt/layers/quantization/fp8_utils.py | 12 +- python/sglang/srt/layers/quantization/gguf.py | 6 +- .../srt/layers/quantization/gptq/gptq.py | 3 +- .../srt/layers/quantization/kv_cache.py | 4 +- .../srt/layers/quantization/kvfp4_tensor.py | 6 +- .../srt/layers/quantization/modelopt_quant.py | 27 ++- .../quantization/modelslim/modelslim.py | 8 +- .../modelslim/schemes/modelslim_mxfp8.py | 1 - .../modelslim/schemes/modelslim_w4a4_int4.py | 1 - .../modelslim/schemes/modelslim_w8a8_int8.py | 1 - .../srt/layers/quantization/moe_wna16.py | 7 +- .../sglang/srt/layers/quantization/mxfp4.py | 22 +- .../layers/quantization/mxfp4_marlin_moe.py | 2 +- .../srt/layers/quantization/mxfp4_tensor.py | 6 +- .../sglang/srt/layers/quantization/petit.py | 2 +- .../srt/layers/quantization/petit_utils.py | 2 +- .../srt/layers/quantization/quark/quark.py | 3 - .../quark/schemes/quark_w4a4_mxfp4.py | 6 +- .../quark/schemes/quark_w4a4_mxfp4_moe.py | 8 +- .../quark/schemes/quark_w8a8_fp8.py | 1 - .../quark/schemes/quark_w8a8_fp8_moe.py | 7 +- .../srt/layers/quantization/quark/weights.py | 6 +- .../layers/quantization/quark_int4fp8_moe.py | 12 +- .../sglang/srt/layers/quantization/unquant.py | 9 +- .../sglang/srt/layers/quantization/utils.py | 34 ++- .../srt/layers/quantization/w8a8_fp8.py | 1 - .../srt/layers/quantization/w8a8_int8.py | 1 - .../srt/layers/rotary_embedding/base.py | 31 ++- .../srt/layers/rotary_embedding/factory.py | 6 +- .../srt/layers/rotary_embedding/mrope.py | 18 +- .../layers/rotary_embedding/rope_variant.py | 14 +- python/sglang/srt/layers/sampler.py | 24 +- python/sglang/srt/layers/utils/cp_utils.py | 6 +- .../srt/layers/vocab_parallel_embedding.py | 10 +- .../sglang/srt/lora/backend/ascend_backend.py | 12 +- .../sglang/srt/lora/backend/base_backend.py | 12 +- .../srt/lora/backend/chunked_backend.py | 6 +- .../sglang/srt/lora/backend/torch_backend.py | 12 +- .../sglang/srt/lora/backend/triton_backend.py | 6 +- python/sglang/srt/lora/layers.py | 6 +- python/sglang/srt/lora/lora_drainer.py | 12 +- python/sglang/srt/lora/lora_manager.py | 40 ++-- .../sglang/srt/lora/lora_moe_runner_marlin.py | 6 +- python/sglang/srt/lora/lora_registry.py | 12 +- .../srt/lora/marlin_lora_temp/moe_runner.py | 3 +- python/sglang/srt/lora/mem_pool.py | 12 +- .../srt/lora/torch_ops/graph_lora_ops.py | 3 - python/sglang/srt/lora/torch_ops/lora_ops.py | 3 - .../srt/lora/trtllm_lora_temp/__init__.py | 11 +- .../lora/trtllm_lora_temp/lora_dispatch.py | 12 +- .../srt/lora/trtllm_lora_temp/lora_layer.py | 18 +- .../srt/lora/trtllm_lora_temp/moe_overlap.py | 12 +- .../srt/lora/trtllm_lora_temp/sgl_fp8_moe.py | 6 +- .../sglang/srt/managers/cache_controller.py | 12 +- .../srt/managers/hisparse_coordinator.py | 12 +- .../srt/managers/multi_tokenizer_mixin.py | 6 +- python/sglang/srt/managers/overlap_utils.py | 2 - python/sglang/srt/managers/prefill_delayer.py | 12 +- python/sglang/srt/managers/schedule_batch.py | 6 +- python/sglang/srt/managers/schedule_policy.py | 10 +- python/sglang/srt/managers/scheduler.py | 33 ++- .../new_token_ratio_tracker.py | 4 +- .../scheduler_components/weight_updater.py | 12 +- .../srt/managers/tokenizer_control_mixin.py | 84 +++---- .../sglang/srt/managers/tokenizer_manager.py | 31 +-- python/sglang/srt/managers/tp_worker.py | 6 +- python/sglang/srt/managers/utils.py | 2 +- .../srt/mem_cache/allocator/hisparse.py | 25 +- .../sglang/srt/mem_cache/allocator/paged.py | 6 +- python/sglang/srt/mem_cache/chunk_cache.py | 6 +- python/sglang/srt/mem_cache/common.py | 24 +- .../mem_cache/deepseek_v4_compress_state.py | 6 +- .../srt/mem_cache/deepseek_v4_memory_pool.py | 14 +- .../srt/mem_cache/dsa_cache_layer_split.py | 6 +- .../mem_cache/embedding_cache_controller.py | 2 +- .../sglang/srt/mem_cache/hicache_storage.py | 1 - python/sglang/srt/mem_cache/hiradix_cache.py | 15 +- .../sglang/srt/mem_cache/kv_cache_builder.py | 1 - .../srt/mem_cache/kv_cache_configurator.py | 36 +-- .../srt/mem_cache/kv_index_translator.py | 3 +- python/sglang/srt/mem_cache/kv_vmm_backing.py | 6 +- .../sglang/srt/mem_cache/mamba_radix_cache.py | 169 +++++++------- python/sglang/srt/mem_cache/memory_pool.py | 62 ++--- .../srt/mem_cache/multi_ended_allocator.py | 18 +- python/sglang/srt/mem_cache/pool_host/base.py | 6 +- .../sglang/srt/mem_cache/pool_host/common.py | 15 +- .../sglang/srt/mem_cache/pool_host/mamba.py | 6 +- .../srt/mem_cache/pure_swa_radix_cache.py | 6 +- python/sglang/srt/mem_cache/radix_cache.py | 19 +- .../sglang/srt/mem_cache/radix_cache_cpp.py | 6 +- .../srt/mem_cache/rust_tree_core/adapter.py | 12 +- .../aibrix_kvcache/aibrix_kvcache_storage.py | 12 +- .../srt/mem_cache/storage/eic/test_unit.py | 18 +- .../mem_cache/storage/flexkv/flexkv_comm.py | 6 +- .../storage/flexkv/flexkv_connector.py | 16 +- .../mem_cache/storage/hf3fs/storage_hf3fs.py | 2 +- .../storage/mooncake_store/mooncake_store.py | 1 - .../mooncake_store/test_mooncake_store.py | 23 +- .../mem_cache/storage/simm/hicache_simm.py | 11 +- .../srt/mem_cache/storage/simm/test_simm.py | 21 +- .../srt/mem_cache/storage/umbp/umbp_store.py | 3 +- .../sglang/srt/mem_cache/swa_radix_cache.py | 163 +++++++------ .../components/full_component.py | 10 +- .../components/mamba_component.py | 12 +- .../unified_cache/components/swa_component.py | 46 ++-- .../components/tree_component.py | 19 +- .../unified_cache/storage_attachment.py | 1 - .../unified_cache/unified_tree_core.py | 4 +- .../unified_tree_core_interface.py | 6 +- .../srt/mem_cache/unified_memory_pool.py | 84 +++---- .../srt/mem_cache/unified_radix_cache.py | 34 +-- .../srt/model_executor/cpu_graph_runner.py | 42 ++-- .../srt/model_executor/forward_batch_info.py | 6 +- .../sglang/srt/model_executor/hook_manager.py | 5 +- .../srt/model_executor/input_buffers.py | 13 +- .../sglang/srt/model_executor/model_runner.py | 9 +- .../cuda_graph_setup.py | 18 +- .../load_model_utils.py | 6 +- .../ngram_embedding_manager.py | 6 +- .../startup_weight_load.py | 3 +- .../weight_exporter.py | 24 +- .../model_runner_components/weight_updater.py | 6 +- .../srt/model_executor/pool_configurator.py | 18 +- .../srt/model_executor/runner/base_runner.py | 12 +- .../runner/decode_cuda_graph_runner.py | 13 +- .../runner/flashinfer_autotune.py | 14 +- .../runner/prefill_cuda_graph_runner.py | 10 +- .../breakable_cuda_graph_backend.py | 13 +- .../runner_backend/cuda_graph_dedup_mixin.py | 6 +- .../breakable_cuda_graph.py | 6 +- .../breakable_cuda_graph/cuda_utils.py | 3 +- .../sglang/srt/model_loader/kimi_k3_gguf.py | 5 +- python/sglang/srt/model_loader/loader.py | 10 +- .../sglang/srt/model_loader/weight_utils.py | 16 +- python/sglang/srt/models/afmoe.py | 6 - python/sglang/srt/models/apertus.py | 2 +- python/sglang/srt/models/baichuan.py | 6 +- python/sglang/srt/models/bailing_moe.py | 5 +- .../sglang/srt/models/bailing_moe_linear.py | 31 +-- python/sglang/srt/models/bailing_moe_v3.py | 8 +- python/sglang/srt/models/bert.py | 12 - python/sglang/srt/models/clip.py | 4 - python/sglang/srt/models/deepseek.py | 9 +- .../attention_forward_methods/forward_mha.py | 6 +- .../forward_mha_rocm.py | 1 - .../forward_mla_fused_rope_cpu.py | 17 +- .../forward_mla_fused_rope_rocm.py | 1 - .../forward_mla_rocm.py | 1 - .../sglang/srt/models/deepseek_janus_pro.py | 14 +- python/sglang/srt/models/deepseek_nextn.py | 1 - python/sglang/srt/models/deepseek_ocr.py | 12 +- python/sglang/srt/models/deepseek_v2.py | 38 ++- python/sglang/srt/models/deepseek_v4.py | 25 +- .../sglang/srt/models/deepseek_v4_dspark.py | 18 +- python/sglang/srt/models/deepseek_v4_nextn.py | 1 - python/sglang/srt/models/deepseek_vl2.py | 1 - python/sglang/srt/models/dflash.py | 3 +- .../models/dots3_common/dots_omni_audio.py | 3 +- .../srt/models/dots3_common/modeling.py | 54 +++-- python/sglang/srt/models/dspark.py | 6 - python/sglang/srt/models/ernie45_moe_vl.py | 1 - python/sglang/srt/models/ernie45_vl.py | 11 +- python/sglang/srt/models/exaone.py | 3 +- python/sglang/srt/models/exaone4.py | 3 +- python/sglang/srt/models/exaone_moe.py | 3 +- python/sglang/srt/models/falcon_h1.py | 15 +- python/sglang/srt/models/gemma2_reward.py | 6 +- python/sglang/srt/models/gemma4_causal.py | 6 +- python/sglang/srt/models/gemma4_vision.py | 1 - python/sglang/srt/models/glm4_moe.py | 20 +- python/sglang/srt/models/glm4v_moe.py | 6 +- python/sglang/srt/models/gpt2.py | 5 - python/sglang/srt/models/gpt_bigcode.py | 4 - python/sglang/srt/models/gpt_j.py | 5 - python/sglang/srt/models/gpt_oss.py | 9 +- python/sglang/srt/models/granite.py | 3 +- python/sglang/srt/models/granitemoe.py | 5 - python/sglang/srt/models/granitemoehybrid.py | 1 - python/sglang/srt/models/hunyuan.py | 10 +- python/sglang/srt/models/hunyuan_v3_nextn.py | 2 - python/sglang/srt/models/idefics2.py | 4 - .../srt/models/inkling_common/dense_mlp.py | 12 +- .../models/inkling_common/kernels/sconv.py | 12 +- .../sglang/srt/models/inkling_common/moe.py | 10 +- .../inkling_common/quantization/quant.py | 6 +- python/sglang/srt/models/internlm2.py | 3 +- python/sglang/srt/models/interns1.py | 1 - python/sglang/srt/models/interns1pro.py | 17 +- python/sglang/srt/models/internvl.py | 1 - python/sglang/srt/models/iquest_loopcoder.py | 6 +- python/sglang/srt/models/kimi_k25.py | 21 +- python/sglang/srt/models/kimi_linear.py | 6 +- python/sglang/srt/models/kimi_vl.py | 3 +- python/sglang/srt/models/kimi_vl_moonvit.py | 17 +- python/sglang/srt/models/llada2.py | 5 +- python/sglang/srt/models/llama.py | 5 +- python/sglang/srt/models/llama4.py | 2 - .../sglang/srt/models/llama_classification.py | 6 +- python/sglang/srt/models/llama_embedding.py | 6 +- python/sglang/srt/models/llama_reward.py | 12 +- python/sglang/srt/models/longcat_flash.py | 5 +- .../sglang/srt/models/longcat_flash_nextn.py | 2 - python/sglang/srt/models/midashenglm.py | 12 +- python/sglang/srt/models/mimo_mtp.py | 1 - python/sglang/srt/models/mimo_v2.py | 30 ++- python/sglang/srt/models/mimo_v2_nextn.py | 2 - python/sglang/srt/models/minicpm.py | 3 +- python/sglang/srt/models/minicpm3.py | 4 +- python/sglang/srt/models/minicpmo.py | 14 +- python/sglang/srt/models/minicpmv.py | 21 +- python/sglang/srt/models/minimax_m2.py | 24 +- python/sglang/srt/models/minimax_m3.py | 12 +- python/sglang/srt/models/minimax_m3_vl.py | 6 +- python/sglang/srt/models/minimax_vl_common.py | 36 +-- python/sglang/srt/models/mistral_eagle.py | 6 +- python/sglang/srt/models/mixtral.py | 1 - python/sglang/srt/models/mllama.py | 2 - python/sglang/srt/models/mllama4.py | 6 - python/sglang/srt/models/moss_vl.py | 1 - python/sglang/srt/models/muse_glimmer.py | 2 - python/sglang/srt/models/nemotron_h_mtp.py | 12 +- python/sglang/srt/models/nemotron_nas.py | 1 - python/sglang/srt/models/nvila.py | 6 +- python/sglang/srt/models/olmo.py | 1 - python/sglang/srt/models/olmo2.py | 1 - python/sglang/srt/models/olmoe.py | 4 - python/sglang/srt/models/opt.py | 7 +- python/sglang/srt/models/orion.py | 3 +- python/sglang/srt/models/paddleocr_vl.py | 7 - python/sglang/srt/models/parakeet.py | 2 +- python/sglang/srt/models/persimmon.py | 5 - python/sglang/srt/models/phi.py | 4 - python/sglang/srt/models/phi3_small.py | 10 +- python/sglang/srt/models/phi4mm.py | 16 +- python/sglang/srt/models/phi4mm_audio.py | 35 ++- python/sglang/srt/models/phi4mm_utils.py | 9 +- python/sglang/srt/models/phimoe.py | 5 - python/sglang/srt/models/qwen.py | 3 +- python/sglang/srt/models/qwen2.py | 5 +- python/sglang/srt/models/qwen2_5_vl.py | 4 - .../sglang/srt/models/qwen2_classification.py | 8 +- python/sglang/srt/models/qwen2_vl.py | 6 - python/sglang/srt/models/qwen3_5.py | 6 +- python/sglang/srt/models/qwen3_5_mtp.py | 1 - python/sglang/srt/models/qwen3_next.py | 10 +- python/sglang/srt/models/qwen3_next_mtp.py | 1 - python/sglang/srt/models/qwen3_omni_moe.py | 1 - python/sglang/srt/models/qwen3_vl.py | 19 +- python/sglang/srt/models/radio.py | 18 +- python/sglang/srt/models/registry.py | 12 +- python/sglang/srt/models/roberta.py | 8 +- python/sglang/srt/models/siglip.py | 4 - python/sglang/srt/models/solar.py | 12 +- python/sglang/srt/models/starcoder2.py | 5 - python/sglang/srt/models/step3_vl.py | 5 +- python/sglang/srt/models/step3p5_mtp.py | 1 - python/sglang/srt/models/step3p7.py | 1 - .../sglang/srt/models/torch_native_llama.py | 3 +- python/sglang/srt/models/transformers.py | 9 +- python/sglang/srt/models/unlimited_ocr.py | 4 +- python/sglang/srt/models/utils.py | 6 +- python/sglang/srt/models/whisper.py | 9 +- python/sglang/srt/models/xverse.py | 3 +- python/sglang/srt/models/xverse_moe.py | 9 +- python/sglang/srt/models/zaya.py | 24 +- .../sglang/srt/multimodal/cache/identity.py | 3 +- .../sglang/srt/multimodal/evs/evs_module.py | 12 +- .../srt/multimodal/evs/evs_processor.py | 2 +- .../srt/multimodal/media_artifacts/base.py | 4 +- python/sglang/srt/multimodal/mm_utils.py | 5 +- .../multimodal/processors/base_processor.py | 6 +- .../srt/multimodal/processors/ernie45_vl.py | 12 +- .../srt/multimodal/processors/internvl.py | 2 +- .../srt/multimodal/processors/mimo_audio.py | 12 +- .../srt/multimodal/processors/mimo_v2.py | 44 ++-- .../srt/multimodal/processors/mimo_v2_asr.py | 13 +- .../srt/multimodal/processors/minicpm.py | 6 +- .../srt/multimodal/processors/moss_vl.py | 7 +- .../srt/multimodal/processors/qwen3_asr.py | 5 +- .../srt/multimodal/processors/qwen_audio.py | 6 +- .../srt/multimodal/processors/qwen_vl.py | 10 +- .../srt/multimodal/processors/step3_vl.py | 1 - .../multimodal/processors/unlimited_ocr.py | 3 +- .../srt/multimodal/transport/memory_pool.py | 1 - .../srt/multimodal/vit_cuda_graph_runner.py | 6 +- .../srt/multiplex/multiplexing_mixin.py | 1 - python/sglang/srt/multiplex/pdmux_context.py | 7 +- .../srt/observability/metrics_collector.py | 1 - .../srt/observability/req_time_stats.py | 26 ++- .../srt/parser/code_completion_parser.py | 6 +- python/sglang/srt/parser/conversation.py | 6 +- .../sglang/srt/parser/template_detection.py | 110 +++++---- python/sglang/srt/parser/template_manager.py | 12 +- python/sglang/srt/plugins/hook_registry.py | 4 +- python/sglang/srt/rust_server/server.py | 10 +- .../srt/sampling/custom_logit_processor.py | 6 +- .../srt/sampling/penaltylib/min_new_tokens.py | 4 +- python/sglang/srt/sampling/sampling_params.py | 5 +- python/sglang/srt/server_args.py | 6 +- python/sglang/srt/speculative/dflash_utils.py | 8 +- .../srt/speculative/dflash_worker_v2.py | 4 +- .../dspark_components/dspark_observability.py | 2 - .../dspark_components/dspark_planner.py | 7 +- .../dspark_components/dspark_sps.py | 1 - .../dspark_components/dspark_verify.py | 2 - .../dspark_components/dspark_worker_v2.py | 1 - .../sglang/srt/speculative/eagle_worker_v2.py | 12 +- .../speculative/frozen_kv_mtp_worker_v2.py | 13 +- .../multi_layer_eagle_worker_v2.py | 12 +- python/sglang/srt/speculative/ngram_worker.py | 6 +- python/sglang/srt/speculative/spec_info.py | 6 +- python/sglang/srt/speculative/spec_utils.py | 6 +- .../srt/speculative/standalone_worker_v2.py | 1 - .../srt/tokenizer/tiktoken_tokenizer.py | 4 +- python/sglang/srt/utils/bench_utils.py | 4 +- python/sglang/srt/utils/common.py | 49 ++-- python/sglang/srt/utils/custom_op.py | 18 +- python/sglang/srt/utils/flatten.py | 6 +- .../utils/hf_transformers/mistral_utils.py | 6 +- .../srt/utils/hf_transformers_patches.py | 4 +- python/sglang/srt/utils/numa_utils.py | 6 +- python/sglang/srt/utils/offloader.py | 49 ++-- python/sglang/srt/utils/patch_tokenizer.py | 12 +- python/sglang/srt/utils/patch_torch.py | 6 +- python/sglang/srt/utils/phase_checker.py | 3 +- python/sglang/srt/utils/profile_utils.py | 6 +- .../srt/utils/rank_consensus_checker.py | 2 +- python/sglang/srt/utils/rpd_utils.py | 8 +- python/sglang/srt/utils/watchdog.py | 6 +- python/sglang/srt/utils/weight_checker.py | 12 +- .../srt/utils/weight_checker_comparator.py | 6 +- python/sglang/srt/weight_cache/daemon.py | 2 +- python/sglang/test/accuracy_test_runner.py | 4 +- .../test/ascend/disaggregation_utils.py | 2 +- .../test/ascend/e2e/gen_dataset_fixed_len.py | 6 +- .../ascend/e2e/test_npu_accuracy_utils.py | 7 +- .../ascend/e2e/test_npu_multi_node_utils.py | 8 +- .../ascend/e2e/test_npu_performance_utils.py | 22 +- .../sglang/test/ascend/gsm8k_ascend_mixin.py | 4 +- .../sglang/test/ascend/test_embedding_base.py | 6 +- python/sglang/test/ascend/test_mmlu.py | 1 - .../sglang/test/cache_consistency_jitter.py | 10 +- python/sglang/test/ci/ci_stress_utils.py | 4 +- python/sglang/test/ci/ci_utils.py | 6 +- python/sglang/test/cpu_test_utils.py | 6 +- .../test/kernels/kv_canary/_canary_helpers.py | 16 +- .../test/kernels/kv_canary/_differential.py | 21 +- .../test/kernels/kv_canary/_invariants.py | 102 ++++---- .../attention_methods/dsa_attention.py | 3 +- .../attention_methods/dsv4_attention.py | 37 +-- .../speculative_draft_extend_runner.py | 5 +- .../runner_modes/speculative_draft_runner.py | 30 +-- .../speculative_target_verify_runner.py | 18 +- .../sglang/test/kits/ebnf_constrained_kit.py | 1 - python/sglang/test/kits/eval_accuracy_kit.py | 14 +- python/sglang/test/kits/fwd_occupancy_kit.py | 4 +- .../sglang/test/kits/json_constrained_kit.py | 1 - python/sglang/test/kits/matched_stop_kit.py | 24 +- python/sglang/test/kits/mmmu_vlm_kit.py | 4 +- .../test/kits/prefix_cache_branching_kit.py | 12 +- .../sglang/test/kits/regex_constrained_kit.py | 1 - .../sglang/test/kits/streaming_session_kit.py | 6 +- python/sglang/test/kl_multiturn_utils.py | 30 +-- python/sglang/test/kl_test_utils.py | 12 +- python/sglang/test/kv_canary/pp_fixture.py | 1 - python/sglang/test/lora_utils.py | 11 +- python/sglang/test/performance_test_runner.py | 4 +- python/sglang/test/runners.py | 6 +- .../background_http_poster.py | 1 - .../test/scripted_runtime/context/api.py | 13 +- .../context/kv_pool_exhauster.py | 7 +- .../context/lock_ref_exhauster.py | 1 - .../test/scripted_runtime/http_server.py | 1 - .../sglang/test/scripted_runtime/io_struct.py | 3 - .../test/scripted_runtime/scheduler_hook.py | 1 - .../sglang/test/scripted_runtime/test_case.py | 1 - .../scripted_runtime/tokenizer_recv_proxy.py | 1 - .../test/scripted_runtime_chunked_helpers.py | 12 +- python/sglang/test/send_one.py | 6 +- .../server_fixtures/disaggregation_fixture.py | 2 +- .../hybrid_attn_backend_fixture.py | 1 - .../test/server_fixtures/pcg_spec_fixture.py | 6 +- .../streaming_session_fixture.py | 7 +- .../test/simple_eval_mixed_prefix_gsm8k.py | 1 - python/sglang/test/test_deepep_utils.py | 6 +- python/sglang/test/test_deterministic.py | 12 +- python/sglang/test/test_marlin_utils.py | 9 +- python/sglang/test/test_programs.py | 22 +- python/sglang/test/test_utils.py | 1 - python/sglang/test/tool_call_test_runner.py | 12 +- python/sglang/test/vlm_utils.py | 38 +-- .../test/xpu/simple_eval_gsm8k_xpu_mixin.py | 4 +- rust/sglang-mm/bench/bench_parity.py | 2 +- rust/sglang-mm/tests/test_resize_parity.py | 2 +- scripts/ci/test_list_stage_models.py | 5 +- scripts/ci/update_est_time.py | 5 +- scripts/ci/utils/ci_coverage_report.py | 6 +- .../diffusion/compute_diffusion_partitions.py | 12 +- .../diffusion/generate_diffusion_dashboard.py | 2 +- scripts/ci/utils/diffusion/run_comparison.py | 12 +- .../diffusion/verify_diffusion_coverage.py | 3 +- scripts/ci/utils/publish_traces.py | 1 - scripts/ci/utils/query_job_status.py | 2 +- scripts/ci/utils/slash_command_handler.py | 7 +- scripts/ci/utils/xpu_job_monitor.py | 2 +- scripts/ci_monitor/ci_failures_analysis.py | 6 +- scripts/code_sync/copy_to_oss.py | 11 +- scripts/code_sync/utils.py | 7 +- scripts/convert_otel_2_perfetto.py | 2 +- scripts/lint/check_registered_tests.py | 6 +- scripts/playground/router/test_tree.py | 2 +- scripts/release/update_others_whl_index.py | 2 +- .../python/src/sglang_router/mini_lb.py | 2 - .../python/src/sglang_router/router_args.py | 1 - .../python/tests/test_startup_sequence.py | 3 +- .../chat_completions/test_enable_thinking.py | 12 +- .../chat_completions/test_function_calling.py | 134 +++++------ .../chat_completions/test_openai_server.py | 30 +-- .../chat_completions/test_validation.py | 12 +- .../e2e_test/embeddings/test_correctness.py | 12 +- .../e2e_test/infra/gpu_monitor.py | 12 +- .../k8s_integration/test_cross_namespace.py | 12 +- .../k8s_integration/test_lifecycle.py | 40 ++-- .../k8s_integration/test_multi_model.py | 6 +- .../k8s_integration/test_pd_type_change.py | 18 +- .../k8s_integration/test_reconciliation.py | 18 +- .../responses/test_streaming_events.py | 48 ++-- .../e2e_test/responses/test_tools_call.py | 96 ++++---- .../e2e_test/router/test_mmlu.py | 12 +- .../e2e_test/router/test_pd_mmlu.py | 6 +- .../e2e_test/router/test_worker_api.py | 6 +- .../8-gpu-models/test_deepseek_v3_basic.py | 4 +- .../8-gpu-models/test_dsa_models_basic.py | 16 +- test/manual/ascend/disaggregation_utils.py | 2 +- .../manual/ascend/test_ascend_deepseek_mtp.py | 1 - .../attention/test_trtllm_mla_backend.py | 14 +- .../test_beam_search_perf_sweep.py | 6 +- .../chunked_prefill/test_scripted_abort.py | 28 +-- .../test_scripted_chunk_size.py | 6 +- .../test_scripted_http_smoke.py | 6 +- .../test_scripted_hybrid_swa.py | 12 +- .../test_scripted_invariants.py | 30 +-- .../test_scripted_kv_pressure.py | 6 +- .../test_scripted_lifecycle.py | 15 +- .../test_scripted_max_new_tokens.py | 5 +- .../test_scripted_multi_req.py | 18 +- .../chunked_prefill/test_scripted_pp.py | 6 +- .../chunked_prefill/test_scripted_priority.py | 6 +- .../chunked_prefill/test_scripted_radix.py | 18 +- .../test_scripted_regression.py | 24 +- .../chunked_prefill/test_scripted_sampling.py | 30 ++- .../test_scripted_special_case.py | 135 ++++++----- .../dsv4/test_fused_compress_attn_hip.py | 2 - test/manual/ep/test_deepep_internode.py | 32 +-- test/manual/ep/test_deepep_intranode.py | 10 +- test/manual/ep/test_deepep_low_latency.py | 37 +-- test/manual/ep/test_elastic_scale.py | 3 +- test/manual/ep/test_flashinfer_dispatcher.py | 7 +- .../hicache/test_disaggregation_hicache.py | 2 +- test/manual/kernels/test_hisparse_prefetch.py | 12 +- test/manual/lang_frontend/test_choices.py | 1 - .../manual/lang_frontend/test_jump_forward.py | 1 + .../dsa/test_get_k_scale_triton_kernel.py | 2 +- test/manual/lora/test_lora_backend.py | 1 - test/manual/lora/test_lora_cuda_graph.py | 1 - test/manual/lora/test_torch_backend.py | 1 - .../manual/minimax_m3/test_npu_memory_pool.py | 4 +- test/manual/models/test_clip_models.py | 13 +- test/manual/models/test_gme_qwen_models.py | 12 +- test/manual/models/test_kimi_k2_models.py | 2 +- .../models/test_mistral_large3_basic.py | 4 +- test/manual/models/test_qwen3_asr.py | 14 +- .../nightly/test_vlms_piecewise_cuda_graph.py | 2 +- .../nightly/test_vlms_vit_cuda_graph.py | 2 +- .../nightly/test_vlms_vit_flashinfer_cudnn.py | 2 +- .../features/test_cache_report.py | 24 +- .../function_call/test_tool_choice.py | 13 +- test/manual/perf/test_bench_one_batch_1gpu.py | 1 - .../prefill_only/test_cross_encoder_models.py | 7 +- .../test_encoder_embedding_models.py | 7 +- test/manual/quant/test_block_fp8.py | 7 +- .../test_block_fp8_deep_gemm_blackwell.py | 1 - .../quant/test_deepseek_v32_fp4_4gpu.py | 4 +- .../quant/test_deepseek_v3_fp4_4gpu_trtllm.py | 4 +- test/manual/quant/test_eval_fp8_accuracy.py | 1 - .../scheduler/test_no_chunked_prefill.py | 1 - test/manual/spec/eagle/test_eagle3_basic.py | 1 - test/manual/test_create_custom_4d_mask.py | 4 +- test/manual/test_fim_completion.py | 6 +- test/manual/test_get_weights_by_name.py | 1 - test/manual/test_kda_target_verify.py | 12 +- test/manual/test_logprobs.py | 1 - test/manual/test_modelopt.py | 1 - test/manual/test_models_from_modelscope.py | 1 - test/manual/test_ray_engine.py | 4 - test/manual/test_sagemaker_server.py | 18 +- test/manual/test_schedule_policy.py | 1 - .../manual/test_srt_engine_with_quant_args.py | 1 - test/manual/test_tokenizer_batch_encode.py | 1 - test/manual/test_triton_attention_rocm_mla.py | 1 - test/manual/test_vlm_accuracy.py | 3 +- test/manual/test_weight_cache_e2e.py | 4 +- test/manual/test_whisper_cuda_graph.py | 4 +- test/manual/vlm/verify_single_image_gather.py | 4 +- .../test_deepseek_v32_indexcache.py | 4 +- .../mi30x/test_deepseek_r1_eval_amd.py | 4 +- .../mi30x/test_deepseek_v32_dp_eval_amd.py | 6 +- .../mi30x/test_deepseek_v32_eval_amd.py | 4 +- .../mi30x/test_deepseek_v32_tc_eval_amd.py | 6 +- .../amd/accuracy/mi30x/test_glm51_eval_amd.py | 4 +- .../mi30x/test_glm51_hisparse_eval_mi30x.py | 3 +- .../amd/accuracy/mi30x/test_glm5_eval_amd.py | 4 +- .../accuracy/mi30x/test_gpt_oss_eval_amd.py | 4 +- .../amd/accuracy/mi30x/test_grok_eval_amd.py | 4 +- .../amd/accuracy/mi30x/test_gsm8k_eval_amd.py | 16 +- .../mi30x/test_minimax_m25_eval_amd.py | 4 +- .../mi30x/test_minimax_m27_eval_amd.py | 4 +- .../accuracy/mi30x/test_vlms_mmmu_eval_amd.py | 18 +- .../mi35x/test_deepseek_r1_eval_mi35x.py | 4 +- ..._deepseek_r1_mxfp4_ar_fusion_eval_mi35x.py | 4 +- .../test_deepseek_r1_mxfp4_eval_mi35x.py | 4 +- ...est_deepseek_r1_mxfp4_kv_fp8_eval_mi35x.py | 4 +- .../mi35x/test_deepseek_v32_dp_eval_mi35x.py | 6 +- .../mi35x/test_deepseek_v32_eval_mi35x.py | 4 +- .../accuracy/mi35x/test_glm51_eval_mi35x.py | 4 +- .../mi35x/test_glm51_hisparse_eval_mi35x.py | 3 +- .../accuracy/mi35x/test_glm5_eval_mi35x.py | 4 +- .../mi35x/test_glm5_mxfp4_eval_mi35x.py | 4 +- .../accuracy/mi35x/test_gpt_oss_eval_mi35x.py | 4 +- .../test_gpt_oss_w4a8_mxfp4_eval_mi35x.py | 4 +- .../test_kimi_k25_aiter_mla_eval_mi35x.py | 4 +- .../mi35x/test_kimi_k25_mxfp4_eval_mi35x.py | 4 +- .../test_kimi_k27_code_mxfp4_eval_mi35x.py | 4 +- .../mi35x/test_minimax_m25_eval_mi35x.py | 4 +- .../mi35x/test_minimax_m25_tp4_eval_mi35x.py | 4 +- .../mi35x/test_minimax_m27_eval_mi35x.py | 4 +- .../mi35x/test_minimax_m3_tp4_eval_mi35x.py | 4 +- .../mi35x/test_qwen35_mxfp4_eval_mi35x.py | 4 +- .../mi35x/test_qwen3_coder_next_eval_mi35x.py | 4 +- .../mi35x/test_qwen3_moe_eval_mi35x.py | 4 +- .../test_deepseek_r1_0528_mxfp4_eval_mi45x.py | 4 +- .../test_deepseek_v4_flash_eval_mi45x.py | 4 +- .../test_gpt_oss_w4a8_mxfp4_eval_mi45x.py | 4 +- .../test_disaggregation_basic.py | 12 +- .../amd/test_deepseek_r1_mxfp4_8gpu.py | 4 +- .../registered/amd/test_deepseek_v32_basic.py | 8 +- test/registered/amd/test_deepseek_v3_basic.py | 4 +- .../amd/test_deepseek_v3_basic_kv_fp8.py | 4 +- test/registered/amd/test_kimi_k25_mxfp4.py | 4 +- test/registered/amd/test_kimi_k2_instruct.py | 6 +- test/registered/amd/test_moriep_small.py | 8 - .../amd/test_qwen3_coder_next_8gpu.py | 4 +- test/registered/amd/test_qwen3_instruct.py | 4 +- .../registered/amd/test_qwen3_instruct_fp8.py | 4 +- .../amd/test_qwen3_instruct_mxfp4.py | 4 +- .../attention/test_kda_decode_flashinfer.py | 18 +- .../attention/test_kda_prefill_flashkda.py | 12 +- .../test_triton_attention_kernels.py | 1 - .../attention/test_triton_sliding_window.py | 7 +- .../attention/test_wave_attention_kernels.py | 1 - .../dense/test_extend_init_contract.py | 1 - .../unittests/dense/test_trtllm_mha.py | 14 +- .../attention/unittests/mla/test_flashmla.py | 5 +- .../test_flashinfer_fusion_preflight.py | 4 +- test/registered/core/test_srt_endpoint.py | 1 - test/registered/core/test_srt_engine.py | 1 - test/registered/cp/test_dsa_prefill_cp.py | 3 +- test/registered/cpu/test_cpu_graph.py | 1 - test/registered/cpu/test_extend.py | 1 - test/registered/cpu/test_flash_attn.py | 1 - test/registered/cpu/test_gemm.py | 1 - .../cpu/test_intel_amx_attention_backend_a.py | 1 - .../cpu/test_intel_amx_attention_backend_b.py | 1 - .../cpu/test_intel_amx_attention_backend_c.py | 1 - test/registered/cpu/test_moe.py | 1 - test/registered/cpu/test_norm.py | 5 - .../cpu/test_rank_consensus_checker.py | 6 +- .../cpu/test_server_args_backend.py | 23 +- test/registered/cpu/test_spec_kernels.py | 161 +++++++++++-- test/registered/dcp/test_dcp_layout_unit.py | 11 +- .../registered/dcp/test_qwen3p5_triton_dcp.py | 2 +- .../dcp/test_reduce_scatter_along_dim.py | 3 +- .../tensor_comparator/test_formatter.py | 7 +- .../debug_utils/comparator/test_entrypoint.py | 24 +- .../comparator/test_meta_overrider.py | 12 +- .../comparator/test_output_types.py | 8 +- .../source_patcher/test_source_editor.py | 76 +++--- test/registered/debug_utils/test_dumper.py | 101 ++++---- .../test_engine_dumper_comparator_e2e.py | 2 +- .../test_tensor_dump_forward_hook.py | 1 - .../test_disaggregation_basic.py | 12 +- .../test_kimi_linear_pd_dcp4.py | 3 +- .../dp_attn/test_dp_attention_bcg_kl.py | 15 +- test/registered/eplb/test_lplb_distributed.py | 16 +- .../function_call/test_kimik2_detector.py | 4 +- .../gemm/test_linear_bf16_fp32_hpc.py | 1 - .../communication/bench_custom_all_reduce.py | 2 +- .../communication/bench_tp_qknorm.py | 2 +- .../diffusion/bench_fused_norm_scale_shift.py | 8 +- .../diffusion/bench_residual_gate_add.py | 4 +- .../benchmark/kv_canary/bench_verify.py | 4 +- .../kernels/ops/activation/test_activation.py | 6 +- .../kernels/ops/attention/test_cutedsl_gdn.py | 10 +- .../attention/test_dsv32_indexer_fusion.py | 6 +- .../attention/test_flash_attention_4_sm120.py | 8 +- .../ops/attention/test_flash_mla_backends.py | 4 +- .../ops/attention/test_fused_metadata_copy.py | 52 +++-- .../attention/test_fused_store_index_cache.py | 6 +- .../test_minimax_decode_topk_page_table.py | 12 +- .../ops/attention/test_paged_mqa_metadata.py | 12 +- .../kernels/ops/attention/test_pos_enc.py | 2 +- .../test_q8kv8_sparse_prefill_backend.py | 7 +- .../ops/attention/test_qprep_bf16_fp8_sm90.py | 6 +- .../attention/test_sm120_paged_mqa_logits.py | 6 +- .../kernels/ops/attention/test_topk_v2.py | 6 +- ...test_amd_deterministic_custom_allreduce.py | 14 +- .../test_amd_nccl_allreduce_determinism.py | 12 +- .../gemm/test_minimax_fused_qkv_index_gemm.py | 6 +- .../ops/kv_canary/test_kernel_config.py | 66 +++--- .../ops/kv_canary/test_pipeline_e2e.py | 12 +- .../kernels/ops/kv_canary/test_plan_hand.py | 54 ++--- .../kernels/ops/kv_canary/test_verify_hand.py | 36 +-- .../kernels/ops/kv_canary/test_write_fuzz.py | 6 +- .../kernels/ops/kv_canary/test_write_hand.py | 42 ++-- .../kernels/ops/kvcache/test_hicache.py | 12 +- .../layernorm/test_gemma4_fused_routing.py | 6 +- .../kernels/ops/layernorm/test_rmsnorm_hf.py | 6 +- .../ops/moe/test_minimax_quant_scatter.py | 6 +- .../ops/moe/test_moe_lora_align_block_size.py | 6 +- .../kernels/ops/moe/test_moe_topk_sigmoid.py | 46 ++-- .../ops/moe/test_shuffle_rows_with_scales.py | 6 +- .../test_per_token_group_quant_8bit_v2.py | 6 +- test/registered/kernels/test_jit_cache.py | 6 +- .../kernels/test_kda_replayssm_ring_ragged.py | 6 +- .../kernels/test_kernel_inventory.py | 10 +- .../registered/kernels/test_lean_attention.py | 6 +- .../test_quick_allreduce_bf16_range.py | 6 +- .../kv_canary/test_self_e2e_pp_baseline.py | 1 - .../kv_canary/test_self_e2e_pp_perturb.py | 1 - .../kv_canary/test_self_unit_endpoint.py | 93 +++++--- .../kv_canary/test_self_unit_perturb.py | 59 +++-- .../test_self_unit_runner_per_forward.py | 32 +-- .../test_self_unit_runner_swa_divergence.py | 14 +- .../layers/mamba/test_mamba2_mixer.py | 6 +- .../layers/mamba/test_mamba_ssm_ssd.py | 3 - ...test_lora_deepseek_v3_base_logprob_diff.py | 1 - .../test_lora_gpt_oss_20b_logprob_diff.py | 1 - .../lora/test_lora_kimi_k25_logprob_diff.py | 1 - .../test_lora_moe_vllm_sgl_logprob_diff.py | 2 - ...nemotron_3_super_120b_a12b_logprob_diff.py | 1 - .../lora/test_lora_overlap_loading.py | 9 +- ...wen3_30b_a3b_instruct_2507_logprob_diff.py | 1 - .../test_lora_qwen3_5_35b_a3b_logprob_diff.py | 1 - .../lora/test_lora_qwen3_5_4b_logprob_diff.py | 1 - .../lora/test_lora_qwen3_8b_logprob_diff.py | 1 - ..._qwen3_vl_30b_a3b_instruct_logprob_diff.py | 1 - test/registered/lora/test_lora_tp.py | 1 - .../test_load_weights_from_remote_instance.py | 1 - ...t_load_weights_from_remote_instance_npu.py | 1 - .../models_e2e/test_dummy_grok_models.py | 1 - .../models_e2e/test_generation_models.py | 1 - .../models_e2e/test_minimax_m25_basic.py | 4 +- test/registered/moe/test_hpc_ops_moe.py | 1 - .../test_npu_deepseek_v3_2_8p_aime25.py | 1 - .../test_npu_glm4_7_flash_1p_aime25.py | 1 - .../test_npu_glm4_7_flash_1p_gsm8k.py | 1 - ...est_npu_glm5_top64_pruned_bf16_8p_gsm8k.py | 1 - ..._k2_6_w4a8_16p_in64k_out1k_100ms_aime25.py | 1 - .../HiCache/test_npu_hicache_mha.py | 1 - .../HiCache/test_npu_hicache_mla.py | 1 - .../dllm/test_npu_llada2_mini.py | 3 +- .../offloading/test_npu_offload_modes.py | 1 - .../test_npu_compile_graph_tp1_bf16.py | 1 - .../test_npu_graph_tp1_bf16.py | 1 - .../test_npu_graph_tp2_bf16.py | 1 - .../expert_parallelism/test_npu_deepep.py | 1 - ...b_min_rebalancing_utilization_threshold.py | 2 +- .../parameter/test_npu_fim_completion.py | 6 +- .../quant/test_npu_autoround_dense.py | 1 - .../quant/test_npu_autoround_moe.py | 1 - .../npu/basic_function/quant/test_npu_gguf.py | 1 - .../basic_function/quant/test_npu_gguf_moe.py | 1 - .../basic_function/quant/test_npu_gptq_moe.py | 1 - .../quant/test_npu_w4a4_quantization.py | 1 - .../runtime_opts/test_npu_mla_fia_w8a8int8.py | 1 - .../runtime_opts/test_npu_mla_w8a8int8.py | 1 - .../runtime_opts/test_npu_tp1_bf16.py | 1 - .../runtime_opts/test_npu_tp2_bf16.py | 1 - .../runtime_opts/test_npu_tp2_fia_bf16.py | 1 - .../runtime_opts/test_npu_tp4_bf16.py | 1 - .../test_npu_speculative_moe_a2a_backend.py | 1 - .../test_npu_bge_large_en_v1_5.py | 7 +- .../interface/test_npu_api_abort_request.py | 1 - .../test_npu_openai_function_calling.py | 12 +- .../test_npu_bge_reranker_v2_m3.py | 6 +- .../test_npu_gemma_2_27b_v0_2.py | 7 +- .../test_npu_llama_3_1_8b_v0_2.py | 6 +- .../npu/test_npu_memory_consumption.py | 1 - .../openai_server/basic/test_openai_server.py | 66 +++--- .../test_openai_server_hidden_states.py | 33 ++- .../test_openai_function_calling.py | 12 +- .../ops/test_aiter_allgather_amd.py | 1 - .../ops/test_aiter_allreduce_fusion_amd.py | 1 - .../perf/test_bench_one_batch_2gpu.py | 1 - .../prefill_only/test_embedding_models.py | 7 +- .../prefill_only/test_reward_models.py | 7 +- .../test_diffusion_stage_profiler_sync.py | 6 +- test/registered/profiling/test_profile_v2.py | 1 - .../profiling/test_start_profile.py | 1 - test/registered/quant/test_block_int8.py | 6 +- test/registered/quant/test_fp8_utils.py | 49 ++-- .../test_unified_radix_cache_hicache_pp_kl.py | 6 +- .../rl/test_lora_load_from_tensor.py | 4 +- .../rl/test_return_routed_experts.py | 10 +- .../test_update_weights_from_distributed.py | 25 +- .../rl/test_update_weights_from_tensor.py | 12 +- test/registered/rust/test_rust_extension.py | 3 +- .../scheduler/test_min_free_slots_delayer.py | 2 - .../scheduler/test_prefill_delayer.py | 6 +- .../test_scripted_runtime_core.py | 217 +++++++++--------- .../sessions/test_session_control.py | 24 +- .../test_dspark_block_accept_estimator.py | 1 - .../spec/dspark/test_dspark_scheduler.py | 1 - .../tokenizer/test_skip_tokenizer_init.py | 8 +- .../test_tbo_filter_batch_marker.py | 9 +- .../unit/bench/test_mmmu_eval_utils.py | 7 +- .../test_multimodal_piecewise_cuda_graph.py | 4 +- .../test_decode_queue_cleanup.py | 3 +- .../test_disaggregation_wire.py | 13 +- .../disaggregation/test_nixl_backend_basic.py | 4 +- .../unit/distributed/test_parallel_state.py | 40 ++-- .../unit/entrypoints/openai/test_protocol.py | 5 +- .../test_effective_state_surfaces.py | 20 +- .../test_function_call_parser.py | 114 +++++---- .../function_call/test_hunyuan_detector.py | 4 +- .../test_kimik3_structural_tag.py | 2 +- .../function_call/test_minicpm5_detector.py | 18 +- .../test_muse_glimmer_detector.py | 3 +- .../test_poolside_v1_detector.py | 4 +- .../mlx/test_attention_patching.py | 6 +- .../mlx/test_attn_dp_request_capacity.py | 2 +- .../hardware_backend/mlx/test_fused_swiglu.py | 6 +- .../mlx/test_max_running_requests.py | 6 +- .../mlx/test_metal_profiler.py | 21 +- .../mlx/test_muse_glimmer_mlx_model.py | 11 +- .../hardware_backend/mlx/test_quantization.py | 4 +- .../attention/test_dsa_mqa_logits_chunking.py | 11 +- .../attention/test_kda_helion_dispatcher.py | 3 +- .../attention/test_vision_strided_qkv.py | 7 +- .../layers/moe/test_mega_moe_deepgemm_api.py | 4 +- .../layers/moe/test_moe_runner_extensions.py | 59 ++--- .../test_deepgemm_ue8m0_requant.py | 70 +++--- .../test_flashinfer_trtllm_fp8_fallback.py | 21 +- .../test_mxfp4_flashinfer_activation_prep.py | 114 +++++---- .../quantization/test_mxfp4_sm90_cutlass.py | 6 +- .../quantization/test_nvfp4_moe_backends.py | 22 +- .../registered/unit/layers/test_conv_layer.py | 3 - .../layers/test_dsv4_kv_splits_heuristic.py | 6 +- .../layers/test_flashinfer_comm_fusion.py | 9 +- .../managers/test_load_snapshot_backends.py | 23 +- .../unit/managers/test_loadstat_wire.py | 6 +- ...test_priority_scheduling_disaggregation.py | 30 ++- .../mem_cache/test_hicache_file_lru_unit.py | 14 +- .../mem_cache/test_hybrid_pool_assembler.py | 4 +- .../mem_cache/test_kv_index_translator.py | 1 - .../test_mamba_donated_alloc_ratio.py | 4 +- .../mem_cache/test_mamba_path_state_cap.py | 9 +- .../unit/mem_cache/test_mamba_unittest.py | 18 +- .../unit/mem_cache/test_mem_pool_host.py | 21 +- .../unit/mem_cache/test_mmap_allocator.py | 9 +- .../mem_cache/test_multi_ended_allocator.py | 16 +- .../test_radix_cache_slru_accuracy.py | 1 - .../test_rust_tree_core_integration.py | 5 +- .../test_session_token_share_unit.py | 1 - .../mem_cache/test_swa_eviction_boundary.py | 1 - .../unit/mem_cache/test_swa_unittest.py | 18 +- .../test_unified_free_no_host_sync.py | 11 +- .../test_unified_radix_cache_bench.py | 6 +- .../test_unified_radix_cache_unittest.py | 81 ++++--- .../unit/mem_cache/test_unified_tri_pool.py | 11 +- .../runner/test_decode_cuda_graph_runner.py | 79 ++++--- .../test_hidden_state_graph_recapture.py | 9 +- .../test_full_cuda_graph_backend.py | 7 +- .../test_model_runner_decode_rows.py | 20 +- .../model_executor/test_pool_configurator.py | 5 +- .../model_loader/test_presharded_loader.py | 169 ++++++++------ .../unit/models/test_kimi_k3_vision.py | 31 +-- test/registered/unit/models/test_kimi_vl.py | 9 +- .../test_qwen3_vl_feature_materialization.py | 22 +- .../test_request_metrics_exporter.py | 12 +- .../unit/parser/test_template_manager.py | 7 +- .../sampling/test_custom_logit_processor.py | 1 - .../unit/sampling/test_penaltylib.py | 6 - .../unit/sampling/test_sampling_batch_info.py | 14 +- .../test_background_http_poster.py | 7 - .../unit/scripted_runtime/test_http_server.py | 5 - .../scripted_runtime/test_scheduler_hook.py | 1 - .../test_scripted_runtime_utils.py | 2 - .../test_tokenizer_recv_proxy.py | 5 - .../test_resolution_declarations.py | 3 +- .../test_resolution_reads_the_declarations.py | 4 +- .../spec/test_dflash_extra_buffer_lazy.py | 13 +- .../test_eagle_worker_v2_topk1_fastpath.py | 15 +- .../spec/test_ngram_mamba_verify_update.py | 15 +- .../test_routed_experts_scattered_a2a.py | 30 +-- test/registered/unit/test_cuda_vmm_utils.py | 4 +- .../unit/test_dsa_tilelang_fp8_validation.py | 1 - .../unit/test_eval_accuracy_kit_sgl_eval.py | 10 +- test/registered/unit/test_fork_test_worker.py | 18 +- .../unit/test_global_config_read_ratchet.py | 3 +- test/registered/unit/test_model_overrides.py | 11 +- .../unit/test_runtime_context_config_bags.py | 22 +- .../unit/test_server_args_migration.py | 1 - ...test_supplied_instance_exposure_ratchet.py | 24 +- .../unit/utils/test_profile_merger.py | 1 - .../unit/utils/test_weight_checker.py | 19 +- .../utils/test_weight_checker_comparator.py | 2 - .../utils/test_bench_typebaseddispatcher.py | 12 +- test/registered/utils/test_numa_utils.py | 15 +- .../vlm/test_vision_chunked_prefill.py | 9 +- .../vlm/test_vision_openai_server_a.py | 6 +- test/registered/xpu/test_intel_xpu_backend.py | 1 - test/registered/xpu/test_moe_ld_padding.py | 26 ++- test/registered/xpu/test_topk.py | 3 - .../xpu/test_triton_attention_backend.py | 1 - test/registered/xpu/test_xpu_basic.py | 1 - 1411 files changed, 7766 insertions(+), 8176 deletions(-) diff --git a/.claude/skills/llm-torch-profiler-analysis/scripts/triage_kernel_helpers.py b/.claude/skills/llm-torch-profiler-analysis/scripts/triage_kernel_helpers.py index 81a71cb96..de702db73 100644 --- a/.claude/skills/llm-torch-profiler-analysis/scripts/triage_kernel_helpers.py +++ b/.claude/skills/llm-torch-profiler-analysis/scripts/triage_kernel_helpers.py @@ -541,7 +541,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( "fused_flashmla_metadata", ), rationale_hint=( - "NSA replay metadata copies are already fused into one-kernel" " families." + "NSA replay metadata copies are already fused into one-kernel families." ), min_share=0.02, likely_share=0.2, @@ -787,7 +787,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( ("softmax", "sampling"), ), rationale_hint=( - "Decode-time sampling already has fused temperature and softmax" " kernels." + "Decode-time sampling already has fused temperature and softmax kernels." ), min_share=0.05, likely_share=0.5, @@ -1218,8 +1218,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( FusionPatternSpec( pattern="vLLM fused residual add + RMSNorm", candidate_path=( - "vllm/_custom_ops.py" - "
vllm/compilation/passes/fusion/rms_quant_fusion.py" + "vllm/_custom_ops.py
vllm/compilation/passes/fusion/rms_quant_fusion.py" ), active_keywords=( "fused_add_rms_norm", @@ -1236,8 +1235,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( FusionPatternSpec( pattern="vLLM fused activation-and-mul", candidate_path=( - "vllm/_custom_ops.py" - "
vllm/compilation/passes/fusion/act_quant_fusion.py" + "vllm/_custom_ops.py
vllm/compilation/passes/fusion/act_quant_fusion.py" ), active_keywords=( "silu_and_mul", diff --git a/.claude/skills/mechanical-refactor-verify/scripts/mechanical_refactor_proof_generator.py b/.claude/skills/mechanical-refactor-verify/scripts/mechanical_refactor_proof_generator.py index 598ac7a82..0aeb0ea5f 100644 --- a/.claude/skills/mechanical-refactor-verify/scripts/mechanical_refactor_proof_generator.py +++ b/.claude/skills/mechanical-refactor-verify/scripts/mechanical_refactor_proof_generator.py @@ -256,7 +256,9 @@ def _module_assign_names(text: str) -> set: targets = ( node.targets if isinstance(node, ast.Assign) - else [node.target] if isinstance(node, ast.AnnAssign) else [] + else [node.target] + if isinstance(node, ast.AnnAssign) + else [] ) names |= {t.id for t in targets if isinstance(t, ast.Name)} return names diff --git a/.claude/skills/mechanical-refactor-verify/scripts/mechanical_refactor_reproduction_utils.py b/.claude/skills/mechanical-refactor-verify/scripts/mechanical_refactor_reproduction_utils.py index fb0284dba..8c2856a7e 100644 --- a/.claude/skills/mechanical-refactor-verify/scripts/mechanical_refactor_reproduction_utils.py +++ b/.claude/skills/mechanical-refactor-verify/scripts/mechanical_refactor_reproduction_utils.py @@ -173,9 +173,9 @@ def _find_unique_def( if isinstance(node, definition) and node.name == name ] assert matches, f"{name} not found in {where}" - assert ( - len(matches) == 1 - ), f"{len(matches)} defs named {name} in {where}; pass from_class to disambiguate" + assert len(matches) == 1, ( + f"{len(matches)} defs named {name} in {where}; pass from_class to disambiguate" + ) return matches[0] @@ -287,9 +287,9 @@ def _lowered_call_text(text: str, node: ast.Call) -> str: """ receiver = node.args[0] receiver_src = _node_slice(text, receiver) - assert ( - "\n" not in receiver_src and "#" not in receiver_src - ), f"receiver {receiver_src!r} must be single-line and comment-free" + assert "\n" not in receiver_src and "#" not in receiver_src, ( + f"receiver {receiver_src!r} must be single-line and comment-free" + ) opener = _slice_span( text, node.func.end_lineno, @@ -722,9 +722,9 @@ class Repro: ) existing = [alias_text(a.name, a.asname) for a in node.names] added = alias_text(name, asname) - assert ( - added not in existing - ), f"{name!r} already imported from {module!r} in {rel}" + assert added not in existing, ( + f"{name!r} already imported from {module!r} in {rel}" + ) rebuilt = f"from {module} import " + ", ".join(existing + [added]) + nl lines[node.lineno - 1 : node.end_lineno] = [rebuilt] _write_source(path, "".join(lines)) @@ -825,9 +825,9 @@ class Repro: for node in tree.body if isinstance(node, (ast.Import, ast.ImportFrom)) ] - assert ( - imports - ), f"no imports to anchor a new `if TYPE_CHECKING:` block in {rel}" + assert imports, ( + f"no imports to anchor a new `if TYPE_CHECKING:` block in {rel}" + ) insert_at = imports[-1].end_lineno lines[insert_at:insert_at] = [ nl, @@ -864,9 +864,9 @@ class Repro: replaced = lines[node.lineno - 1].replace( f"from {spelled} import", f"from {new_module} import", 1 ) - assert ( - replaced != lines[node.lineno - 1] - ), f"import spelling {spelled!r} not found on its line in {rel}" + assert replaced != lines[node.lineno - 1], ( + f"import spelling {spelled!r} not found on its line in {rel}" + ) lines[node.lineno - 1] = replaced changed = True assert changed, f"nested import of {name} from {old_module} not in {rel}" @@ -974,9 +974,9 @@ class Repro: ``self: Target`` annotation is dropped (redundant inside the class). The body is moved verbatim; the formatter normalises the surrounding blank lines. """ - assert ( - before is None or after is None - ), "move_symbol: before and after are mutually exclusive" + assert before is None or after is None, ( + "move_symbol: before and after are mutually exclusive" + ) def op(root: Path) -> None: src_path = root / src @@ -1255,15 +1255,17 @@ class Repro: targets = ( node.targets if isinstance(node, ast.Assign) - else [node.target] if isinstance(node, ast.AnnAssign) else [] + else [node.target] + if isinstance(node, ast.AnnAssign) + else [] ) names = {t.id for t in targets if isinstance(t, ast.Name)} hit = names & dropped if not hit: continue - assert len(names) == len( - targets - ), f"drop_assigns {sorted(hit)}: non-name targets in {src}" + assert len(names) == len(targets), ( + f"drop_assigns {sorted(hit)}: non-name targets in {src}" + ) value_src = ast.unparse(node.value) if node.value is not None else None for dropped_name in hit: removed_assigns[dropped_name] = value_src @@ -1289,15 +1291,17 @@ class Repro: else: assign_spans.append((node.lineno, node.end_lineno)) found_assigns |= hit - assert ( - found_assigns == dropped - ), f"{dropped - found_assigns} not assigned in {src}" + assert found_assigns == dropped, ( + f"{dropped - found_assigns} not assigned in {src}" + ) rederivable: dict[str, str | None] = {} for node in tree.body: targets = ( node.targets if isinstance(node, ast.Assign) - else [node.target] if isinstance(node, ast.AnnAssign) else [] + else [node.target] + if isinstance(node, ast.AnnAssign) + else [] ) names = [t.id for t in targets if isinstance(t, ast.Name)] if not names or set(names) & dropped: @@ -1383,9 +1387,9 @@ class Repro: src_text = _read_source(src_path) assert src_text.count(body) == 1, f"block not found uniquely in {src}" at = src_text.find(body) - assert ( - at == 0 or src_text[at - 1] == "\n" - ), f"block matches mid-line in {src}; it must start at a line boundary" + assert at == 0 or src_text[at - 1] == "\n", ( + f"block matches mid-line in {src}; it must start at a line boundary" + ) _write_source(src_path, src_text.replace(body, call, 1)) dst_path = root / dst diff --git a/.claude/skills/mechanical-refactor-verify/scripts/tests/proof_generator/test_infer_moves.py b/.claude/skills/mechanical-refactor-verify/scripts/tests/proof_generator/test_infer_moves.py index 0b9e5ead4..8e5e4b601 100644 --- a/.claude/skills/mechanical-refactor-verify/scripts/tests/proof_generator/test_infer_moves.py +++ b/.claude/skills/mechanical-refactor-verify/scripts/tests/proof_generator/test_infer_moves.py @@ -370,12 +370,7 @@ def test_infer_recipe_module_level_def_shadowed_by_method_name(repo: Path) -> No " return foo(x=self.x)\n" ), "util.py": ( - "def keep():\n" - " return 1\n" - "\n" - "\n" - "def foo(*, x):\n" - " return x + 1\n" + "def keep():\n return 1\n\n\ndef foo(*, x):\n return x + 1\n" ), }, ) @@ -437,9 +432,7 @@ def test_infer_recipe_move_leaving_a_forwarding_delegate(repo: Path) -> None: _write( repo, **{ - "model.py": ( - "class M:\n" " def work(self, x):\n" " return x + 1\n" - ), + "model.py": ("class M:\n def work(self, x):\n return x + 1\n"), "comp.py": "class C:\n def keep(self):\n return 1\n", }, ) @@ -448,9 +441,7 @@ def test_infer_recipe_move_leaving_a_forwarding_delegate(repo: Path) -> None: repo, **{ "model.py": ( - "class M:\n" - " def work(self, x):\n" - " return self.comp.work(x)\n" + "class M:\n def work(self, x):\n return self.comp.work(x)\n" ), "comp.py": ( "class C:\n" diff --git a/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_cli/cli_testlib.py b/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_cli/cli_testlib.py index fb1eef19b..e9490fdfa 100644 --- a/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_cli/cli_testlib.py +++ b/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_cli/cli_testlib.py @@ -2,13 +2,9 @@ import subprocess from pathlib import Path _PASSING_PROOF = ( - "import sys\n" - 'print("PASS: reproduces the commit byte-for-byte.")\n' - "sys.exit(0)\n" -) -_FAILING_PROOF = ( - "import sys\n" 'print("RESIDUAL (2 lines):\\n+x\\n-y")\n' "sys.exit(1)\n" + 'import sys\nprint("PASS: reproduces the commit byte-for-byte.")\nsys.exit(0)\n' ) +_FAILING_PROOF = 'import sys\nprint("RESIDUAL (2 lines):\\n+x\\n-y")\nsys.exit(1)\n' def _git(repo: Path, *args: str) -> str: diff --git a/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_add_imports.py b/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_add_imports.py index 0cf0bd58e..3b63e3ad4 100644 --- a/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_add_imports.py +++ b/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_add_imports.py @@ -110,13 +110,7 @@ def test_add_typechecking_import_inserts_in_block(tmp_path: Path) -> None: def test_add_typechecking_import_creates_missing_block(tmp_path: Path) -> None: """With no TYPE_CHECKING block, one is created after the trailing module import.""" (tmp_path / "m.py").write_text( - "from typing import TYPE_CHECKING\n" - "\n" - "from a import X\n" - "\n" - "\n" - "def f():\n" - " pass\n" + "from typing import TYPE_CHECKING\n\nfrom a import X\n\n\ndef f():\n pass\n" ) r = Repro("b", "t").add_typechecking_import("m.py", "from b import Y") _apply(r, tmp_path) @@ -215,12 +209,7 @@ def test_add_typechecking_import_raises_without_imports(tmp_path: Path) -> None: def test_add_typechecking_import_drops_a_lone_pass_placeholder(tmp_path: Path) -> None: """Populating a `pass`-only TYPE_CHECKING block replaces the placeholder.""" (tmp_path / "m.py").write_text( - "from typing import TYPE_CHECKING\n" - "\n" - "if TYPE_CHECKING:\n" - " pass\n" - "\n" - "x = 1\n" + "from typing import TYPE_CHECKING\n\nif TYPE_CHECKING:\n pass\n\nx = 1\n" ) r = Repro("b", "t").add_typechecking_import("m.py", "from b import Y") _apply(r, tmp_path) diff --git a/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_extract_symbols_to_new_module.py b/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_extract_symbols_to_new_module.py index 26d9fae2f..7e0b1ebbf 100644 --- a/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_extract_symbols_to_new_module.py +++ b/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_extract_symbols_to_new_module.py @@ -111,11 +111,7 @@ def test_extract_symbols_to_new_module_drops_relocated_assigns(tmp_path: Path) - " return _FLAG\n" ) header = ( - "from __future__ import annotations\n" - "\n" - "import os\n" - "\n" - "_FLAG = os.cpu_count()\n" + "from __future__ import annotations\n\nimport os\n\n_FLAG = os.cpu_count()\n" ) r = Repro("b", "t").extract_symbols_to_new_module( "src.py", diff --git a/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_move_assign.py b/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_move_assign.py index 762eab97f..1ba4e192e 100644 --- a/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_move_assign.py +++ b/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_move_assign.py @@ -19,13 +19,7 @@ def test_move_assign_relocates_a_module_constant(tmp_path: Path) -> None: _apply(r, tmp_path) assert "LIMIT" not in (tmp_path / "src.py").read_text().split("def stay")[0] assert (tmp_path / "dst.py").read_text() == ( - "import sys\n" - "\n" - "LIMIT = 480 # seconds\n" - "\n" - "\n" - "def keep():\n" - " return 1\n" + "import sys\n\nLIMIT = 480 # seconds\n\n\ndef keep():\n return 1\n" ) @@ -50,13 +44,7 @@ def test_move_assign_relocates_an_annotated_constant(tmp_path: Path) -> None: _apply(r, tmp_path) assert "LIMIT" not in (tmp_path / "src.py").read_text().split("def stay")[0] assert (tmp_path / "dst.py").read_text() == ( - "import sys\n" - "\n" - "LIMIT: int = 480\n" - "\n" - "\n" - "def keep():\n" - " return 1\n" + "import sys\n\nLIMIT: int = 480\n\n\ndef keep():\n return 1\n" ) diff --git a/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_move_symbol.py b/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_move_symbol.py index a5f74a687..fc808a998 100644 --- a/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_move_symbol.py +++ b/.claude/skills/mechanical-refactor-verify/scripts/tests/reproduction_utils/test_move_symbol.py @@ -260,13 +260,7 @@ def test_move_symbol_dedent_leaves_string_literal_interior_lines( ) _apply(r, tmp_path) assert (tmp_path / "dst.py").read_text() == ( - "import os\n" - "\n" - "def foo(self):\n" - " s = '''raw\n" - " partial\n" - "'''\n" - " return s\n" + "import os\n\ndef foo(self):\n s = '''raw\n partial\n'''\n return s\n" ) diff --git a/.claude/skills/sglang-prod-incident-triage/scripts/incident_artifact_tool.py b/.claude/skills/sglang-prod-incident-triage/scripts/incident_artifact_tool.py index 2a9dacc69..fdc4aa875 100755 --- a/.claude/skills/sglang-prod-incident-triage/scripts/incident_artifact_tool.py +++ b/.claude/skills/sglang-prod-incident-triage/scripts/incident_artifact_tool.py @@ -101,8 +101,7 @@ def format_summary_line(filename: str, result: Dict[str, Any]) -> str: if result.get("ok"): return f"{filename}: ok" return ( - f"{filename}: failed status={result.get('status')} " - f"error={result.get('error')}" + f"{filename}: failed status={result.get('status')} error={result.get('error')}" ) @@ -617,7 +616,9 @@ def summarize_dump_file(path: Path, max_requests: int, preview_chars: int) -> st time_span = ( max(timestamps) - min(timestamps) if len(timestamps) >= 2 - else 0.0 if len(timestamps) == 1 else None + else 0.0 + if len(timestamps) == 1 + else None ) lines = [ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index cc6d6c1ed..32ac18680 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -48,10 +48,7 @@ repos: python/sglang/srt/grpc/.*_pb2\.pyi$| python/sglang/srt/grpc/.*_pb2_grpc\.pyi$| )$ - - repo: https://github.com/psf/black - rev: 26.1.0 - hooks: - - id: black-jupyter + - id: ruff-format exclude: '^python/sglang/srt/grpc/.*_pb2\.py$|^python/sglang/srt/grpc/.*_pb2_grpc\.py$|^python/sglang/srt/grpc/.*_pb2\.pyi$|^python/sglang/srt/grpc/.*_pb2_grpc\.pyi$' - repo: https://github.com/codespell-project/codespell rev: v2.4.1 diff --git a/3rdparty/amd/tuning/benchmark_moe_rocm.py b/3rdparty/amd/tuning/benchmark_moe_rocm.py index d7ea67c5f..71b4d1710 100644 --- a/3rdparty/amd/tuning/benchmark_moe_rocm.py +++ b/3rdparty/amd/tuning/benchmark_moe_rocm.py @@ -187,8 +187,10 @@ def run_grid(bs, model, method, tp_size, dtype: str): configs = union_of_list_of_dicts(prune_configs_1, prune_configs_2) - print(f"{bs=} || {len(full_configs)=} | {len(prune_configs_1)=} | \ - {len(prune_configs_2)=} | {len(configs)=}") + print( + f"{bs=} || {len(full_configs)=} | {len(prune_configs_1)=} | \ + {len(prune_configs_2)=} | {len(configs)=}" + ) best_config = None best_time_us = 1e20 diff --git a/benchmark/asr/bench_sglang.py b/benchmark/asr/bench_sglang.py index 875ed952b..c60685d62 100644 --- a/benchmark/asr/bench_sglang.py +++ b/benchmark/asr/bench_sglang.py @@ -343,7 +343,7 @@ def run_evaluation(args): print("\n" + "=" * 20 + " Sample Predictions " + "=" * 20) num_to_show = min(args.print_n, len(results)) for i in range(num_to_show): - print(f"Sample {i+1}:") + print(f"Sample {i + 1}:") print(f" REF: {references[i]}") print(f" PRED: {predictions[i]}") print("-" * 40) diff --git a/benchmark/bench_linear_attention/bench_kda_flashinfer_mtp.py b/benchmark/bench_linear_attention/bench_kda_flashinfer_mtp.py index 5e988392e..903841d57 100644 --- a/benchmark/bench_linear_attention/bench_kda_flashinfer_mtp.py +++ b/benchmark/bench_linear_attention/bench_kda_flashinfer_mtp.py @@ -243,9 +243,7 @@ def run(task, fi, tri, device, dtype, args): ) # noqa: E731 else: inp = make_decode_inputs(B, H, HV, K, V, pool, device, dtype) - corr = lambda kern: call_decode( - kern, inp, inp["ssm"].clone() - ) # noqa: E731 + corr = lambda kern: call_decode(kern, inp, inp["ssm"].clone()) # noqa: E731 ssm_t = inp["ssm"].clone() timed = lambda kern: call_decode(kern, inp, ssm_t) # noqa: E731 diff --git a/benchmark/gsm8k/bench_sglang.py b/benchmark/gsm8k/bench_sglang.py index a0e09a39c..b3989673a 100644 --- a/benchmark/gsm8k/bench_sglang.py +++ b/benchmark/gsm8k/bench_sglang.py @@ -53,9 +53,9 @@ def main(args): if args.enable_thinking: from transformers import AutoTokenizer - assert ( - args.tokenizer_path is not None - ), "--tokenizer-path is required when --enable-thinking is set" + assert args.tokenizer_path is not None, ( + "--tokenizer-path is required when --enable-thinking is set" + ) tokenizer = AutoTokenizer.from_pretrained( args.tokenizer_path, trust_remote_code=True ) diff --git a/benchmark/hf3fs/bench_client.py b/benchmark/hf3fs/bench_client.py index ef4a967c4..7b18a8b2c 100644 --- a/benchmark/hf3fs/bench_client.py +++ b/benchmark/hf3fs/bench_client.py @@ -14,11 +14,11 @@ def print_stats(x: List[int]): x = sorted(x) lenx = len(x) print( - f"mean = {sum(x)/len(x):.2f}, " + f"mean = {sum(x) / len(x):.2f}, " f"min = {min(x):.2f}, " - f"p25 = {x[int(lenx*0.25)]:.2f}, " - f"p50 = {x[int(lenx*0.5)]:.2f}, " - f"p75 = {x[int(lenx*0.75)]:.2f}, " + f"p25 = {x[int(lenx * 0.25)]:.2f}, " + f"p50 = {x[int(lenx * 0.5)]:.2f}, " + f"p75 = {x[int(lenx * 0.75)]:.2f}, " f"max = {max(x):.2f}" ) diff --git a/benchmark/hf3fs/bench_storage.py b/benchmark/hf3fs/bench_storage.py index f0ce171bf..71122a41d 100644 --- a/benchmark/hf3fs/bench_storage.py +++ b/benchmark/hf3fs/bench_storage.py @@ -18,11 +18,11 @@ def print_stats(x: List[int]): x = sorted(x) lenx = len(x) print( - f"mean = {sum(x)/len(x):.2f}, " + f"mean = {sum(x) / len(x):.2f}, " f"min = {min(x):.2f}, " - f"p25 = {x[int(lenx*0.25)]:.2f}, " - f"p50 = {x[int(lenx*0.5)]:.2f}, " - f"p75 = {x[int(lenx*0.75)]:.2f}, " + f"p25 = {x[int(lenx * 0.25)]:.2f}, " + f"p50 = {x[int(lenx * 0.5)]:.2f}, " + f"p75 = {x[int(lenx * 0.75)]:.2f}, " f"max = {max(x):.2f}" ) diff --git a/benchmark/hf3fs/bench_zerocopy.py b/benchmark/hf3fs/bench_zerocopy.py index 0f9c5e150..a4ed66912 100644 --- a/benchmark/hf3fs/bench_zerocopy.py +++ b/benchmark/hf3fs/bench_zerocopy.py @@ -109,7 +109,7 @@ elif hicache_mem_layout == "layer_first": for operation in operations: cache_controller.generic_page_backup(operation, batch_size=128) tok = time.monotonic() -print(f"{tok-tik:.6f} s") +print(f"{tok - tik:.6f} s") operations = [ PrefetchOperation( @@ -137,4 +137,4 @@ elif hicache_mem_layout == "layer_first": for operation in operations: cache_controller.generic_page_transfer(operation, batch_size=128) tok = time.monotonic() -print(f"{tok-tik:.6f} s") +print(f"{tok - tik:.6f} s") diff --git a/benchmark/hicache/bench_mix.py b/benchmark/hicache/bench_mix.py index 2a65574ea..5d2615186 100644 --- a/benchmark/hicache/bench_mix.py +++ b/benchmark/hicache/bench_mix.py @@ -457,7 +457,7 @@ class WorkloadGenerator: try: user_data, response = self.response_queue.get(timeout=10) logger.info( - f"{((time.perf_counter()-self.start_time)/self.duration*100):.2f}%" + f"{((time.perf_counter() - self.start_time) / self.duration * 100):.2f}%" ) if not response.success: raise ValueError(f"Request failed with error: {response.error}") @@ -540,10 +540,10 @@ class WorkloadGenerator: output_stats = self.user_generator.output_stats print(f"round_ratios: {user_stats}") print( - f"mean_new_tokens_per_round: {[int(a/b) if b > 0 else 0 for a, b in input_stats]}" + f"mean_new_tokens_per_round: {[int(a / b) if b > 0 else 0 for a, b in input_stats]}" ) print( - f"mean_return_tokens_per_round: {[int(a/b) if b > 0 else 0 for a, b in output_stats]}" + f"mean_return_tokens_per_round: {[int(a / b) if b > 0 else 0 for a, b in output_stats]}" ) return performance_data diff --git a/benchmark/hicache/bench_serving.py b/benchmark/hicache/bench_serving.py index a80c059e5..4e3285858 100644 --- a/benchmark/hicache/bench_serving.py +++ b/benchmark/hicache/bench_serving.py @@ -75,9 +75,9 @@ async def async_request_openai_completions( pbar: Optional[tqdm] = None, ) -> RequestFuncOutput: api_url = request_func_input.api_url - assert api_url.endswith( - "completions" - ), "OpenAI Completions API URL must end with 'completions'." + assert api_url.endswith("completions"), ( + "OpenAI Completions API URL must end with 'completions'." + ) async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: payload = { diff --git a/benchmark/hicache/nextqa.py b/benchmark/hicache/nextqa.py index 4db6caa1c..1c74dac71 100644 --- a/benchmark/hicache/nextqa.py +++ b/benchmark/hicache/nextqa.py @@ -120,7 +120,7 @@ class NExTQALoader(VideoLoader): video = Video(video_path, num_frames) prompt = entry["question"] + "?" if self.task == "MC": # add choices - prompt += f' a0: {entry["a0"]}, a1: {entry["a1"]}, a2: {entry["a2"]}, a3: {entry["a3"]}' + prompt += f" a0: {entry['a0']}, a1: {entry['a1']}, a2: {entry['a2']}, a3: {entry['a3']}" return VideoPrompt(video_path, num_frames, prompt) def __iter__(self): diff --git a/benchmark/kernels/bench_fused_rmsnorm_fp8_quant.py b/benchmark/kernels/bench_fused_rmsnorm_fp8_quant.py index 53e0b8c75..f3230134f 100644 --- a/benchmark/kernels/bench_fused_rmsnorm_fp8_quant.py +++ b/benchmark/kernels/bench_fused_rmsnorm_fp8_quant.py @@ -149,9 +149,9 @@ def _check_correctness(): cos = torch.nn.functional.cosine_similarity( (q.float() * scale).flatten(), ref_deq.flatten(), dim=0 ).item() - assert ( - cos > 0.99 - ), f"{provider} h={hidden_size} residual={add_residual} cos={cos:.4f}" + assert cos > 0.99, ( + f"{provider} h={hidden_size} residual={add_residual} cos={cos:.4f}" + ) print("correctness check passed (all fused providers vs unfused within FP8)") diff --git a/benchmark/kernels/deepep/deepep_utils.py b/benchmark/kernels/deepep/deepep_utils.py index 169529ef6..573d68deb 100644 --- a/benchmark/kernels/deepep/deepep_utils.py +++ b/benchmark/kernels/deepep/deepep_utils.py @@ -191,9 +191,9 @@ def bench_kineto( kernel_names = (kernel_names,) if isinstance(kernel_names, str) else kernel_names assert all([isinstance(name, str) for name in kernel_names]) for name in kernel_names: - assert ( - sum([name in line for line in prof_lines]) == 1 - ), f"Errors of the kernel {name} in the profiling table" + assert sum([name in line for line in prof_lines]) == 1, ( + f"Errors of the kernel {name} in the profiling table" + ) # Save chrome traces if trace_path is not None: diff --git a/benchmark/kernels/deepep/tuning_deepep.py b/benchmark/kernels/deepep/tuning_deepep.py index 191819d2c..6a1401408 100644 --- a/benchmark/kernels/deepep/tuning_deepep.py +++ b/benchmark/kernels/deepep/tuning_deepep.py @@ -155,7 +155,7 @@ def test_main( for with_topk in (False, True): if local_rank == 0: print( - f'[testing] Running with {"FP8" if isinstance(current_x, tuple) else "BF16"}, {"with" if with_topk else "without"} top-k (async={async_mode}, previous={previous_mode}) ...', + f"[testing] Running with {'FP8' if isinstance(current_x, tuple) else 'BF16'}, {'with' if with_topk else 'without'} top-k (async={async_mode}, previous={previous_mode}) ...", flush=True, end="", ) @@ -198,9 +198,9 @@ def test_main( # Checks recv_gbl_rank_prefix_sum = handle[-4] - assert gbl_num_tokens_per_rank[rank].item() == recv_x.size( - 0 - ), f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}" + assert gbl_num_tokens_per_rank[rank].item() == recv_x.size(0), ( + f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}" + ) assert ( gbl_num_tokens_per_expert.view(num_ranks, -1)[rank].tolist() == recv_num_tokens_per_expert_list @@ -325,11 +325,14 @@ def test_main( tune_args = {"x": current_x, "handle": handle, "config": config} t = bench(lambda: buffer.dispatch(**tune_args))[0] if t < best_time: - best_time, best_results = t, ( - num_sms, - nvl_chunk_size, - rdma_chunk_size, - config_kwargs, + best_time, best_results = ( + t, + ( + num_sms, + nvl_chunk_size, + rdma_chunk_size, + config_kwargs, + ), ) if local_rank == 0: print( @@ -338,7 +341,7 @@ def test_main( ) if local_rank == 0: print( - f'[tuning] Best dispatch ({"FP8" if isinstance(current_x, tuple) else "BF16"}): SMs {best_results[0]}, NVL chunk {best_results[1]}, RDMA chunk {best_results[2]}: {rdma_send_bytes / 1e9 / best_time:.2f} GB/s (RDMA), {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)', + f"[tuning] Best dispatch ({'FP8' if isinstance(current_x, tuple) else 'BF16'}): SMs {best_results[0]}, NVL chunk {best_results[1]}, RDMA chunk {best_results[2]}: {rdma_send_bytes / 1e9 / best_time:.2f} GB/s (RDMA), {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)", flush=True, ) print("", flush=True) @@ -399,11 +402,14 @@ def test_main( flush=True, ) if t < best_time: - best_time, best_results = t, ( - num_sms, - nvl_chunk_size, - rdma_chunk_size, - config_kwargs, + best_time, best_results = ( + t, + ( + num_sms, + nvl_chunk_size, + rdma_chunk_size, + config_kwargs, + ), ) if local_rank == 0: diff --git a/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py index 0b958de68..bae6904e2 100644 --- a/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py +++ b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py @@ -59,7 +59,6 @@ def tl_gemm( bx, by, ): - A_shared = T.alloc_shared(A_shared_shape, in_dtype) B_shared = T.alloc_shared(B_shared_shape, in_dtype) C_shared = T.alloc_shared(C_shared_shape, out_dtype) @@ -350,7 +349,7 @@ def get_benchmark(tp_size): tflops = flops / (ms * 1e-3) / 1e12 # Print shape-specific results with TFLOPS - print(f"Time: {ms*1000:.2f} ms, TFLOPS: {tflops:.2f}") + print(f"Time: {ms * 1000:.2f} ms, TFLOPS: {tflops:.2f}") return ms * 1000, max_ms * 1000, min_ms * 1000 # convert to ms return benchmark diff --git a/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm_blackwell.py b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm_blackwell.py index 5670f3c33..280f6b97d 100644 --- a/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm_blackwell.py +++ b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm_blackwell.py @@ -224,7 +224,7 @@ def _benchmark(m, n, k, tp_size, provider): tflops = flops / (ms * 1e-3) / 1e12 # Print shape-specific results with TFLOPS - print(f"Time: {ms*1000:.2f} us, TFLOPS: {tflops:.2f}") + print(f"Time: {ms * 1000:.2f} us, TFLOPS: {tflops:.2f}") return ms, max_ms, min_ms diff --git a/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py index 8b1be7b88..1aace8087 100644 --- a/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py +++ b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py @@ -435,7 +435,7 @@ def get_benchmark(tp_size): flops = 2 * m * n * k # multiply-adds tflops = flops / (ms * 1e-3) / 1e12 - print(f"Time: {ms*1000:.2f} ms, TFLOPS: {tflops:.2f}") + print(f"Time: {ms * 1000:.2f} ms, TFLOPS: {tflops:.2f}") return ms * 1000, max_ms * 1000, min_ms * 1000 # convert to ms return benchmark diff --git a/benchmark/kernels/deepseek/benchmark_q8kv8_kv_gather.py b/benchmark/kernels/deepseek/benchmark_q8kv8_kv_gather.py index 2f885b0ee..12a73faf8 100755 --- a/benchmark/kernels/deepseek/benchmark_q8kv8_kv_gather.py +++ b/benchmark/kernels/deepseek/benchmark_q8kv8_kv_gather.py @@ -243,8 +243,7 @@ def main(): else: speedup = f"{legacy_us / us:.2f}x" print( - f"{s_q:>5} {kv_len:>7} {name:>10} {us:>9.1f} " - f"{tbps:>9.3f} {speedup:>8}" + f"{s_q:>5} {kv_len:>7} {name:>10} {us:>9.1f} {tbps:>9.3f} {speedup:>8}" ) print() diff --git a/benchmark/kernels/elementwise/benchmark_concat_mla.py b/benchmark/kernels/elementwise/benchmark_concat_mla.py index 7bc51d3da..8ae687aad 100644 --- a/benchmark/kernels/elementwise/benchmark_concat_mla.py +++ b/benchmark/kernels/elementwise/benchmark_concat_mla.py @@ -143,9 +143,7 @@ output_exp = execute_and_get_output(fn_cuda, data) if not torch.all(output_ref == output_exp): abs_delta = torch.abs(output_ref - output_exp) raise AssertionError( - f"{output_ref=} {output_exp=} " - f"{abs_delta=} " - f"{torch.argwhere(abs_delta != 0.0)=} " + f"{output_ref=} {output_exp=} {abs_delta=} {torch.argwhere(abs_delta != 0.0)=} " ) diff --git a/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py b/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py index a74972c33..1c5880147 100644 --- a/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py +++ b/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py @@ -535,7 +535,6 @@ class BestConfigTrace: class BenchmarkWorker: - def __init__(self, seed: int, server_args: ServerArgs) -> None: torch.set_default_device("cuda") torch.cuda.manual_seed_all(0) @@ -729,8 +728,7 @@ class BenchmarkWorker: down_use_tma_map[block_m] = time_cost_all[2] > time_cost_all[3] print( - f"Round 1 done. Down TMA decisions per BLOCK_SIZE_M: " - f"{down_use_tma_map}" + f"Round 1 done. Down TMA decisions per BLOCK_SIZE_M: {down_use_tma_map}" ) # === Round 2: Up with c_sorted from round 1 === diff --git a/benchmark/kernels/lora_csgmv/tune_lora_csgmv.py b/benchmark/kernels/lora_csgmv/tune_lora_csgmv.py index aa684dfa8..84a2d8f5a 100755 --- a/benchmark/kernels/lora_csgmv/tune_lora_csgmv.py +++ b/benchmark/kernels/lora_csgmv/tune_lora_csgmv.py @@ -470,9 +470,9 @@ def _tune_shrink( device: torch.device, ) -> tuple: """Tune shrink kernel for one layer type. Returns (best_configs, results).""" - print(f"\n{'='*80}") + print(f"\n{'=' * 80}") print(f"Tuning SHRINK — {label} (K={K}, N={N}, slices={num_slices})") - print(f"{'='*80}") + print(f"{'=' * 80}") search = get_shrink_search_space() print(f"Search space: {len(search)} configs") @@ -508,7 +508,7 @@ def _tune_shrink( best_config = config if (i + 1) % 20 == 0: print( - f" chunk={chunk_size}: {i+1}/{len(search)} tested, best={best_time:.3f}ms" + f" chunk={chunk_size}: {i + 1}/{len(search)} tested, best={best_time:.3f}ms" ) best_configs[chunk_size] = sort_config(best_config) @@ -533,9 +533,9 @@ def _tune_expand( device: torch.device, ) -> tuple: """Tune expand kernel for one layer type. Returns (best_configs, results).""" - print(f"\n{'='*80}") + print(f"\n{'=' * 80}") print(f"Tuning EXPAND — {label} (output_dim={output_dim}, slices={num_slices})") - print(f"{'='*80}") + print(f"{'=' * 80}") search = get_expand_search_space() print(f"Search space: {len(search)} configs") @@ -584,7 +584,7 @@ def _tune_expand( best_config = config if (i + 1) % 50 == 0: print( - f" chunk={chunk_size}: {i+1}/{len(search)} tested, best={best_time:.3f}ms" + f" chunk={chunk_size}: {i + 1}/{len(search)} tested, best={best_time:.3f}ms" ) best_configs[chunk_size] = sort_config(best_config) @@ -673,9 +673,9 @@ def main(args: argparse.Namespace): ) # --- Summary --- - print(f"\n{'='*80}") + print(f"\n{'=' * 80}") print(f"SUMMARY") - print(f"{'='*80}") + print(f"{'=' * 80}") print( f"\n{'layer':<10} {'kernel':<8} {'K/dim':>6} {'chunk':>6}" f" {'baseline':>10} {'tuned':>10} {'speedup':>8} config" diff --git a/benchmark/lean_kernel_sweep.py b/benchmark/lean_kernel_sweep.py index f6a6679ac..6e5ce6a8f 100755 --- a/benchmark/lean_kernel_sweep.py +++ b/benchmark/lean_kernel_sweep.py @@ -137,19 +137,21 @@ def main(): # b32 x 128K on the 8-KV-head config exceeds the microbench's single # contiguous KV tensor (faults the GPU); real serving uses a paged pool. if H_KV == 8 and B == 32 and S == 131072: - print(f"{B:>5} {S//1024:>5}K {'skipped (contiguous-KV limit)':>30}") + print( + f"{B:>5} {S // 1024:>5}K {'skipped (contiguous-KV limit)':>30}" + ) rows.append(f"{name},{H_Q},{H_KV},{B},{S},,,,skip,") continue try: std, lean, cos, gate = run(H_Q, H_KV, B, S) except torch.cuda.OutOfMemoryError: torch.cuda.empty_cache() - print(f"{B:>5} {S//1024:>5}K {'OOM':>9}") + print(f"{B:>5} {S // 1024:>5}K {'OOM':>9}") rows.append(f"{name},{H_Q},{H_KV},{B},{S},,,,OOM,") continue sp = std / lean print( - f"{B:>5} {S//1024:>5}K {std:>9.3f} {lean:>9.3f} {sp:>7.2f}x {('ON' if gate else 'OFF'):>5} {cos:>7.4f}" + f"{B:>5} {S // 1024:>5}K {std:>9.3f} {lean:>9.3f} {sp:>7.2f}x {('ON' if gate else 'OFF'):>5} {cos:>7.4f}" ) rows.append( f"{name},{H_Q},{H_KV},{B},{S},{std:.4f},{lean:.4f},{sp:.4f},{cos:.4f},{int(gate)}" diff --git a/benchmark/mmlu/bench_sglang.py b/benchmark/mmlu/bench_sglang.py index 9a2006e3d..b86aec46c 100644 --- a/benchmark/mmlu/bench_sglang.py +++ b/benchmark/mmlu/bench_sglang.py @@ -163,7 +163,7 @@ def main(args): pt = 0 for subject, num_qs in zip(subjects[: args.nsub], num_questions): print( - f"subject: {subject}, #q:{num_qs}, acc: {np.mean(cors[pt: pt + num_qs]):.3f}" + f"subject: {subject}, #q:{num_qs}, acc: {np.mean(cors[pt : pt + num_qs]):.3f}" ) pt += num_qs assert pt == len(cors) diff --git a/benchmark/ocr/bench_sglang.py b/benchmark/ocr/bench_sglang.py index 651f61417..f210c6232 100644 --- a/benchmark/ocr/bench_sglang.py +++ b/benchmark/ocr/bench_sglang.py @@ -502,7 +502,7 @@ async def process_sample( } ) print( - f"[INPUT ] [{i+1}] type={ttype!r:12s} expected: {expected[:120]}", + f"[INPUT ] [{i + 1}] type={ttype!r:12s} expected: {expected[:120]}", flush=True, ) # Print OCR output (truncate long outputs) diff --git a/benchmark/ocr/generate_report.py b/benchmark/ocr/generate_report.py index 770c827ee..95bba28d6 100644 --- a/benchmark/ocr/generate_report.py +++ b/benchmark/ocr/generate_report.py @@ -159,12 +159,12 @@ def _render_sample(sample: dict, failures_only: bool) -> str: if failures_only and passed == total and not error: return "" - pct = f"{100*passed//total}%" if total else "—" + pct = f"{100 * passed // total}%" if total else "—" header_cls = "fail" if (error or passed < total) else "pass" parts = [f'
'] parts.append( - f'
' + f"
" f'' f"📄 {html.escape(pdf)}  ·  page {page}" f"" @@ -218,7 +218,7 @@ def _render_sample(sample: dict, failures_only: bool) -> str: f'
{_latex_to_display(latex)}
' ) elif ttype in ("present", "absent", "text_presence", "text_absence"): - parts.append(f'
{html.escape(ti.get("text", ""))}
') + parts.append(f"
{html.escape(ti.get('text', ''))}
") elif ttype in ("order", "natural_reading_order"): before = ti.get("before", "") after = ti.get("after", "") @@ -236,7 +236,7 @@ def _render_sample(sample: dict, failures_only: bool) -> str: parts.append(f'
{m}
') if len(matches) > 6: parts.append( - f'

… and {len(matches)-6} more

' + f'

… and {len(matches) - 6} more

' ) else: parts.append( diff --git a/benchmark/prefill_only/util.py b/benchmark/prefill_only/util.py index 2451239d6..ad21dc1b3 100644 --- a/benchmark/prefill_only/util.py +++ b/benchmark/prefill_only/util.py @@ -383,14 +383,14 @@ async def send_warmup_requests( http_url, data=request_json, headers=headers ) as resp: if resp.status == 200: - print(f"Warmup request {i+1}/{num_warmup} completed successfully") + print(f"Warmup request {i + 1}/{num_warmup} completed successfully") else: print( - f"Warmup request {i+1}/{num_warmup} failed with status {resp.status}" + f"Warmup request {i + 1}/{num_warmup} failed with status {resp.status}" ) except Exception as e: - print(f"Warmup request {i+1}/{num_warmup} failed with error: {e}") + print(f"Warmup request {i + 1}/{num_warmup} failed with error: {e}") print("HTTP warmup requests completed") @@ -745,7 +745,6 @@ async def run_generic_benchmark( async with aiohttp.ClientSession( timeout=aiohttp.ClientTimeout(total=300) ) as session: - # Send START_PROFILE if profiling is enabled if config.profile: await send_profile_request("START_PROFILE", http_url, session=session) diff --git a/benchmark/scheduler/bench_token_storage.py b/benchmark/scheduler/bench_token_storage.py index ea8ef7418..79ee6aa6d 100644 --- a/benchmark/scheduler/bench_token_storage.py +++ b/benchmark/scheduler/bench_token_storage.py @@ -254,7 +254,7 @@ def print_breakdown(title: str, results: dict[str, dict[str, float]]) -> None: def microbench_torch_tensor_paths( - sizes: tuple[int, ...] = (1_000, 10_000, 100_000) + sizes: tuple[int, ...] = (1_000, 10_000, 100_000), ) -> None: """Compare three CPU-buffer -> pinned cuda tensor paths. @@ -293,9 +293,11 @@ def microbench_torch_tensor_paths( ), ( "(C) from_numpy(frombuf(array('q'))).pin() -> cuda", - lambda x: torch.from_numpy(np.frombuffer(x, dtype=np.int64)) - .pin_memory() - .to("cuda", non_blocking=True), + lambda x: ( + torch.from_numpy(np.frombuffer(x, dtype=np.int64)) + .pin_memory() + .to("cuda", non_blocking=True) + ), ), ]: cells = [] diff --git a/docs/demo/deepseek_v4_flash.ipynb b/docs/demo/deepseek_v4_flash.ipynb index 5f2f03495..8ffde4076 100644 --- a/docs/demo/deepseek_v4_flash.ipynb +++ b/docs/demo/deepseek_v4_flash.ipynb @@ -347,11 +347,11 @@ " try:\n", " paper[\"full_text\"] = download_and_extract(paper)\n", " print(\n", - " f\"[{i+1}/{N_FULL_PAPERS}] {paper['title'][:70]} — {len(paper['full_text']):,} chars\"\n", + " f\"[{i + 1}/{N_FULL_PAPERS}] {paper['title'][:70]} — {len(paper['full_text']):,} chars\"\n", " )\n", " except Exception as e:\n", " paper[\"full_text\"] = None\n", - " print(f\"[{i+1}/{N_FULL_PAPERS}] {paper['title'][:70]} — failed: {e}\")" + " print(f\"[{i + 1}/{N_FULL_PAPERS}] {paper['title'][:70]} — failed: {e}\")" ] }, { diff --git a/examples/frontend_language/usage/llava_video/srt_example_llava_v.py b/examples/frontend_language/usage/llava_video/srt_example_llava_v.py index ec5b334b0..d042a09da 100644 --- a/examples/frontend_language/usage/llava_video/srt_example_llava_v.py +++ b/examples/frontend_language/usage/llava_video/srt_example_llava_v.py @@ -124,7 +124,6 @@ def batch(video_dir, save_dir, cur_chunk, num_chunks, num_frames=16, batch_size= if __name__ == "__main__": - url = "https://raw.githubusercontent.com/EvolvingLMMs-Lab/sglang/dev/onevision_local/assets/jobs.mp4" cache_dir = os.path.expanduser("~/.cache") diff --git a/examples/frontend_language/usage/readme_examples.py b/examples/frontend_language/usage/readme_examples.py index 7269ef148..0b66e00fc 100644 --- a/examples/frontend_language/usage/readme_examples.py +++ b/examples/frontend_language/usage/readme_examples.py @@ -31,7 +31,7 @@ def tip_suggestion(s): forks = s.fork(2) for i, f in enumerate(forks): - f += f"Now, expand tip {i+1} into a paragraph:\n" + f += f"Now, expand tip {i + 1} into a paragraph:\n" f += sgl.gen(f"detailed_tip", max_tokens=256, stop="\n\n") s += "Tip 1:" + forks[0]["detailed_tip"] + "\n" diff --git a/examples/profiler/nsys_profile_tools/gputrc2graph.py b/examples/profiler/nsys_profile_tools/gputrc2graph.py index 4bfc4340e..68cde5045 100755 --- a/examples/profiler/nsys_profile_tools/gputrc2graph.py +++ b/examples/profiler/nsys_profile_tools/gputrc2graph.py @@ -86,7 +86,7 @@ class GPUTrace2Graph: # Update current_end for overlapping intervals for i in range(1, len(df)): if i % display_units == 0: - print(f"processing trace: {int(i/len(df) * 100)} %", end="\r") + print(f"processing trace: {int(i / len(df) * 100)} %", end="\r") if starts[i] <= current_end: if ends[i] > current_end: # Partial overlap @@ -182,9 +182,9 @@ class GPUTrace2Graph: def is_valid_file(self, base_file): """asserts if base_file is non-existent or is empty""" - assert ( - os.path.isfile(base_file) and os.path.getsize(base_file) > 0 - ), f"{base_file} doesn't exist or is empty" + assert os.path.isfile(base_file) and os.path.getsize(base_file) > 0, ( + f"{base_file} doesn't exist or is empty" + ) def should_gen_file(self, new_file, base_file): """figure out if new file should be generated from base_file""" diff --git a/examples/runtime/engine/fastapi_engine_inference.py b/examples/runtime/engine/fastapi_engine_inference.py index f5da9d715..66eb8ce12 100644 --- a/examples/runtime/engine/fastapi_engine_inference.py +++ b/examples/runtime/engine/fastapi_engine_inference.py @@ -130,7 +130,7 @@ def send_requests(server_url, prompts, max_new_tokens, temperature): """Sends generation requests to the running server for a list of prompts.""" # Iterate through prompts and send requests for i, prompt in enumerate(prompts): - print(f"\n[{i+1}/{len(prompts)}] Sending prompt: '{prompt}'") + print(f"\n[{i + 1}/{len(prompts)}] Sending prompt: '{prompt}'") payload = { "prompt": prompt, "max_new_tokens": max_new_tokens, diff --git a/examples/runtime/engine/offline_batch_inference_qwen_1m.py b/examples/runtime/engine/offline_batch_inference_qwen_1m.py index 664efa6d7..5505bf7b4 100644 --- a/examples/runtime/engine/offline_batch_inference_qwen_1m.py +++ b/examples/runtime/engine/offline_batch_inference_qwen_1m.py @@ -17,8 +17,7 @@ def load_prompt() -> str: # https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/1m.txt with urlopen( - "https://qianwen-res.oss-cn-beijing.aliyuncs.com" - "/Qwen2.5-1M/test-data/64k.txt", + "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-1M/test-data/64k.txt", timeout=5, ) as response: prompt = response.read().decode("utf-8") @@ -41,9 +40,7 @@ def process_requests(llm: sgl.Engine, prompts: list[str]) -> None: for output in outputs: prompt_token_ids = output["meta_info"]["prompt_tokens"] generated_text = output["text"] - print( - f"Prompt length: {prompt_token_ids}, " f"Generated text: {generated_text!r}" - ) + print(f"Prompt length: {prompt_token_ids}, Generated text: {generated_text!r}") # Create an LLM. diff --git a/examples/runtime/qwen3_vl_reranker.py b/examples/runtime/qwen3_vl_reranker.py index 09779996f..b94e27631 100644 --- a/examples/runtime/qwen3_vl_reranker.py +++ b/examples/runtime/qwen3_vl_reranker.py @@ -44,7 +44,7 @@ def rerank_text_only(): print("Results (sorted by relevance):") for i, result in enumerate(results): - print(f" {i+1}. Score: {result['score']:.4f} - {result['document'][:60]}...") + print(f" {i + 1}. Score: {result['score']:.4f} - {result['document'][:60]}...") print() @@ -99,7 +99,7 @@ def rerank_with_images(): print("Results (sorted by relevance):") for i, result in enumerate(results): - print(f" {i+1}. Index: {result['index']}, Score: {result['score']:.4f}") + print(f" {i + 1}. Index: {result['index']}, Score: {result['score']:.4f}") print() @@ -149,7 +149,7 @@ def rerank_multimodal_query(): print("Results (sorted by relevance):") for i, result in enumerate(results): - print(f" {i+1}. Index: {result['index']}, Score: {result['score']:.4f}") + print(f" {i + 1}. Index: {result['index']}, Score: {result['score']:.4f}") print() diff --git a/examples/usage/modelopt_quantize_and_export.py b/examples/usage/modelopt_quantize_and_export.py index 4394d917c..b3f26ede4 100755 --- a/examples/usage/modelopt_quantize_and_export.py +++ b/examples/usage/modelopt_quantize_and_export.py @@ -213,7 +213,7 @@ def deploy_exported_model( outputs = llm.generate(prompts, sampling_params) for i, output in enumerate(outputs): - print(f"Prompt {i+1}: {prompts[i]}") + print(f"Prompt {i + 1}: {prompts[i]}") print(f"Output: {output['text']}") print() diff --git a/examples/usage/reasoning_aware_compression/rac_collect_traces.py b/examples/usage/reasoning_aware_compression/rac_collect_traces.py index a9be8f69e..4c83b1618 100755 --- a/examples/usage/reasoning_aware_compression/rac_collect_traces.py +++ b/examples/usage/reasoning_aware_compression/rac_collect_traces.py @@ -335,7 +335,7 @@ def report(manifest: TraceManifest) -> None: print(f" decode tokens : {manifest.num_decode_tokens}") print(f" total tokens : {total}") print(f" decode share : {decode_share:.1%}") - print(f" wall clock : {manifest.elapsed_seconds/60:.1f} min") + print(f" wall clock : {manifest.elapsed_seconds / 60:.1f} min") if manifest.calibration_mode == "rac": print( "\nThe decode share is the activation mass that prompt-only " diff --git a/examples/usage/reasoning_aware_compression/rac_serve_and_eval.py b/examples/usage/reasoning_aware_compression/rac_serve_and_eval.py index 7448cf4db..27a70ab0a 100755 --- a/examples/usage/reasoning_aware_compression/rac_serve_and_eval.py +++ b/examples/usage/reasoning_aware_compression/rac_serve_and_eval.py @@ -209,7 +209,7 @@ def report(results: List[EvalResult]) -> None: f"{result.model_path:<{width}} " f"{result.accuracy:>7.3f} " f"{result.mean_completion_tokens:>16.0f} " - f"{result.elapsed_seconds/60:>10.1f}m" + f"{result.elapsed_seconds / 60:>10.1f}m" ) if len(results) > 1: diff --git a/experimental/sgl-router/tests/e2e/chat_completions/test_two_router_convergence.py b/experimental/sgl-router/tests/e2e/chat_completions/test_two_router_convergence.py index 5534b428d..e92bbdb22 100644 --- a/experimental/sgl-router/tests/e2e/chat_completions/test_two_router_convergence.py +++ b/experimental/sgl-router/tests/e2e/chat_completions/test_two_router_convergence.py @@ -194,9 +194,9 @@ def _direct_warm(worker_url: str, model_id: str, prefix: str) -> None: }, timeout=60.0, ) - assert ( - r.status_code == 200 - ), f"direct warm to {worker_url} failed: HTTP {r.status_code} {r.text!r}" + assert r.status_code == 200, ( + f"direct warm to {worker_url} failed: HTTP {r.status_code} {r.text!r}" + ) def _route_through(router_url: str, model_id: str, prompt: str) -> str: @@ -212,9 +212,9 @@ def _route_through(router_url: str, model_id: str, prompt: str) -> str: after = _success_counts_by_worker(router_url) deltas = {w: after.get(w, 0) - before.get(w, 0) for w in set(after) | set(before)} winners = [w for w, d in deltas.items() if d > 0] - assert ( - len(winners) == 1 - ), f"expected exactly one worker delta on {router_url}, got {deltas}" + assert len(winners) == 1, ( + f"expected exactly one worker delta on {router_url}, got {deltas}" + ) return winners[0] @@ -309,15 +309,15 @@ def test_routers_route_by_prefix_content( landed = _route_through( router.base_url, spec["model"], PREFIX_X ) - assert ( - landed == worker_x.url - ), f"router {label}: PREFIX_X must route to {worker_x.url}; landed on {landed}" + assert landed == worker_x.url, ( + f"router {label}: PREFIX_X must route to {worker_x.url}; landed on {landed}" + ) landed = _route_through( router.base_url, spec["model"], PREFIX_Y ) - assert ( - landed == worker_y.url - ), f"router {label}: PREFIX_Y must route to {worker_y.url}; landed on {landed}" + assert landed == worker_y.url, ( + f"router {label}: PREFIX_Y must route to {worker_y.url}; landed on {landed}" + ) except Exception: _dump_logs(logs) raise diff --git a/experimental/sgl-router/tests/e2e/chat_completions/test_validation.py b/experimental/sgl-router/tests/e2e/chat_completions/test_validation.py index 4b9ec369a..5216831fd 100644 --- a/experimental/sgl-router/tests/e2e/chat_completions/test_validation.py +++ b/experimental/sgl-router/tests/e2e/chat_completions/test_validation.py @@ -50,9 +50,9 @@ def test_chat_non_streaming_returns_assistant_message( body = resp.json() choice = body["choices"][0] assert choice["message"]["role"] == "assistant" - assert choice["message"][ - "content" - ], f"empty assistant content: {choice!r}" + assert choice["message"]["content"], ( + f"empty assistant content: {choice!r}" + ) assert choice.get("finish_reason"), choice finally: gpu_allocator.release(gpu) @@ -91,8 +91,8 @@ def test_chat_streaming_emits_sse_chunks_with_done( if line.startswith("data:"): chunks.append(line.strip()) assert len(chunks) >= 2, f"expected >=2 SSE chunks, got: {chunks}" - assert any( - "[DONE]" in c for c in chunks - ), f"no [DONE] terminator in stream: {chunks}" + assert any("[DONE]" in c for c in chunks), ( + f"no [DONE] terminator in stream: {chunks}" + ) finally: gpu_allocator.release(gpu) diff --git a/experimental/sgl-router/tests/e2e/k8s_integration/test_discovery.py b/experimental/sgl-router/tests/e2e/k8s_integration/test_discovery.py index 83f0faab0..400d5cd41 100644 --- a/experimental/sgl-router/tests/e2e/k8s_integration/test_discovery.py +++ b/experimental/sgl-router/tests/e2e/k8s_integration/test_discovery.py @@ -66,15 +66,17 @@ def test_router_discovers_multiple_workers(router_url): # Scale down to 1 — router should still route after reconverging _scale_fake_worker(1) _poll_until( - lambda: httpx.post( - f"{router_url}/v1/chat/completions", - json={ - "model": "tiny", - "messages": [{"role": "user", "content": "post-scale-down"}], - }, - timeout=10.0, - ).status_code - == 200, + lambda: ( + httpx.post( + f"{router_url}/v1/chat/completions", + json={ + "model": "tiny", + "messages": [{"role": "user", "content": "post-scale-down"}], + }, + timeout=10.0, + ).status_code + == 200 + ), "router routes after scale-down to 1", timeout=60, interval=3, diff --git a/experimental/sgl-router/tests/e2e/test_chat_smoke.py b/experimental/sgl-router/tests/e2e/test_chat_smoke.py index 16d030bd3..82847e1f9 100644 --- a/experimental/sgl-router/tests/e2e/test_chat_smoke.py +++ b/experimental/sgl-router/tests/e2e/test_chat_smoke.py @@ -16,9 +16,9 @@ def test_models(router: str) -> None: assert resp.status_code == 200, resp.text data = resp.json() ids = [m["id"] for m in data.get("data", [])] - assert any( - MODEL in mid for mid in ids - ), f"Model {MODEL!r} not found in /v1/models response: {ids}" + assert any(MODEL in mid for mid in ids), ( + f"Model {MODEL!r} not found in /v1/models response: {ids}" + ) def test_chat_non_streaming(router: str) -> None: @@ -59,6 +59,6 @@ def test_chat_streaming(router: str) -> None: chunks.append(line) assert len(chunks) >= 2, f"Expected >=2 SSE chunks, got {len(chunks)}: {chunks}" - assert any( - "[DONE]" in c for c in chunks - ), f"No [DONE] chunk found in SSE stream: {chunks}" + assert any("[DONE]" in c for c in chunks), ( + f"No [DONE] chunk found in SSE stream: {chunks}" + ) diff --git a/experimental/sgl-router/tests/e2e/test_tokenize_smoke.py b/experimental/sgl-router/tests/e2e/test_tokenize_smoke.py index cee049f87..1e9babf6e 100644 --- a/experimental/sgl-router/tests/e2e/test_tokenize_smoke.py +++ b/experimental/sgl-router/tests/e2e/test_tokenize_smoke.py @@ -20,9 +20,9 @@ def test_tokenize_round_trip(router: str) -> None: ) assert tok_resp.status_code == 200, tok_resp.text tokens = tok_resp.json()["tokens"] - assert ( - isinstance(tokens, list) and len(tokens) > 0 - ), f"Expected non-empty token list, got: {tokens}" + assert isinstance(tokens, list) and len(tokens) > 0, ( + f"Expected non-empty token list, got: {tokens}" + ) # Detokenize detok_resp = httpx.post( @@ -32,6 +32,6 @@ def test_tokenize_round_trip(router: str) -> None: ) assert detok_resp.status_code == 200, detok_resp.text recovered = detok_resp.json()["text"] - assert ( - TEXT in recovered or recovered in TEXT - ), f"Round-trip mismatch: original={TEXT!r}, recovered={recovered!r}" + assert TEXT in recovered or recovered in TEXT, ( + f"Round-trip mismatch: original={TEXT!r}, recovered={recovered!r}" + ) diff --git a/experimental/sgl-router/tests/scripts/generate_parity_fixtures.py b/experimental/sgl-router/tests/scripts/generate_parity_fixtures.py index 5d33813f2..4e985823f 100644 --- a/experimental/sgl-router/tests/scripts/generate_parity_fixtures.py +++ b/experimental/sgl-router/tests/scripts/generate_parity_fixtures.py @@ -85,7 +85,7 @@ def load_tokenizer_with_fallback(primary, fallback, slug): raise continue raise RuntimeError( - f"No accessible tokenizer for slug={slug} " f"(tried: {primary}, {fallback})" + f"No accessible tokenizer for slug={slug} (tried: {primary}, {fallback})" ) diff --git a/python/sglang/benchmark/dspark_sps_profiler.py b/python/sglang/benchmark/dspark_sps_profiler.py index 1f914fa7c..88a3a13ca 100644 --- a/python/sglang/benchmark/dspark_sps_profiler.py +++ b/python/sglang/benchmark/dspark_sps_profiler.py @@ -687,8 +687,7 @@ def run_one_round( rank_rows = fetch_rank_rows(base_url=context.base_url) if len(rank_rows) != len(watermarks): raise RuntimeError( - f"DP rank count changed mid-profile: {len(watermarks)} -> " - f"{len(rank_rows)}." + f"DP rank count changed mid-profile: {len(watermarks)} -> {len(rank_rows)}." ) new_rank_rows = [ [row for row in rows if row.forward_ct > watermark] diff --git a/python/sglang/benchmark/offline_throughput.py b/python/sglang/benchmark/offline_throughput.py index 37acaad6f..0114c8a15 100644 --- a/python/sglang/benchmark/offline_throughput.py +++ b/python/sglang/benchmark/offline_throughput.py @@ -136,13 +136,13 @@ class BenchArgs: "--gsp-system-prompt-len", type=int, default=BenchArgs.gsp_system_prompt_len, - help="System prompt length, used" "only for generate-shared-prefix", + help="System prompt length, usedonly for generate-shared-prefix", ) parser.add_argument( "--gsp-question-len", type=int, default=BenchArgs.gsp_question_len, - help="Question length, used" "only for generate-shared-prefix", + help="Question length, usedonly for generate-shared-prefix", ) parser.add_argument( "--gsp-output-len", @@ -259,9 +259,9 @@ def throughput_test_once( ] if profile: - assert ( - "SGLANG_TORCH_PROFILER_DIR" in os.environ - ), "Please set SGLANG_TORCH_PROFILER_DIR." + assert "SGLANG_TORCH_PROFILER_DIR" in os.environ, ( + "Please set SGLANG_TORCH_PROFILER_DIR." + ) os.makedirs(os.environ["SGLANG_TORCH_PROFILER_DIR"], exist_ok=True) known_files = None backend.start_profile( diff --git a/python/sglang/benchmark/one_batch_server.py b/python/sglang/benchmark/one_batch_server.py index d904992c9..27a426336 100644 --- a/python/sglang/benchmark/one_batch_server.py +++ b/python/sglang/benchmark/one_batch_server.py @@ -486,7 +486,7 @@ def _warmup_cache( return print( - f"Warming up cache with {cache_hit_rate*100:.1f}% hit rate " + f"Warming up cache with {cache_hit_rate * 100:.1f}% hit rate " f"({cached_token_len} tokens per request)" ) # Create prefix input_ids for cache warming @@ -1024,7 +1024,7 @@ def get_report_summary( f"\nInput lens: {bench_args.input_len}. Output lens: {bench_args.output_len}." ) if bench_args.cache_hit_rate > 0.0: - summary += f" Cache hit rate: {bench_args.cache_hit_rate*100:.1f}%." + summary += f" Cache hit rate: {bench_args.cache_hit_rate * 100:.1f}%." summary += "\n" if is_blackwell(): @@ -1241,9 +1241,9 @@ def run_benchmark_internal( skip_max_running_requests_threshold = float("inf") skip_token_capacity_threshold = float("inf") else: - assert ( - max_running_requests_per_dp > 0 - ), f"effective_max_running_requests_per_dp is not set, {max_running_requests_per_dp=}" + assert max_running_requests_per_dp > 0, ( + f"effective_max_running_requests_per_dp is not set, {max_running_requests_per_dp=}" + ) skip_max_running_requests_threshold = max_running_requests_per_dp * dp_size print(f"{max_running_requests_per_dp=}") @@ -1288,9 +1288,9 @@ def run_benchmark_internal( "--lora-request-distribution=distinct/skewed requires more than " "one adapter via --lora-name." ) - assert ( - bench_args.lora_zipf_alpha > 1 - ), f"--lora-zipf-alpha must be > 1, got {bench_args.lora_zipf_alpha}" + assert bench_args.lora_zipf_alpha > 1, ( + f"--lora-zipf-alpha must be > 1, got {bench_args.lora_zipf_alpha}" + ) if bench_args.apply_chat_template and not ( bench_args.fixed_prompt_file or bench_args.dataset_name in REPLAY_TEXT_DATASETS diff --git a/python/sglang/benchmark/serving.py b/python/sglang/benchmark/serving.py index 9ba76b1cb..4362619d5 100644 --- a/python/sglang/benchmark/serving.py +++ b/python/sglang/benchmark/serving.py @@ -261,9 +261,9 @@ async def async_request_openai_completions( pbar: Optional[tqdm] = None, ) -> RequestFuncOutput: api_url = request_func_input.api_url - assert api_url.endswith( - "completions" - ), "OpenAI Completions API URL must end with 'completions'." + assert api_url.endswith("completions"), ( + "OpenAI Completions API URL must end with 'completions'." + ) prompt = request_func_input.prompt @@ -392,9 +392,9 @@ async def async_request_openai_chat_completions( latency, TTFT, ITL, and success status. """ api_url = request_func_input.api_url - assert api_url.endswith( - "chat/completions" - ), "OpenAI Chat Completions API URL must end with 'chat/completions'." + assert api_url.endswith("chat/completions"), ( + "OpenAI Chat Completions API URL must end with 'chat/completions'." + ) # TODO put it to other functions when `pbar` logic is refactored if getattr(args, "print_requests", False): @@ -1296,9 +1296,9 @@ def _normalize_round_messages(turn: Any) -> Optional[List[Dict[str, str]]]: def wrap_multi_turn_request_func(request_func: Callable, backend: str) -> Callable: - assert ( - backend in MULTI_TURN_BACKENDS - ), f"Multi-turn only supports chat backends: {MULTI_TURN_BACKENDS}, got {backend}" + assert backend in MULTI_TURN_BACKENDS, ( + f"Multi-turn only supports chat backends: {MULTI_TURN_BACKENDS}, got {backend}" + ) async def f( request_func_input: RequestFuncInput, @@ -1534,9 +1534,9 @@ async def benchmark( lora_name = lora_names[lora_idx] lora_idx = (lora_idx + 1) % len(lora_names) else: - assert ( - lora_request_distribution == "skewed" - ), f"Unexpected lora_request_distribution: {lora_request_distribution}. Expected 'skewed'." + assert lora_request_distribution == "skewed", ( + f"Unexpected lora_request_distribution: {lora_request_distribution}. Expected 'skewed'." + ) lora_name = np.random.choice(lora_names, p=lora_probs) else: @@ -2000,9 +2000,9 @@ def run_benchmark(args_: argparse.Namespace): extra_request_body["bootstrap_room"] = 0 if args.tokenize_prompt: - assert ( - args.backend == "sglang" - ), "`--tokenize-prompt` only compatible with `--backend sglang` currently" + assert args.backend == "sglang", ( + "`--tokenize-prompt` only compatible with `--backend sglang` currently" + ) # Set url if args.port is None: @@ -2079,18 +2079,18 @@ def run_benchmark(args_: argparse.Namespace): if args.dataset_name in ["image", "mmmu"]: args.apply_chat_template = True - assert ( - not args.tokenize_prompt - ), "`--tokenize-prompt` not compatible with image dataset" + assert not args.tokenize_prompt, ( + "`--tokenize-prompt` not compatible with image dataset" + ) if args.lora_request_distribution in ["distinct", "skewed"]: - assert ( - args.lora_name is not None and len(args.lora_name) > 1 - ), "More than 1 LoRA adapter must be specified via --lora-name to use 'distinct' or 'skewed' request distribution." + assert args.lora_name is not None and len(args.lora_name) > 1, ( + "More than 1 LoRA adapter must be specified via --lora-name to use 'distinct' or 'skewed' request distribution." + ) - assert ( - args.lora_zipf_alpha > 1 - ), f"Got invalid value for --lora-zipf-alpha of {args.lora_zipf_alpha}. It must be greater than 1." + assert args.lora_zipf_alpha > 1, ( + f"Got invalid value for --lora-zipf-alpha of {args.lora_zipf_alpha}. It must be greater than 1." + ) print(f"{args}\n") @@ -2364,13 +2364,13 @@ def cli_main(): "--image-format", type=str, default="jpeg", - help=("Format of images for image dataset. " "Supports jpeg and png."), + help=("Format of images for image dataset. Supports jpeg and png."), ) parser.add_argument( "--image-content", type=str, default="random", - help=("Content for images for image dataset. " "Supports random and blank."), + help=("Content for images for image dataset. Supports random and blank."), ) parser.add_argument( "--request-rate", diff --git a/python/sglang/cli/killall.py b/python/sglang/cli/killall.py index 1e672df2c..631988416 100755 --- a/python/sglang/cli/killall.py +++ b/python/sglang/cli/killall.py @@ -315,8 +315,7 @@ def _print_diagnostics(unkillable_pids): print(f" {line}") else: print( - "\n[killall] Diagnostic — no sglang/python/gpu processes " - "in this container" + "\n[killall] Diagnostic — no sglang/python/gpu processes in this container" ) diff --git a/python/sglang/cli/serve.py b/python/sglang/cli/serve.py index 3a934ada5..6e913622d 100644 --- a/python/sglang/cli/serve.py +++ b/python/sglang/cli/serve.py @@ -195,7 +195,7 @@ def serve(args, extra_argv): else: registered = registry.get(backend_name) logger.info( - "Dispatch override enabled: --model-type=%s " "(skip auto detection)", + "Dispatch override enabled: --model-type=%s (skip auto detection)", backend_name, ) diff --git a/python/sglang/cli/utils.py b/python/sglang/cli/utils.py index 26258be2b..f27ad6cf0 100644 --- a/python/sglang/cli/utils.py +++ b/python/sglang/cli/utils.py @@ -126,8 +126,7 @@ def get_model_path(extra_argv): ) else: raise Exception( - "Error: --model-path is required. " - "Please provide the path to the model." + "Error: --model-path is required. Please provide the path to the model." ) return model_path diff --git a/python/sglang/kernels/aot/benchmark/bench_amd_deterministic_allreduce.py b/python/sglang/kernels/aot/benchmark/bench_amd_deterministic_allreduce.py index 234d09342..60b658698 100644 --- a/python/sglang/kernels/aot/benchmark/bench_amd_deterministic_allreduce.py +++ b/python/sglang/kernels/aot/benchmark/bench_amd_deterministic_allreduce.py @@ -245,7 +245,7 @@ def worker(world_size, rank, port, results_queue): if input_size_bytes > custom_ar.max_size: if rank == 0: print( - f" Deterministic kernel skipped: input size ({input_size_bytes/(1024*1024):.1f} MB) > buffer size ({custom_ar.max_size/(1024*1024):.1f} MB)" + f" Deterministic kernel skipped: input size ({input_size_bytes / (1024 * 1024):.1f} MB) > buffer size ({custom_ar.max_size / (1024 * 1024):.1f} MB)" ) deterministic_kernel_available = False else: @@ -412,17 +412,17 @@ def worker(world_size, rank, port, results_queue): } print( - f" All-Reduce: {lat_ar_median*1000:.3f}ms, Deterministic: {ar_deterministic}, Max variance: {ar_max_variance:.6f}" + f" All-Reduce: {lat_ar_median * 1000:.3f}ms, Deterministic: {ar_deterministic}, Max variance: {ar_max_variance:.6f}" ) print( - f" RS+All-Gather: {lat_rs_ag_median*1000:.3f}ms, Deterministic: {rs_ag_deterministic}, Max variance: {rs_ag_max_variance:.6f}" + f" RS+All-Gather: {lat_rs_ag_median * 1000:.3f}ms, Deterministic: {rs_ag_deterministic}, Max variance: {rs_ag_max_variance:.6f}" ) if custom_ar is not None and lat_custom_ar_median is not None: overhead_custom = ( (lat_custom_ar_median - lat_ar_median) / lat_ar_median ) * 100 print( - f" Custom AR: {lat_custom_ar_median*1000:.3f}ms, Deterministic: {custom_ar_deterministic}, Max variance: {custom_ar_max_variance:.6f}, Overhead: {overhead_custom:+.1f}%" + f" Custom AR: {lat_custom_ar_median * 1000:.3f}ms, Deterministic: {custom_ar_deterministic}, Max variance: {custom_ar_max_variance:.6f}, Overhead: {overhead_custom:+.1f}%" ) if lat_deterministic_kernel_median is not None: overhead_kernel = ( @@ -433,7 +433,7 @@ def worker(world_size, rank, port, results_queue): / lat_rs_ag_median ) * 100 print( - f" Deterministic Kernel: {lat_deterministic_kernel_median*1000:.3f}ms, Deterministic: {deterministic_kernel_deterministic}, Max variance: {deterministic_kernel_max_variance:.6f}, Overhead: {overhead_kernel:+.1f}%, Speedup vs RS+AG: {speedup_kernel_vs_rs_ag:+.1f}%" + f" Deterministic Kernel: {lat_deterministic_kernel_median * 1000:.3f}ms, Deterministic: {deterministic_kernel_deterministic}, Max variance: {deterministic_kernel_max_variance:.6f}, Overhead: {overhead_kernel:+.1f}%, Speedup vs RS+AG: {speedup_kernel_vs_rs_ag:+.1f}%" ) if lat_optimized_rs_ag_median is not None: overhead_opt = ( @@ -443,7 +443,7 @@ def worker(world_size, rank, port, results_queue): (lat_rs_ag_median - lat_optimized_rs_ag_median) / lat_rs_ag_median ) * 100 print( - f" Optimized RS+AG: {lat_optimized_rs_ag_median*1000:.3f}ms, Deterministic: {optimized_rs_ag_deterministic}, Max variance: {optimized_rs_ag_max_variance:.6f}, Overhead: {overhead_opt:+.1f}%, Speedup vs RS+AG: {speedup_vs_rs_ag:+.1f}%" + f" Optimized RS+AG: {lat_optimized_rs_ag_median * 1000:.3f}ms, Deterministic: {optimized_rs_ag_deterministic}, Max variance: {optimized_rs_ag_max_variance:.6f}, Overhead: {overhead_opt:+.1f}%, Speedup vs RS+AG: {speedup_vs_rs_ag:+.1f}%" ) print(f" RS+AG Overhead: {overhead_rs_ag:+.1f}%") @@ -515,8 +515,8 @@ def main(): ar_det_str = "✓" if r["all_reduce"]["deterministic"] else "✗" rs_ag_det_str = "✓" if r["rs_ag"]["deterministic"] else "✗" line = ( - f"{bs:<8} {r['all_reduce']['latency_median']*1000:<12.3f} {ar_det_str:<8} " - f"{r['rs_ag']['latency_median']*1000:<15.3f} {rs_ag_det_str:<10} " + f"{bs:<8} {r['all_reduce']['latency_median'] * 1000:<12.3f} {ar_det_str:<8} " + f"{r['rs_ag']['latency_median'] * 1000:<15.3f} {rs_ag_det_str:<10} " f"{r['overhead_rs_ag_pct']:<12.1f}" ) if r.get("custom_ar") is not None: @@ -526,7 +526,7 @@ def main(): (custom_ar["latency_median"] - r["all_reduce"]["latency_median"]) / r["all_reduce"]["latency_median"] ) * 100 - line += f" {custom_ar['latency_median']*1000:<18.3f} {custom_ar_det_str:<15} {custom_ar_overhead:<15.1f}" + line += f" {custom_ar['latency_median'] * 1000:<18.3f} {custom_ar_det_str:<15} {custom_ar_overhead:<15.1f}" if r.get("deterministic_kernel") is not None: det_kernel = r["deterministic_kernel"] det_kernel_det_str = "✓" if det_kernel["deterministic"] else "✗" @@ -538,7 +538,7 @@ def main(): (r["rs_ag"]["latency_median"] - det_kernel["latency_median"]) / r["rs_ag"]["latency_median"] ) * 100 - line += f" {det_kernel['latency_median']*1000:<18.3f} {det_kernel_det_str:<15} {det_kernel_overhead:<15.1f} {speedup_kernel:<10.1f}" + line += f" {det_kernel['latency_median'] * 1000:<18.3f} {det_kernel_det_str:<15} {det_kernel_overhead:<15.1f} {speedup_kernel:<10.1f}" if r.get("optimized_rs_ag") is not None: opt_rs_ag = r["optimized_rs_ag"] opt_rs_ag_det_str = "✓" if opt_rs_ag["deterministic"] else "✗" @@ -550,7 +550,7 @@ def main(): (r["rs_ag"]["latency_median"] - opt_rs_ag["latency_median"]) / r["rs_ag"]["latency_median"] ) * 100 - line += f" {opt_rs_ag['latency_median']*1000:<18.3f} {opt_rs_ag_det_str:<15} {opt_rs_ag_overhead:<15.1f} {speedup:<10.1f}" + line += f" {opt_rs_ag['latency_median'] * 1000:<18.3f} {opt_rs_ag_det_str:<15} {opt_rs_ag_overhead:<15.1f} {speedup:<10.1f}" print(line) print("=" * 80) diff --git a/python/sglang/kernels/aot/benchmark/bench_cutlass_mla.py b/python/sglang/kernels/aot/benchmark/bench_cutlass_mla.py index 461d8862b..9d24443f6 100644 --- a/python/sglang/kernels/aot/benchmark/bench_cutlass_mla.py +++ b/python/sglang/kernels/aot/benchmark/bench_cutlass_mla.py @@ -114,10 +114,8 @@ def benchmark(batch_size, seq_len, provider, block_size, num_kv_splits): q_size = qn.numel() * qn.element_size() + qr.numel() * qr.element_size() - gbps = ( - lambda ms: ( - q_size + q_size * dv / d + kv_cache.numel() * kv_cache.element_size() - ) + gbps = lambda ms: ( + (q_size + q_size * dv / d + kv_cache.numel() * kv_cache.element_size()) * 1e-9 / (ms * 1e-3) ) diff --git a/python/sglang/kernels/aot/benchmark/bench_fp4_gemm.py b/python/sglang/kernels/aot/benchmark/bench_fp4_gemm.py index df536f508..26d8aad30 100755 --- a/python/sglang/kernels/aot/benchmark/bench_fp4_gemm.py +++ b/python/sglang/kernels/aot/benchmark/bench_fp4_gemm.py @@ -368,9 +368,9 @@ def benchmark(batch_size, provider, N, K, dtype, correctness, csv_file): res_fi, backend="cudnn", ) - assert torch.allclose( - res_fi, res_cutlass, atol=1e-3, rtol=1e-3 - ), "cudnn fp4 doesn't match cutlass fp4" + assert torch.allclose(res_fi, res_cutlass, atol=1e-3, rtol=1e-3), ( + "cudnn fp4 doesn't match cutlass fp4" + ) mm_fp4( a_fp4, b_fp4_T, @@ -381,9 +381,9 @@ def benchmark(batch_size, provider, N, K, dtype, correctness, csv_file): res_fi, backend="trtllm", ) - assert torch.allclose( - res_fi, res_cutlass, atol=1e-3, rtol=1e-3 - ), "trtllm fp4 doesn't match cutlass fp4" + assert torch.allclose(res_fi, res_cutlass, atol=1e-3, rtol=1e-3), ( + "trtllm fp4 doesn't match cutlass fp4" + ) if csv_file: with open(csv_file, "a", newline="") as f: diff --git a/python/sglang/kernels/aot/benchmark/bench_int8_gemm.py b/python/sglang/kernels/aot/benchmark/bench_int8_gemm.py index 722d89d7c..0d0cb2a15 100644 --- a/python/sglang/kernels/aot/benchmark/bench_int8_gemm.py +++ b/python/sglang/kernels/aot/benchmark/bench_int8_gemm.py @@ -127,8 +127,8 @@ def benchmark(batch_size, provider, N, K): lambda: vllm_scaled_mm(a, b, scale_a, scale_b, torch.float16, bias), quantiles=quantiles, ) - gbps = ( - lambda ms: ( + gbps = lambda ms: ( + ( (2 * M * N * K - M * N) * a.element_size() + (3 * M * N) * scale_a.element_size() ) diff --git a/python/sglang/kernels/aot/python/sgl_kernel/attention.py b/python/sglang/kernels/aot/python/sgl_kernel/attention.py index faf23a4f0..3351d6d6f 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/attention.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/attention.py @@ -38,9 +38,9 @@ def cutlass_mla_decode( ) -> torch.Tensor: assert q_nope.ndim == 3, f"q_nope must be a 3D tensor, but got {q_nope.ndim}" assert q_pe.ndim == 3, f"q_pe must be a 3D tensor, but got {q_pe.ndim}" - assert ( - kv_c_and_k_pe_cache.ndim == 3 - ), f"kv_c_and_k_pe_cache must be a 3D tensor, but got {kv_c_and_k_pe_cache.ndim}" + assert kv_c_and_k_pe_cache.ndim == 3, ( + f"kv_c_and_k_pe_cache must be a 3D tensor, but got {kv_c_and_k_pe_cache.ndim}" + ) B_q, H, D_q_nope = q_nope.shape B_q_2, H_2, D_q_pe = q_pe.shape @@ -77,12 +77,12 @@ def cutlass_mla_decode( torch.bfloat16, ), f"q_nope.dtype needs to be fp16 or bf16 but got {q_nope.dtype}." assert q_nope.dtype == q_pe.dtype == kv_c_and_k_pe_cache.dtype - assert ( - seq_lens.dtype == torch.int32 - ), f"seq_lens.dtype needs to be int32 but got {seq_lens.dtype}." - assert ( - page_table.dtype == torch.int32 - ), f"page_table.dtype needs to be int32 but got {page_table.dtype}." + assert seq_lens.dtype == torch.int32, ( + f"seq_lens.dtype needs to be int32 but got {seq_lens.dtype}." + ) + assert page_table.dtype == torch.int32, ( + f"page_table.dtype needs to be int32 but got {page_table.dtype}." + ) out = q_nope.new_empty((B_q, MAX_HEADS, D_latent)) diff --git a/python/sglang/kernels/aot/python/sgl_kernel/elementwise.py b/python/sglang/kernels/aot/python/sgl_kernel/elementwise.py index aa325b277..d2315761c 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/elementwise.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/elementwise.py @@ -247,12 +247,12 @@ def gemma_fused_add_rmsnorm( def _check_shape(input: torch.Tensor, output: torch.Tensor) -> None: assert input.ndim == output.ndim, f"{input.ndim} != {output.ndim}" - assert ( - input.shape[:-1] == output.shape[:-1] - ), f"{input.shape[:-1]} != {output.shape[:-1]}" - assert ( - input.shape[-1] == 2 * output.shape[-1] - ), f"{input.shape[-1]} != {2 * output.shape[-1]}" + assert input.shape[:-1] == output.shape[:-1], ( + f"{input.shape[:-1]} != {output.shape[:-1]}" + ) + assert input.shape[-1] == 2 * output.shape[-1], ( + f"{input.shape[-1]} != {2 * output.shape[-1]}" + ) def silu_and_mul(input: torch.Tensor, out: torch.Tensor = None) -> torch.Tensor: diff --git a/python/sglang/kernels/aot/python/sgl_kernel/flash_mla.py b/python/sglang/kernels/aot/python/sgl_kernel/flash_mla.py index bd8d558a6..2a15aad9c 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/flash_mla.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/flash_mla.py @@ -159,9 +159,9 @@ def flash_mla_with_kvcache( assert extra_topk_length is None if indices is not None: assert causal == False, "causal must be `false` if sparse attention is enabled." - assert (descale_q is None) == ( - descale_k is None - ), "descale_q and descale_k should be both None or both not None" + assert (descale_q is None) == (descale_k is None), ( + "descale_q and descale_k should be both None or both not None" + ) if indices is None and q.element_size() == 1: out, softmax_lse = torch.ops.sgl_kernel.fwd_kvcache_mla_fp8.default( @@ -257,9 +257,9 @@ def _flash_mla_with_kvcache_sched_meta( assert sched_meta.config.causal == causal, helper_msg assert sched_meta.config.is_fp8_kvcache == is_fp8_kvcache, helper_msg assert sched_meta.config.topk == topk, helper_msg - assert ( - sched_meta.config.extra_page_block_size == extra_page_block_size - ), helper_msg + assert sched_meta.config.extra_page_block_size == extra_page_block_size, ( + helper_msg + ) assert sched_meta.config.extra_topk == extra_topk, helper_msg if topk is not None: diff --git a/python/sglang/kernels/aot/python/sgl_kernel/metal.py b/python/sglang/kernels/aot/python/sgl_kernel/metal.py index 8edec27c8..a98b976c8 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/metal.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/metal.py @@ -76,11 +76,11 @@ def rope_pool_fused( if q_shape != (q_shape[0], num_qo_heads, head_dim): raise ValueError( - "q shape must be [num_tokens, num_qo_heads, head_dim], " f"got {q.shape}" + f"q shape must be [num_tokens, num_qo_heads, head_dim], got {q.shape}" ) if k_shape != (q_shape[0], num_kv_heads, head_dim): raise ValueError( - "k shape must be [num_tokens, num_kv_heads, head_dim], " f"got {k.shape}" + f"k shape must be [num_tokens, num_kv_heads, head_dim], got {k.shape}" ) if v_shape != k_shape: raise ValueError(f"v shape must match k shape, got {v.shape} vs {k.shape}") diff --git a/python/sglang/kernels/aot/python/sgl_kernel/musa.py b/python/sglang/kernels/aot/python/sgl_kernel/musa.py index 49dd825ca..220bb4858 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/musa.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/musa.py @@ -86,9 +86,9 @@ def musa_fused_gemv( out_shape = x.shape[:-1] + ( qweight.shape[0] if not use_swigelu else qweight.shape[0] // 2, ) - assert not ( - use_swigelu and use_rms_norm - ), "gemv only fused one activation (swigelu or rms_norm)!" + assert not (use_swigelu and use_rms_norm), ( + "gemv only fused one activation (swigelu or rms_norm)!" + ) if use_rms_norm: if gamma is None: @@ -113,9 +113,9 @@ def musa_fused_gemv( return output # w4a16 gemv elif qweight_scales is not None: - assert ( - x.dtype == torch.bfloat16 or x.dtype == torch.float16 - ), "W4A16 gemv only support bfloat16 or float16!" + assert x.dtype == torch.bfloat16 or x.dtype == torch.float16, ( + "W4A16 gemv only support bfloat16 or float16!" + ) use_int4_w4a16 = True out_shape = x.shape[:-1] + ( qweight.shape[0] if not use_swigelu else qweight.shape[0] // 2, diff --git a/python/sglang/kernels/aot/python/sgl_kernel/scalar_type.py b/python/sglang/kernels/aot/python/sgl_kernel/scalar_type.py index 571c1dca7..24dbaa93f 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/scalar_type.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/scalar_type.py @@ -70,9 +70,9 @@ class ScalarType: """ def _floating_point_max_int(self) -> int: - assert ( - self.mantissa <= 52 and self.exponent <= 11 - ), f"Cannot represent max/min as a double for type {self.__str__()}" + assert self.mantissa <= 52 and self.exponent <= 11, ( + f"Cannot represent max/min as a double for type {self.__str__()}" + ) max_mantissa = (1 << self.mantissa) - 1 if self.nan_repr == NanRepr.EXTD_RANGE_MAX_MIN: @@ -80,9 +80,9 @@ class ScalarType: max_exponent = (1 << self.exponent) - 2 if self.nan_repr == NanRepr.EXTD_RANGE_MAX_MIN or self.nan_repr == NanRepr.NONE: - assert ( - self.exponent < 11 - ), f"Cannot represent max/min as a double for type {self.__str__()}" + assert self.exponent < 11, ( + f"Cannot represent max/min as a double for type {self.__str__()}" + ) max_exponent = max_exponent + 1 # adjust the exponent to match that of a double @@ -109,25 +109,25 @@ class ScalarType: if self.is_floating_point(): return self._floating_point_max() else: - assert ( - self.size_bits < 64 or self.size_bits == 64 and self.is_signed() - ), "Cannot represent max as an int" + assert self.size_bits < 64 or self.size_bits == 64 and self.is_signed(), ( + "Cannot represent max as an int" + ) return (1 << self.mantissa) - 1 def _raw_min(self) -> Union[int, float]: if self.is_floating_point(): - assert ( - self.is_signed() - ), "We currently assume all floating point types are signed" + assert self.is_signed(), ( + "We currently assume all floating point types are signed" + ) sign_bit_double = 1 << 63 max_raw = self._floating_point_max_int() min_raw = max_raw | sign_bit_double return struct.unpack("!d", struct.pack("!Q", min_raw))[0] else: - assert ( - not self.is_signed() or self.size_bits <= 64 - ), "Cannot represent min as a int64_t" + assert not self.is_signed() or self.size_bits <= 64, ( + "Cannot represent min as a int64_t" + ) if self.is_signed(): return -(1 << (self.size_bits - 1)) diff --git a/python/sglang/kernels/aot/python/sgl_kernel/test_utils.py b/python/sglang/kernels/aot/python/sgl_kernel/test_utils.py index ede113fd0..4d7d1ca25 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/test_utils.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/test_utils.py @@ -94,9 +94,9 @@ def _compute_imbalanced_split( def assert_all_close_or_tiny_diff(a: torch.Tensor, b: torch.Tensor): - assert (a.shape == b.shape) and ( - a.dtype == b.dtype - ), f"{a.shape=} {b.shape=} {a.dtype=} {b.dtype=}" + assert (a.shape == b.shape) and (a.dtype == b.dtype), ( + f"{a.shape=} {b.shape=} {a.dtype=} {b.dtype=}" + ) numel = a.numel() if a.dtype == torch.float8_e4m3fn: diff --git a/python/sglang/kernels/aot/python/sgl_kernel/testing/rotary_embedding.py b/python/sglang/kernels/aot/python/sgl_kernel/testing/rotary_embedding.py index 1a4d90a9d..5a9eb6526 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/testing/rotary_embedding.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/testing/rotary_embedding.py @@ -112,9 +112,9 @@ class RotaryEmbedding(torch.nn.Module): fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """A PyTorch-native implementation of forward().""" - assert ( - fused_set_kv_buffer_arg is None - ), "fused_set_kv_buffer_arg is not supported for native implementation" + assert fused_set_kv_buffer_arg is None, ( + "fused_set_kv_buffer_arg is not supported for native implementation" + ) if offsets is not None: positions = positions + offsets @@ -182,9 +182,9 @@ class SglKernelRotaryEmbedding(RotaryEmbedding): offsets: Optional[torch.Tensor] = None, fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: - assert ( - fused_set_kv_buffer_arg is None - ), "fused_set_kv_buffer_arg is not supported for sgl-kernel implementation" + assert fused_set_kv_buffer_arg is None, ( + "fused_set_kv_buffer_arg is not supported for sgl-kernel implementation" + ) if self.cos_sin_cache.dtype != query.dtype: self.cos_sin_cache = self.cos_sin_cache.to(query.dtype) torch.ops.sgl_kernel.rotary_embedding( diff --git a/python/sglang/kernels/aot/python/sgl_kernel/top_k.py b/python/sglang/kernels/aot/python/sgl_kernel/top_k.py index 4b842499a..c93dcd05c 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/top_k.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/top_k.py @@ -33,9 +33,9 @@ def fast_topk_v2( Returns: The topk indices tensor of shape (B, topk) """ - assert ( - topk == 2048 - ), "fast_topk_v2 is only optimized for deepseek v3.2 model, where topk=2048" + assert topk == 2048, ( + "fast_topk_v2 is only optimized for deepseek v3.2 model, where topk=2048" + ) assert score.dim() == 2 topk_indices = score.new_empty((score.size(0), topk), dtype=torch.int32) torch.ops.sgl_kernel.fast_topk(score, topk_indices, lengths, row_starts) @@ -68,9 +68,9 @@ def fast_topk_transform_fused( Returns: The topk indices tensor of shape (B, topk) """ - assert ( - topk == 2048 - ), "fast_topk_transform_fused is only optimized for deepseek v3.2 model, where topk=2048" + assert topk == 2048, ( + "fast_topk_transform_fused is only optimized for deepseek v3.2 model, where topk=2048" + ) assert score.dim() == 2 src_page_table = page_table_size_1 dst_page_table = score.new_empty((score.shape[0], topk), dtype=torch.int32) @@ -138,9 +138,9 @@ def fast_topk_transform_ragged_fused( Returns: The topk indices tensor of shape (B, topk) """ - assert ( - topk == 2048 - ), "fast_topk_transform_ragged_fused is only optimized for deepseek v3.2 model, where topk=2048" + assert topk == 2048, ( + "fast_topk_transform_ragged_fused is only optimized for deepseek v3.2 model, where topk=2048" + ) assert score.dim() == 2 topk_indices_ragged = score.new_empty((score.shape[0], topk), dtype=torch.int32) torch.ops.sgl_kernel.fast_topk_transform_ragged_fused( diff --git a/python/sglang/kernels/aot/tests/speculative/test_speculative_sampling.py b/python/sglang/kernels/aot/tests/speculative/test_speculative_sampling.py index 5a95f6e15..a828749d8 100644 --- a/python/sglang/kernels/aot/tests/speculative/test_speculative_sampling.py +++ b/python/sglang/kernels/aot/tests/speculative/test_speculative_sampling.py @@ -116,15 +116,15 @@ def test_tree_speculative_sampling_target_only( deterministic=True, ) - assert ( - predicts.tolist() == expected_predicts - ), f"Predicts mismatch for thresholds ({threshold_single}, {threshold_acc})" - assert ( - accept_index.tolist() == expected_accept_index - ), f"Accept index mismatch for thresholds ({threshold_single}, {threshold_acc})" - assert ( - accept_token_num.tolist() == expected_accept_token_num - ), f"Accept token num mismatch for thresholds ({threshold_single}, {threshold_acc})" + assert predicts.tolist() == expected_predicts, ( + f"Predicts mismatch for thresholds ({threshold_single}, {threshold_acc})" + ) + assert accept_index.tolist() == expected_accept_index, ( + f"Accept index mismatch for thresholds ({threshold_single}, {threshold_acc})" + ) + assert accept_token_num.tolist() == expected_accept_token_num, ( + f"Accept token num mismatch for thresholds ({threshold_single}, {threshold_acc})" + ) if __name__ == "__main__": diff --git a/python/sglang/kernels/aot/tests/test_custom_allreduce.py b/python/sglang/kernels/aot/tests/test_custom_allreduce.py index d729392c3..a60b4ae31 100644 --- a/python/sglang/kernels/aot/tests/test_custom_allreduce.py +++ b/python/sglang/kernels/aot/tests/test_custom_allreduce.py @@ -92,9 +92,9 @@ def multi_process_parallel( for i in range(world_size): procs[i].join() - assert ( - procs[i].exitcode == 0 - ), f"Process {i} failed with exit code {procs[i].exitcode}" + assert procs[i].exitcode == 0, ( + f"Process {i} failed with exit code {procs[i].exitcode}" + ) class TestCustomAllReduce(unittest.TestCase): diff --git a/python/sglang/kernels/aot/tests/test_flash_attn_sparse.py b/python/sglang/kernels/aot/tests/test_flash_attn_sparse.py index 1126894a4..0932ca18a 100644 --- a/python/sglang/kernels/aot/tests/test_flash_attn_sparse.py +++ b/python/sglang/kernels/aot/tests/test_flash_attn_sparse.py @@ -251,12 +251,14 @@ def test_sparse_attention( ref_out, ref_lse = ref_attn(q, k, v) - torch.testing.assert_close( - out, ref_out, atol=2e-2, rtol=1e-2 - ), f"{torch.max(torch.abs(out - ref_out))}" - torch.testing.assert_close( - lse, ref_lse, atol=2e-2, rtol=1e-2 - ), f"{torch.max(torch.abs(lse - ref_lse))}" + ( + torch.testing.assert_close(out, ref_out, atol=2e-2, rtol=1e-2), + f"{torch.max(torch.abs(out - ref_out))}", + ) + ( + torch.testing.assert_close(lse, ref_lse, atol=2e-2, rtol=1e-2), + f"{torch.max(torch.abs(lse - ref_lse))}", + ) # sparse attention utils diff --git a/python/sglang/kernels/aot/tests/test_flashmla.py b/python/sglang/kernels/aot/tests/test_flashmla.py index 3afdd7866..f9b7574b6 100644 --- a/python/sglang/kernels/aot/tests/test_flashmla.py +++ b/python/sglang/kernels/aot/tests/test_flashmla.py @@ -198,9 +198,7 @@ def reference_torch_prefill( kvs = torch.index_select( kvs, 0, indices.masked_fill(invalid_indices_mask, 0).flatten() - ).view( - s_q, topk, 576 - ) # [s_q, topk, d_qk] + ).view(s_q, topk, 576) # [s_q, topk, d_qk] attn_score = qs @ kvs.transpose(1, 2) # [s_q, h_q, topk] attn_score.masked_fill_(invalid_indices_mask.unsqueeze(1), float("-inf")) attn_score *= sm_scale * math.log2(math.e) diff --git a/python/sglang/kernels/aot/tests/test_fused_qk_norm_rope.py b/python/sglang/kernels/aot/tests/test_fused_qk_norm_rope.py index 7f1daf5aa..57419b63b 100644 --- a/python/sglang/kernels/aot/tests/test_fused_qk_norm_rope.py +++ b/python/sglang/kernels/aot/tests/test_fused_qk_norm_rope.py @@ -76,9 +76,9 @@ def torch_ref_rms_norm_rope( v_size = num_heads_v * head_dim # Verify dimensions match - assert ( - hidden_size == q_size + k_size + v_size - ), f"Hidden size {hidden_size} doesn't match Q+K+V dimensions {q_size + k_size + v_size}" + assert hidden_size == q_size + k_size + v_size, ( + f"Hidden size {hidden_size} doesn't match Q+K+V dimensions {q_size + k_size + v_size}" + ) # Split the tensor into Q, K, V parts q = qkv[:, :q_size] diff --git a/python/sglang/kernels/aot/tests/test_moe_topk_sigmoid.py b/python/sglang/kernels/aot/tests/test_moe_topk_sigmoid.py index 6b2cba858..46a5cd5dc 100644 --- a/python/sglang/kernels/aot/tests/test_moe_topk_sigmoid.py +++ b/python/sglang/kernels/aot/tests/test_moe_topk_sigmoid.py @@ -44,13 +44,13 @@ def test_topk_sigmoid(num_tokens, num_experts, topk): topk_weights_ref, topk_indices_ref = torch.topk(sigmoid_output, topk, dim=-1) # Verify the top-k weights and indices match the torch native ones - assert torch.allclose( - topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3 - ), f"Weights mismatch: torch={topk_weights_ref} vs SGLang={topk_weights}" + assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), ( + f"Weights mismatch: torch={topk_weights_ref} vs SGLang={topk_weights}" + ) - assert torch.allclose( - topk_indices_ref.int(), topk_indices, atol=0, rtol=0 - ), f"Indices mismatch: torch={topk_indices_ref}, SGLang={topk_indices}" + assert torch.allclose(topk_indices_ref.int(), topk_indices, atol=0, rtol=0), ( + f"Indices mismatch: torch={topk_indices_ref}, SGLang={topk_indices}" + ) @pytest.mark.parametrize( @@ -87,13 +87,13 @@ def test_topk_sigmoid_dtype_regression(num_tokens, num_experts, topk, dtype): gating_output.float(), ) - assert torch.allclose( - topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3 - ), f"Weights mismatch: SGLang old interface={topk_weights_ref} vs SGLang new interface={topk_weights}" + assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), ( + f"Weights mismatch: SGLang old interface={topk_weights_ref} vs SGLang new interface={topk_weights}" + ) - assert torch.allclose( - topk_indices_ref.int(), topk_indices, atol=0, rtol=0 - ), f"Indices mismatch: SGLang old interface={topk_indices_ref}, SGLang new interface={topk_indices}" + assert torch.allclose(topk_indices_ref.int(), topk_indices, atol=0, rtol=0), ( + f"Indices mismatch: SGLang old interface={topk_indices_ref}, SGLang new interface={topk_indices}" + ) @pytest.mark.parametrize( @@ -136,13 +136,13 @@ def test_topk_sigmoid_renormalize(num_tokens, num_experts, topk): ) topk_weights_ref = topk_weights_ref / topk_weights_ref.sum(dim=-1, keepdim=True) - assert torch.allclose( - topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3 - ), f"Weights mismatch: SGLang w/o fused renormalize={topk_weights_ref} vs SGLang w/ fused renormalize={topk_weights}" + assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), ( + f"Weights mismatch: SGLang w/o fused renormalize={topk_weights_ref} vs SGLang w/ fused renormalize={topk_weights}" + ) - assert torch.allclose( - topk_indices_ref.int(), topk_indices, atol=0, rtol=0 - ), f"Indices mismatch: SGLang w/o fused renormalize={topk_indices_ref}, SGLang w/ fused renormalize={topk_indices}" + assert torch.allclose(topk_indices_ref.int(), topk_indices, atol=0, rtol=0), ( + f"Indices mismatch: SGLang w/o fused renormalize={topk_indices_ref}, SGLang w/ fused renormalize={topk_indices}" + ) @pytest.mark.parametrize( @@ -180,13 +180,13 @@ def test_topk_sigmoid_renormalize_correction_bias(num_tokens, num_experts, topk) topk_weights_ref = topk_weights_ref / topk_weights_ref.sum(dim=-1, keepdim=True) # Verify the top-k weights and indices match the torch native ones - assert torch.allclose( - topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3 - ), f"Weights mismatch: torch={topk_weights_ref} vs SGLang={topk_weights}" + assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), ( + f"Weights mismatch: torch={topk_weights_ref} vs SGLang={topk_weights}" + ) - assert torch.allclose( - topk_indices_ref.int(), topk_indices, atol=0, rtol=0 - ), f"Indices mismatch: torch={topk_indices_ref}, SGLang={topk_indices}" + assert torch.allclose(topk_indices_ref.int(), topk_indices, atol=0, rtol=0), ( + f"Indices mismatch: torch={topk_indices_ref}, SGLang={topk_indices}" + ) if __name__ == "__main__": diff --git a/python/sglang/kernels/aot/tests/test_moe_topk_softmax.py b/python/sglang/kernels/aot/tests/test_moe_topk_softmax.py index 77ffe8a46..80c04f983 100644 --- a/python/sglang/kernels/aot/tests/test_moe_topk_softmax.py +++ b/python/sglang/kernels/aot/tests/test_moe_topk_softmax.py @@ -41,13 +41,13 @@ def test_topkfast_softmax(num_tokens, num_experts, topk): topk_weights_ref, topk_indices_ref = torch.topk(softmax_output, topk, dim=-1) # Verify the top-k weights and indices match the torch native ones - assert torch.allclose( - topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3 - ), f"Weights mismatch: torch={topk_indices_ref} vs SGLang={topk_weights}" + assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), ( + f"Weights mismatch: torch={topk_indices_ref} vs SGLang={topk_weights}" + ) - assert compare_topk_values( - gating_output, topk_indices_ref.int(), topk_indices - ), f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}" + assert compare_topk_values(gating_output, topk_indices_ref.int(), topk_indices), ( + f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}" + ) @pytest.mark.parametrize( @@ -79,13 +79,13 @@ def test_topk_softmax(num_tokens, num_experts, topk): topk_weights_ref, topk_indices_ref = torch.topk(softmax_output, topk, dim=-1) # Verify the top-k weights and indices match the torch native ones - assert torch.allclose( - topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3 - ), f"Weights mismatch: torch={topk_indices_ref} vs SGLang={topk_weights}" + assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), ( + f"Weights mismatch: torch={topk_indices_ref} vs SGLang={topk_weights}" + ) - assert compare_topk_values( - gating_output, topk_indices_ref.int(), topk_indices - ), f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}" + assert compare_topk_values(gating_output, topk_indices_ref.int(), topk_indices), ( + f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}" + ) @pytest.mark.parametrize( @@ -122,13 +122,13 @@ def test_topk_softmax_dtype_regression(num_tokens, num_experts, topk, dtype): gating_output.float(), ) - assert torch.allclose( - topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3 - ), f"Weights mismatch: SGLang old interface={topk_indices_ref} vs SGLang new interface={topk_weights}" + assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), ( + f"Weights mismatch: SGLang old interface={topk_indices_ref} vs SGLang new interface={topk_weights}" + ) - assert compare_topk_values( - gating_output, topk_indices_ref.int(), topk_indices - ), f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}" + assert compare_topk_values(gating_output, topk_indices_ref.int(), topk_indices), ( + f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}" + ) @pytest.mark.parametrize( @@ -171,13 +171,13 @@ def test_topk_softmax_renormalize(num_tokens, num_experts, topk): ) topk_weights_ref = topk_weights_ref / topk_weights_ref.sum(dim=-1, keepdim=True) - assert torch.allclose( - topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3 - ), f"Weights mismatch: SGLang w/o fused renormalize={topk_indices_ref} vs SGLang w/ fused renormalize={topk_weights}" + assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), ( + f"Weights mismatch: SGLang w/o fused renormalize={topk_indices_ref} vs SGLang w/ fused renormalize={topk_weights}" + ) - assert compare_topk_values( - gating_output, topk_indices_ref.int(), topk_indices - ), f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}" + assert compare_topk_values(gating_output, topk_indices_ref.int(), topk_indices), ( + f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}" + ) if __name__ == "__main__": diff --git a/python/sglang/kernels/jit/__main__.py b/python/sglang/kernels/jit/__main__.py index b69b5fef7..6507e71d4 100644 --- a/python/sglang/kernels/jit/__main__.py +++ b/python/sglang/kernels/jit/__main__.py @@ -72,9 +72,9 @@ def generate_clangd(): arch = make_jit_cuda_arch(int(major), int(minor)) else: arch = get_jit_cuda_arch() - assert ( - arch.major > 0 - ), "Cannot detect CUDA architecture, please specify --cuda-target explicitly." + assert arch.major > 0, ( + "Cannot detect CUDA architecture, please specify --cuda-target explicitly." + ) compile_flags = [ "-xcuda", diff --git a/python/sglang/kernels/jit/benchmark/marker.py b/python/sglang/kernels/jit/benchmark/marker.py index e34a00e31..71bf3fd9b 100644 --- a/python/sglang/kernels/jit/benchmark/marker.py +++ b/python/sglang/kernels/jit/benchmark/marker.py @@ -253,9 +253,9 @@ class Benchmark(Generic[F]): f"parametrize name {name!r} is not a parameter of " f"{self._fn.__name__}; available: {list(self._fn_params)}" ) - assert ( - name not in self._seen_args - ), f"parametrize name {name!r} is already used" + assert name not in self._seen_args, ( + f"parametrize name {name!r} is already used" + ) self._seen_args.add(name) self._configs.insert(0, (names, vals)) @@ -305,8 +305,7 @@ class Benchmark(Generic[F]): if p.default is inspect.Parameter.empty and p.kind in kinds } - (set(flat_names) | {self._line_arg}) assert not missing, ( - f"parameters not parametrized for {self._fn.__name__}: " - f"{sorted(missing)}" + f"parameters not parametrized for {self._fn.__name__}: {sorted(missing)}" ) results, bandwidths, should_log_bw = self._collect_results() @@ -360,13 +359,13 @@ def parametrize(names: str, vals: List[Any], ci_vals: Optional[List[Any]] = None return [(v,) for v in vs] out: List[Tuple[Any, ...]] = [] for v in vs: - assert isinstance( - v, (tuple, list) - ), f"parametrize: multi-name values must be tuples, got {v!r}" + assert isinstance(v, (tuple, list)), ( + f"parametrize: multi-name values must be tuples, got {v!r}" + ) t = tuple(v) - assert ( - len(t) == arity - ), f"parametrize: each value must have length {arity}, got {t!r}" + assert len(t) == arity, ( + f"parametrize: each value must have length {arity}, got {t!r}" + ) out.append(t) return out diff --git a/python/sglang/kernels/jit/utils/compile/loader.py b/python/sglang/kernels/jit/utils/compile/loader.py index 89fecce45..7b810a445 100644 --- a/python/sglang/kernels/jit/utils/compile/loader.py +++ b/python/sglang/kernels/jit/utils/compile/loader.py @@ -121,7 +121,7 @@ def load_jit( # Also the benign case where a concurrent GC unlinked the leaf # between the lookup and the load. logger.warning( - "Cached JIT module %s failed to load; rebuilding. " "Got error: %s", + "Cached JIT module %s failed to load; rebuilding. Got error: %s", spec.module_name, e, ) diff --git a/python/sglang/kernels/kda_kernels/causal_conv3d_cat_pad_jit.py b/python/sglang/kernels/kda_kernels/causal_conv3d_cat_pad_jit.py index 5d5bebd69..896873a4d 100644 --- a/python/sglang/kernels/kda_kernels/causal_conv3d_cat_pad_jit.py +++ b/python/sglang/kernels/kda_kernels/causal_conv3d_cat_pad_jit.py @@ -25,7 +25,7 @@ def _jit_causal_conv3d_cat_pad_module(dtype: torch.dtype) -> Module: cuda_wrappers=[ ( "causal_conv3d_cat_pad", - "causal_conv3d_cat_pad::" f"CausalConv3dCatPadKernel<{args}>::run", + f"causal_conv3d_cat_pad::CausalConv3dCatPadKernel<{args}>::run", ) ], ) diff --git a/python/sglang/kernels/kda_kernels/qwen3x_nvfp4_gemm_sm120.py b/python/sglang/kernels/kda_kernels/qwen3x_nvfp4_gemm_sm120.py index 607d1e634..45aaa9375 100644 --- a/python/sglang/kernels/kda_kernels/qwen3x_nvfp4_gemm_sm120.py +++ b/python/sglang/kernels/kda_kernels/qwen3x_nvfp4_gemm_sm120.py @@ -353,9 +353,9 @@ class _Qwen3xNvfp4Sm120Kernel: self.occupancy, ) - assert ( - self.epi_stage > 0 - ), "epi_stage <= 0, not enough shared memory. This configuration will be skipped." + assert self.epi_stage > 0, ( + "epi_stage <= 0, not enough shared memory. This configuration will be skipped." + ) ( self.a_smem_layout_staged, diff --git a/python/sglang/kernels/kda_kernels/residual_gate_add_jit.py b/python/sglang/kernels/kda_kernels/residual_gate_add_jit.py index fc17fa6c5..594b0cad8 100644 --- a/python/sglang/kernels/kda_kernels/residual_gate_add_jit.py +++ b/python/sglang/kernels/kda_kernels/residual_gate_add_jit.py @@ -34,11 +34,11 @@ def _jit_residual_gate_add_module(dtype: torch.dtype) -> Module: cuda_wrappers=[ ( "residual_gate_add", - "residual_gate_add::" f"ResidualGateAddKernel<{args}>::run", + f"residual_gate_add::ResidualGateAddKernel<{args}>::run", ), ( "residual_gate_add_transposed", - "residual_gate_add::" f"ResidualGateAddKernel<{args}>::run_transposed", + f"residual_gate_add::ResidualGateAddKernel<{args}>::run_transposed", ), ], ) diff --git a/python/sglang/kernels/ops/activation/activation.py b/python/sglang/kernels/ops/activation/activation.py index bf2f12cc8..f5765181e 100644 --- a/python/sglang/kernels/ops/activation/activation.py +++ b/python/sglang/kernels/ops/activation/activation.py @@ -157,9 +157,9 @@ def run_unary_activation( Unlike :func:`run_activation`, there is no gate/up split — ``input`` and ``out`` share the same shape. """ - assert ( - op_name in SUPPORTED_UNARY_ACTIVATIONS - ), f"Unsupported unary activation: {op_name}" + assert op_name in SUPPORTED_UNARY_ACTIVATIONS, ( + f"Unsupported unary activation: {op_name}" + ) if out is None: out = torch.empty_like(input) _run_unary_activation_inplace(op_name, input, out) diff --git a/python/sglang/kernels/ops/activation/softcap.py b/python/sglang/kernels/ops/activation/softcap.py index 5d39de6f8..8b719c9b9 100644 --- a/python/sglang/kernels/ops/activation/softcap.py +++ b/python/sglang/kernels/ops/activation/softcap.py @@ -101,9 +101,9 @@ def softcap_inplace_logits(full_logits, final_logit_softcapping): row_stride = ncols else: assert full_logits.ndim == 2, "non-contiguous softcap requires 2D tensor" - assert ( - full_logits.stride(1) == 1 - ), "non-contiguous softcap requires contiguous columns" + assert full_logits.stride(1) == 1, ( + "non-contiguous softcap requires contiguous columns" + ) nrows, ncols = full_logits.shape row_stride = full_logits.stride(0) diff --git a/python/sglang/kernels/ops/attention/cutedsl_fp8_paged_mqa_logits.py b/python/sglang/kernels/ops/attention/cutedsl_fp8_paged_mqa_logits.py index b74401168..e0450405e 100644 --- a/python/sglang/kernels/ops/attention/cutedsl_fp8_paged_mqa_logits.py +++ b/python/sglang/kernels/ops/attention/cutedsl_fp8_paged_mqa_logits.py @@ -221,12 +221,12 @@ class FP8MQALogitsKernel: self.block_kv = block_kv self.phys_block_kv = phys_block_kv self.num_blocks_per_mma = block_kv // phys_block_kv - assert ( - block_kv % phys_block_kv == 0 - ), f"block_kv={block_kv} must be divisible by phys_block_kv={phys_block_kv}" - assert ( - self.num_blocks_per_mma <= 4 - ), f"num_blocks_per_mma={self.num_blocks_per_mma} exceeds max 4" + assert block_kv % phys_block_kv == 0, ( + f"block_kv={block_kv} must be divisible by phys_block_kv={phys_block_kv}" + ) + assert self.num_blocks_per_mma <= 4, ( + f"num_blocks_per_mma={self.num_blocks_per_mma} exceeds max 4" + ) self.remove_kv_wait_in_epilogue = remove_kv_wait_in_epilogue self.early_tmem_copy = early_tmem_copy self.smem_subpartition_opt = smem_subpartition_opt diff --git a/python/sglang/kernels/ops/attention/cutedsl_gdn_mtp_ring.py b/python/sglang/kernels/ops/attention/cutedsl_gdn_mtp_ring.py index cf50bb4cb..64c0817b0 100644 --- a/python/sglang/kernels/ops/attention/cutedsl_gdn_mtp_ring.py +++ b/python/sglang/kernels/ops/attention/cutedsl_gdn_mtp_ring.py @@ -3080,9 +3080,9 @@ def gated_delta_rule_mtp_wide_vec( assert K_val == 128 and V_val == 128 assert initial_state_source.dtype == torch.bfloat16 assert tile_v in (32, 64, 128), f"tile_v must be 32/64/128, got {tile_v}" - assert ( - V_val % tile_v == 0 and (tile_v // NUM_GROUPS) % ILP_ROWS == 0 - ), f"tile_v={tile_v} incompatible with 8 groups × ILP=4 layout" + assert V_val % tile_v == 0 and (tile_v // NUM_GROUPS) % ILP_ROWS == 0, ( + f"tile_v={tile_v} incompatible with 8 groups × ILP=4 layout" + ) if cache_ring: assert replayssm_rawv is not None and replayssm_rawk is not None @@ -3194,13 +3194,13 @@ def gated_delta_rule_mtp_wide_vec( ) # Validate recovery_steps for fused recovery+decode mode. - assert ( - 0 <= recovery_steps <= T_val - ), f"recovery_steps must be in [0, T={T_val}], got {recovery_steps}" + assert 0 <= recovery_steps <= T_val, ( + f"recovery_steps must be in [0, T={T_val}], got {recovery_steps}" + ) if recovery_steps > 0: - assert ( - not cache_intermediate_states - ), "recovery_steps > 0 is incompatible with intermediate state caching" + assert not cache_intermediate_states, ( + "recovery_steps > 0 is incompatible with intermediate state caching" + ) assert not disable_state_update, ( "recovery_steps > 0 requires state writeback " "(disable_state_update=False); the boundary writeback at i_t=K-1 " @@ -3220,12 +3220,12 @@ def gated_delta_rule_mtp_wide_vec( # accepted_steps[i] is the per-request phase boundary. per_request_accepted_steps = accepted_steps is not None if per_request_accepted_steps: - assert accepted_steps.shape == ( - B_val, - ), f"accepted_steps must have shape [B={B_val}], got {accepted_steps.shape}" - assert ( - accepted_steps.dtype == torch.int32 - ), f"accepted_steps must be int32, got {accepted_steps.dtype}" + assert accepted_steps.shape == (B_val,), ( + f"accepted_steps must have shape [B={B_val}], got {accepted_steps.shape}" + ) + assert accepted_steps.dtype == torch.int32, ( + f"accepted_steps must be int32, got {accepted_steps.dtype}" + ) assert accepted_steps.device == q.device # FLA-style per-token pool scatter (vLLM API compat). When the public @@ -3236,15 +3236,15 @@ def gated_delta_rule_mtp_wide_vec( # this entry point hit the same fail-fast errors. per_token_pool_scatter = ssm_state_indices is not None if per_token_pool_scatter: - assert ( - intermediate_states_buffer is None - ), "ssm_state_indices and intermediate_states_buffer are mutually exclusive" - assert ( - not disable_state_update - ), "ssm_state_indices requires state writes; disable_state_update must be False" - assert ( - recovery_steps == 0 - ), "ssm_state_indices + recovery_steps>0 not yet supported (MVP exclusion)" + assert intermediate_states_buffer is None, ( + "ssm_state_indices and intermediate_states_buffer are mutually exclusive" + ) + assert not disable_state_update, ( + "ssm_state_indices requires state writes; disable_state_update must be False" + ) + assert recovery_steps == 0, ( + "ssm_state_indices + recovery_steps>0 not yet supported (MVP exclusion)" + ) assert T_val >= 2, ( f"ssm_state_indices requires T >= 2 (got T={T_val}); " f"for T=1 use output_state_indices" @@ -3253,9 +3253,9 @@ def gated_delta_rule_mtp_wide_vec( f"ssm_state_indices must have shape [B={B_val}, T={T_val}], " f"got {tuple(ssm_state_indices.shape)}" ) - assert ( - ssm_state_indices.dtype == torch.int32 - ), f"ssm_state_indices must be int32, got {ssm_state_indices.dtype}" + assert ssm_state_indices.dtype == torch.int32, ( + f"ssm_state_indices must be int32, got {ssm_state_indices.dtype}" + ) assert ssm_state_indices.device == q.device phase_b_unroll = _select_wide_vec_phase_b_unroll( @@ -3517,9 +3517,9 @@ def gated_delta_rule_t1_wide_vec( assert K_val == 128 and V_val == 128 assert initial_state_source.dtype == torch.bfloat16 assert tile_v in (32, 64, 128), f"tile_v must be 32/64/128, got {tile_v}" - assert ( - V_val % tile_v == 0 and (tile_v // NUM_GROUPS) % ILP_ROWS == 0 - ), f"tile_v={tile_v} incompatible with 8 groups × ILP=4 layout" + assert V_val % tile_v == 0 and (tile_v // NUM_GROUPS) % ILP_ROWS == 0, ( + f"tile_v={tile_v} incompatible with 8 groups × ILP=4 layout" + ) if scale is None: scale = 1.0 / math.sqrt(K_val) @@ -3827,9 +3827,9 @@ def gated_delta_rule_mtp( f"intermediate_states_buffer dim 0 ({buffer_size}) must equal " f"batch size B={B}; the buffer is batch-scoped, not pool-scoped" ) - assert ( - cache_steps >= T - ), f"intermediate_states_buffer dim 1 ({cache_steps}) must be >= T={T}" + assert cache_steps >= T, ( + f"intermediate_states_buffer dim 1 ({cache_steps}) must be >= T={T}" + ) assert intermediate_states_buffer.dtype == torch.bfloat16 intermediate_states = intermediate_states_buffer.reshape( B * cache_steps * HV, V, K @@ -3860,28 +3860,28 @@ def gated_delta_rule_mtp( # results/2026-06-03/FLA_SCATTER_MODE_PLAN.md. per_token_pool_scatter = ssm_state_indices is not None if per_token_pool_scatter: - assert ( - intermediate_states_buffer is None - ), "ssm_state_indices and intermediate_states_buffer are mutually exclusive" - assert ( - not disable_state_update - ), "ssm_state_indices requires state writes; disable_state_update must be False" - assert ( - recovery_steps == 0 - ), "ssm_state_indices + recovery_steps>0 not yet supported (MVP exclusion)" - assert ( - T >= 2 - ), f"ssm_state_indices requires T >= 2 (got T={T}); for T=1 use output_state_indices" + assert intermediate_states_buffer is None, ( + "ssm_state_indices and intermediate_states_buffer are mutually exclusive" + ) + assert not disable_state_update, ( + "ssm_state_indices requires state writes; disable_state_update must be False" + ) + assert recovery_steps == 0, ( + "ssm_state_indices + recovery_steps>0 not yet supported (MVP exclusion)" + ) + assert T >= 2, ( + f"ssm_state_indices requires T >= 2 (got T={T}); for T=1 use output_state_indices" + ) assert ssm_state_indices.shape == (B, T), ( f"ssm_state_indices must have shape [B={B}, T={T}], " f"got {tuple(ssm_state_indices.shape)}" ) - assert ( - ssm_state_indices.dtype == torch.int32 - ), f"ssm_state_indices must be int32, got {ssm_state_indices.dtype}" - assert ( - ssm_state_indices.device == q.device - ), f"ssm_state_indices device {ssm_state_indices.device} != q device {q.device}" + assert ssm_state_indices.dtype == torch.int32, ( + f"ssm_state_indices must be int32, got {ssm_state_indices.dtype}" + ) + assert ssm_state_indices.device == q.device, ( + f"ssm_state_indices device {ssm_state_indices.device} != q device {q.device}" + ) # Dispatch to the wide_vec kernel when work_units (B*HV) amortizes its # lower per-CTA parallelism. ``_select_wide_vec_tile_v`` picks tile_v @@ -3960,12 +3960,12 @@ def gated_delta_rule_mtp( # Per-request K opt-in (see gated_delta_rule_mtp_wide_vec for full rationale). per_request_accepted_steps = accepted_steps is not None if per_request_accepted_steps: - assert accepted_steps.shape == ( - B, - ), f"accepted_steps must have shape [B={B}], got {accepted_steps.shape}" - assert ( - accepted_steps.dtype == torch.int32 - ), f"accepted_steps must be int32, got {accepted_steps.dtype}" + assert accepted_steps.shape == (B,), ( + f"accepted_steps must have shape [B={B}], got {accepted_steps.shape}" + ) + assert accepted_steps.dtype == torch.int32, ( + f"accepted_steps must be int32, got {accepted_steps.dtype}" + ) assert accepted_steps.device == q.device # Contiguous pool -> sentinel keys + slot dim marked dynamic (pool-size diff --git a/python/sglang/kernels/ops/attention/cutedsl_kda.py b/python/sglang/kernels/ops/attention/cutedsl_kda.py index 78c7b8495..312769e9d 100644 --- a/python/sglang/kernels/ops/attention/cutedsl_kda.py +++ b/python/sglang/kernels/ops/attention/cutedsl_kda.py @@ -1427,12 +1427,12 @@ def cutedsl_fused_sigmoid_gating_kda_update( N = initial_state_indices.shape[0] assert K == TILE_K, f"Current CuTe DSL KDA kernel requires K={TILE_K}, got {K}" - assert ( - V % TILE_V_SMALL == 0 - ), f"Current CuTe DSL KDA kernel requires V % {TILE_V_SMALL} == 0, got V={V}" - assert ( - V % TILE_V == 0 - ), f"Current CuTe DSL KDA kernel requires V % {TILE_V} == 0, got V={V}" + assert V % TILE_V_SMALL == 0, ( + f"Current CuTe DSL KDA kernel requires V % {TILE_V_SMALL} == 0, got V={V}" + ) + assert V % TILE_V == 0, ( + f"Current CuTe DSL KDA kernel requires V % {TILE_V} == 0, got V={V}" + ) assert (V // TILE_V_SMALL) % NUM_BLOCKS_PER_STATE_SMALL == 0, ( "Small-batch KDA kernel requires num_v_tiles_small divisible by " f"{NUM_BLOCKS_PER_STATE_SMALL}, got V={V}" diff --git a/python/sglang/kernels/ops/attention/decode_attention.py b/python/sglang/kernels/ops/attention/decode_attention.py index e09ce7fef..b535e19e2 100644 --- a/python/sglang/kernels/ops/attention/decode_attention.py +++ b/python/sglang/kernels/ops/attention/decode_attention.py @@ -1483,7 +1483,6 @@ def _lean_attention_decode_kernel( # Use a regular while loop instead of tl.static_range with a dynamic bound to avoid # Triton compiler crashes in the Coalesce pass (max_output_tile_cnt is runtime-computed). while iter < cta_end_tile_gid: - tile_row_idx = iter // tiles_per_khead tile_idx = tile_row_idx * batch_size tile_iter = tile_row_idx * tiles_per_khead diff --git a/python/sglang/kernels/ops/attention/deepseek_v4_rope.py b/python/sglang/kernels/ops/attention/deepseek_v4_rope.py index 1dd647465..fb0ee2c79 100644 --- a/python/sglang/kernels/ops/attention/deepseek_v4_rope.py +++ b/python/sglang/kernels/ops/attention/deepseek_v4_rope.py @@ -354,9 +354,9 @@ def apply_rotary_emb_triton( grid = (batch_size, n_heads if is_3d else 1, num_blocks_dim) if positions is not None: - assert positions.shape == ( - batch_size, - ), f"positions shape {positions.shape} != ({batch_size},)" + assert positions.shape == (batch_size,), ( + f"positions shape {positions.shape} != ({batch_size},)" + ) apply_rotary_emb_triton_kernel[grid]( x, @@ -374,9 +374,9 @@ def apply_rotary_emb_triton( BLOCK_SIZE=BLOCK_SIZE, ) else: - assert ( - freqs_real.shape[0] == batch_size - ), f"freqs_cis batch size {freqs_real.shape[0]} != x batch size {batch_size}" + assert freqs_real.shape[0] == batch_size, ( + f"freqs_cis batch size {freqs_real.shape[0]} != x batch size {batch_size}" + ) apply_rotary_emb_triton_kernel[grid]( x, @@ -621,9 +621,9 @@ def fused_norm_rope_inplace_triton( if weight is not None: assert weight.shape == (head_dim,) if positions is None: - assert ( - freqs_real.shape[0] == M - ), f"freqs_cis row count {freqs_real.shape[0]} != M={M}" + assert freqs_real.shape[0] == M, ( + f"freqs_cis row count {freqs_real.shape[0]} != M={M}" + ) else: assert positions.shape == (M,) and positions.dim() == 1 diff --git a/python/sglang/kernels/ops/attention/dsa/dequant_k_cache.py b/python/sglang/kernels/ops/attention/dsa/dequant_k_cache.py index af01f8a11..55562283f 100644 --- a/python/sglang/kernels/ops/attention/dsa/dequant_k_cache.py +++ b/python/sglang/kernels/ops/attention/dsa/dequant_k_cache.py @@ -181,9 +181,9 @@ def dequantize_k_cache_paged( output: [num_tokens, 1, dim_nope + dim_rope], the de-quantized k-cache """ dim_quant = quant_k_cache.shape[-1] - assert ( - dim_quant == 656 - ), f"dim_quant: {dim_quant} != 656 detected in dequantize_k_cache_paged" + assert dim_quant == 656, ( + f"dim_quant: {dim_quant} != 656 detected in dequantize_k_cache_paged" + ) quant_k_cache = quant_k_cache.view((-1, dim_quant)) # num_tokens can exceed kv_cache_size due to prefix sharing (multiple seqs share same KV slots) diff --git a/python/sglang/kernels/ops/attention/dsa/index_buf_accessor.py b/python/sglang/kernels/ops/attention/dsa/index_buf_accessor.py index 2cd4dd539..6a7165b7d 100644 --- a/python/sglang/kernels/ops/attention/dsa/index_buf_accessor.py +++ b/python/sglang/kernels/ops/attention/dsa/index_buf_accessor.py @@ -308,9 +308,9 @@ def _set_k_and_s_triton( assert scale_dim == 1 if _is_hip: if _use_aiter_preshuffle: - assert ( - page_size % 16 == 0 - ), f"HIP preshuffle requires page_size to be a multiple of 16, got {page_size}" + assert page_size % 16 == 0, ( + f"HIP preshuffle requires page_size to be a multiple of 16, got {page_size}" + ) else: assert page_size == 64 diff --git a/python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py b/python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py index d2dc85132..23d5e3025 100644 --- a/python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py +++ b/python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py @@ -155,9 +155,9 @@ def act_quant( - A tensor of scaling factors with dtype `torch.float32`. """ assert x.is_contiguous(), "Input tensor must be contiguous" - assert ( - x.size(-1) % block_size == 0 - ), f"Last dimension size must be divisible by block_size (block_size={block_size})" + assert x.size(-1) % block_size == 0, ( + f"Last dimension size must be divisible by block_size (block_size={block_size})" + ) N = x.size(-1) if _is_fp8_fnuz: y = torch.empty_like(x, dtype=torch.float8_e4m3fnuz) @@ -272,16 +272,16 @@ def sparse_attention_fwd_kernel_v1( num_stages=2, threads=256, ): - assert dim == tilelang.math.next_power_of_2( - dim - ), f"haven't check padding correctness yet, dim={dim}" - assert tail_dim == tilelang.math.next_power_of_2( - tail_dim - ), f"haven't check padding correctness yet, dim={tail_dim}" + assert dim == tilelang.math.next_power_of_2(dim), ( + f"haven't check padding correctness yet, dim={dim}" + ) + assert tail_dim == tilelang.math.next_power_of_2(tail_dim), ( + f"haven't check padding correctness yet, dim={tail_dim}" + ) assert is_causal == True, "non-casual is not supported" - assert ( - topk % block_I == 0 - ), "otherwise will load some index=0 thus causing wrong kv to be loaded" + assert topk % block_I == 0, ( + "otherwise will load some index=0 thus causing wrong kv to be loaded" + ) if sm_scale is None: sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504 # log2(e) else: @@ -361,7 +361,6 @@ def sparse_attention_fwd_kernel_v1( T.copy(Q[b_i, s_i, H0:H1, D:], Q_tail_shared) for i_i in T.Pipelined(NI, num_stages=num_stages): - for bi_i in T.Parallel(BI): mask[bi_i] = Indices[b_i, s_i, g_i, i_i * BI + bi_i] >= 0 @@ -446,15 +445,15 @@ def sparse_attention_fwd_kernel_v2( sm_scale: Optional[float] = None, block_I: int = 64, ): - assert dim == tilelang.math.next_power_of_2( - dim - ), f"haven't check padding correctness yet, dim={dim}" - assert tail_dim == tilelang.math.next_power_of_2( - tail_dim - ), f"haven't check padding correctness yet, dim={tail_dim}" - assert ( - topk % block_I == 0 - ), "otherwise will load some index=0 thus causing wrong kv to be loaded" + assert dim == tilelang.math.next_power_of_2(dim), ( + f"haven't check padding correctness yet, dim={dim}" + ) + assert tail_dim == tilelang.math.next_power_of_2(tail_dim), ( + f"haven't check padding correctness yet, dim={tail_dim}" + ) + assert topk % block_I == 0, ( + "otherwise will load some index=0 thus causing wrong kv to be loaded" + ) if sm_scale is None: sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504 # log2(e) else: @@ -1078,9 +1077,9 @@ def sparse_mla_fwd_decode_partial_fp8( threads=256, ): assert d_v == 512, f"only support d_v=512" - assert ( - topk % block_I == 0 - ), "otherwise will load some index=0 thus causing wrong kv to be loaded" + assert topk % block_I == 0, ( + "otherwise will load some index=0 thus causing wrong kv to be loaded" + ) # Softmax scores are in [0, 1]. We scale by fp8_max_val before FP8 cast # to better utilize FP8 dynamic range, then apply the inverse scale after GEMM. @@ -1104,9 +1103,9 @@ def sparse_mla_fwd_decode_partial_fp8( h_per_block = 16 # Match bf16 partial behavior: keep fixed 16-head tiles and use # sliced T.copy on H0:H1 for tail handling. - assert ( - num_heads <= h_per_block or num_heads % h_per_block == 0 - ), "num_heads must be <=16 or divisible by 16" + assert num_heads <= h_per_block or num_heads % h_per_block == 0, ( + "num_heads must be <=16 or divisible by 16" + ) head_blocks_per_seq = (num_heads + h_per_block - 1) // h_per_block batch = 1 @@ -1594,9 +1593,7 @@ def dpsk_v4_fp8_partial_kernel( sm_scale = sm_scale * log2e assert dim == 448 and tail_dim == 64 assert topk_1 % block_I == 0 - assert ( - topk_1 // block_I - ) % inner_iter_1 == 0, ( + assert (topk_1 // block_I) % inner_iter_1 == 0, ( f"NI_1={topk_1 // block_I} must be divisible by inner_iter_1={inner_iter_1}" ) assert block_size_kv_1 > 0 and (block_size_kv_1 & (block_size_kv_1 - 1)) == 0 @@ -1605,9 +1602,7 @@ def dpsk_v4_fp8_partial_kernel( if is_dual: assert inner_iter_2 > 0, "dual-cache call requires inner_iter_2 > 0" assert topk_2 % block_I == 0 - assert ( - topk_2 // block_I - ) % inner_iter_2 == 0, ( + assert (topk_2 // block_I) % inner_iter_2 == 0, ( f"NI_2={topk_2 // block_I} must be divisible by inner_iter_2={inner_iter_2}" ) assert block_size_kv_2 > 0 and (block_size_kv_2 & (block_size_kv_2 - 1)) == 0 @@ -2256,12 +2251,8 @@ def dpsk_v4_combine_kernel( @T.prim_func def main( - Partial_O: T.Tensor( - [batch, seq_len, n_groups, num_heads, DT], BF16 - ), # type: ignore - Partial_LSE: T.Tensor( - [batch, seq_len, n_groups, num_heads], accum_dtype - ), # type: ignore + Partial_O: T.Tensor([batch, seq_len, n_groups, num_heads, DT], BF16), # type: ignore + Partial_LSE: T.Tensor([batch, seq_len, n_groups, num_heads], accum_dtype), # type: ignore Topk_length_1: T.Tensor([batch], INT32), # type: ignore Topk_length_2: T.Tensor([batch], INT32), # type: ignore Attn_sink: T.Tensor([num_heads], FP32), # type: ignore @@ -2369,12 +2360,8 @@ def dpsk_v4_combine_kernel( @T.prim_func def main( - Partial_O: T.Tensor( - [batch, seq_len, n_groups, num_heads, DT], BF16 - ), # type: ignore - Partial_LSE: T.Tensor( - [batch, seq_len, n_groups, num_heads], accum_dtype - ), # type: ignore + Partial_O: T.Tensor([batch, seq_len, n_groups, num_heads, DT], BF16), # type: ignore + Partial_LSE: T.Tensor([batch, seq_len, n_groups, num_heads], accum_dtype), # type: ignore Attn_sink: T.Tensor([num_heads], FP32), # type: ignore Output: T.Tensor([batch, seq_len, num_heads, DT], BF16), # type: ignore LSE: T.Tensor([batch, seq_len, num_heads], accum_dtype), # type: ignore diff --git a/python/sglang/kernels/ops/attention/dsa/triton_kernel.py b/python/sglang/kernels/ops/attention/dsa/triton_kernel.py index 0d2969804..4ce7a4c1e 100644 --- a/python/sglang/kernels/ops/attention/dsa/triton_kernel.py +++ b/python/sglang/kernels/ops/attention/dsa/triton_kernel.py @@ -99,9 +99,9 @@ def act_quant( - A tensor of scaling factors with dtype `torch.float32`. """ assert x.is_contiguous(), "Input tensor must be contiguous" - assert ( - x.size(-1) % block_size == 0 - ), f"Last dimension size must be divisible by block_size (block_size={block_size})" + assert x.size(-1) % block_size == 0, ( + f"Last dimension size must be divisible by block_size (block_size={block_size})" + ) # Flatten all dims except last N = x.size(-1) diff --git a/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla.py b/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla.py index 2c6ea0c9e..0a3c06310 100644 --- a/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla.py +++ b/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla.py @@ -69,9 +69,7 @@ def _sparse_mla_fwd_kernel( ) # [H, D_V] q_tail = tl.load( q_rope_ptr + s_i * H * D_TAIL + h[:, None] * D_TAIL + dt[None, :] - ).to( - q_nope_ptr.dtype.element_ty - ) # [H, D_TAIL] + ).to(q_nope_ptr.dtype.element_ty) # [H, D_TAIL] m_i = tl.full([H], -float("inf"), tl.float32) l_i = tl.zeros([H], tl.float32) @@ -89,9 +87,7 @@ def _sparse_mla_fwd_kernel( ) # [BLOCK_N, D_V] -- reused as V kv_tail = tl.load( kbase + (D_V + dt)[None, :], mask=valid[:, None], other=0.0 - ).to( - q_nope_ptr.dtype.element_ty - ) # [BLOCK_N, D_TAIL] + ).to(q_nope_ptr.dtype.element_ty) # [BLOCK_N, D_TAIL] qk = tl.dot(q_main, tl.trans(kv_main)).to(tl.float32) qk += tl.dot(q_tail, tl.trans(kv_tail)).to(tl.float32) diff --git a/python/sglang/kernels/ops/attention/dsv4/index_buf_accessor.py b/python/sglang/kernels/ops/attention/dsv4/index_buf_accessor.py index 8536e18a9..50738d8f8 100644 --- a/python/sglang/kernels/ops/attention/dsv4/index_buf_accessor.py +++ b/python/sglang/kernels/ops/attention/dsv4/index_buf_accessor.py @@ -209,7 +209,9 @@ def _set_k_and_s_torch( == num_tokens_to_write_nope == num_tokens_to_write_rope == num_tokens_to_write_scale - ), f"{num_tokens_to_write=} {num_tokens_to_write_nope=} {num_tokens_to_write_rope=} {num_tokens_to_write_scale=}" + ), ( + f"{num_tokens_to_write=} {num_tokens_to_write_nope=} {num_tokens_to_write_rope=} {num_tokens_to_write_scale=}" + ) assert buf.dtype == torch.uint8 assert loc.dtype in [ diff --git a/python/sglang/kernels/ops/attention/dsv4/metadata_kernel.py b/python/sglang/kernels/ops/attention/dsv4/metadata_kernel.py index efe579f1a..d477a0238 100644 --- a/python/sglang/kernels/ops/attention/dsv4/metadata_kernel.py +++ b/python/sglang/kernels/ops/attention/dsv4/metadata_kernel.py @@ -110,9 +110,9 @@ def _init_compressed_attn_metadata_triton( # no cache-write locations. Keep the write buffers unpadded and mask those # rows in the kernel. num_write_tokens = raw_out_loc.shape[0] - assert ( - num_write_tokens <= bs - ), f"raw_out_loc has {num_write_tokens} rows, expected at most {bs} metadata rows" + assert num_write_tokens <= bs, ( + f"raw_out_loc has {num_write_tokens} rows, expected at most {bs} metadata rows" + ) device = seq_lens.device c4_out_loc = torch.empty(num_write_tokens, dtype=torch.int64, device=device) @@ -126,12 +126,12 @@ def _init_compressed_attn_metadata_triton( c128_seq_lens_clamp1 = torch.empty(bs, dtype=torch.int32, device=device) if compute_page_indices: - assert ( - page_table is not None - ), "page_table required when compute_page_indices=True" - assert ( - page_size >= 128 and page_size % 128 == 0 - ), "page_size must be a multiple of 128 when compute_page_indices=True" + assert page_table is not None, ( + "page_table required when compute_page_indices=True" + ) + assert page_size >= 128 and page_size % 128 == 0, ( + "page_size must be a multiple of 128 when compute_page_indices=True" + ) max_pages = page_table.shape[1] c128_page_size = page_size // 128 c128_cur_max_seq_len = c128_page_size * max_pages diff --git a/python/sglang/kernels/ops/attention/fa4_sm120/flash_fwd.py b/python/sglang/kernels/ops/attention/fa4_sm120/flash_fwd.py index 4188b95c4..5245639f5 100644 --- a/python/sglang/kernels/ops/attention/fa4_sm120/flash_fwd.py +++ b/python/sglang/kernels/ops/attention/fa4_sm120/flash_fwd.py @@ -1090,9 +1090,9 @@ class FlashAttentionForwardSm120(FlashAttentionForwardBase): ): assert blocksparse_tensors is None, "Block sparsity is not supported on SM120" assert (mBias is not None) == self.has_bias - assert ( - mPageTable is None or self.paged_kv - ), "SM120 paged KV requires the dedicated DMA-warp specialization" + assert mPageTable is None or self.paged_kv, ( + "SM120 paged KV requires the dedicated DMA-warp specialization" + ) self._check_type( *( t.element_type if t is not None else None @@ -1251,7 +1251,9 @@ class FlashAttentionForwardSm120(FlashAttentionForwardBase): TileScheduler = ( Sm120UniformBatchScheduler if is_varlen and self.direct_uniform_batch - else SingleTileVarlenScheduler if is_varlen else SingleTileScheduler + else SingleTileVarlenScheduler + if is_varlen + else SingleTileScheduler ) tile_sched_args = TileSchedulerArguments( num_block=cute.ceil_div(cute.size(mQ.shape[0]), self.tile_m), diff --git a/python/sglang/kernels/ops/attention/fa4_sm120/scheduler.py b/python/sglang/kernels/ops/attention/fa4_sm120/scheduler.py index 9cc7cb365..5fe3df516 100644 --- a/python/sglang/kernels/ops/attention/fa4_sm120/scheduler.py +++ b/python/sglang/kernels/ops/attention/fa4_sm120/scheduler.py @@ -75,9 +75,9 @@ class Sm120UniformBatchScheduler: loc=None, ip=None, ) -> Params: - assert ( - scheduling_mode == SchedulingMode.STATIC - ), f"SM120 uniform-batch scheduler only supports STATIC, got {scheduling_mode!r}" + assert scheduling_mode == SchedulingMode.STATIC, ( + f"SM120 uniform-batch scheduler only supports STATIC, got {scheduling_mode!r}" + ) return Sm120UniformBatchScheduler.Params.create(args, loc=loc, ip=ip) @staticmethod diff --git a/python/sglang/kernels/ops/attention/fla/chunk.py b/python/sglang/kernels/ops/attention/fla/chunk.py index 8acb9d14e..80cfb4e1f 100644 --- a/python/sglang/kernels/ops/attention/fla/chunk.py +++ b/python/sglang/kernels/ops/attention/fla/chunk.py @@ -85,7 +85,6 @@ def chunk_gated_delta_rule_fwd( class ChunkGatedDeltaRuleFunction(torch.autograd.Function): - @staticmethod @input_guard @autocast_custom_fwd @@ -207,12 +206,12 @@ def chunk_gated_delta_rule( ) """ assert q.dtype == k.dtype == v.dtype - assert ( - q.dtype != torch.float32 - ), "ChunkGatedDeltaRuleFunction does not support float32. Please use bfloat16." - assert ( - len(beta.shape) == 3 - ), "beta must be of shape [B, T, H] if head_first=False, or [B, H, T] otherwise." + assert q.dtype != torch.float32, ( + "ChunkGatedDeltaRuleFunction does not support float32. Please use bfloat16." + ) + assert len(beta.shape) == 3, ( + "beta must be of shape [B, T, H] if head_first=False, or [B, H, T] otherwise." + ) if head_first: raise DeprecationWarning( diff --git a/python/sglang/kernels/ops/attention/fla/chunk_delta_h.py b/python/sglang/kernels/ops/attention/fla/chunk_delta_h.py index 2fe5e623d..d80aa25d1 100644 --- a/python/sglang/kernels/ops/attention/fla/chunk_delta_h.py +++ b/python/sglang/kernels/ops/attention/fla/chunk_delta_h.py @@ -82,9 +82,10 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( i_v, i_nh = tl.program_id(0), tl.program_id(1) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( - cu_seqlens + i_n + 1 - ).to(tl.int32) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) T = eos - bos NT = tl.cdiv(T, BT) boh = tl.load(chunk_offsets + i_n).to(tl.int32) @@ -326,9 +327,9 @@ def chunk_gated_delta_rule_fwd_h( chunk_indices: Optional[torch.LongTensor] = None, use_exp2: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - assert not ( - use_exp2 and g is not None - ), "use_exp2 covers only the per-channel gk path; scalar g stays natural-exp" + assert not (use_exp2 and g is not None), ( + "use_exp2 covers only the per-channel gk path; scalar g stays natural-exp" + ) B, T, Hg, K, V = *k.shape, u.shape[-1] H = u.shape[-2] BT = CHUNK_SIZE diff --git a/python/sglang/kernels/ops/attention/fla/chunk_fwd.py b/python/sglang/kernels/ops/attention/fla/chunk_fwd.py index 828ddc4c9..4e7b06810 100644 --- a/python/sglang/kernels/ops/attention/fla/chunk_fwd.py +++ b/python/sglang/kernels/ops/attention/fla/chunk_fwd.py @@ -71,12 +71,14 @@ def chunk_gated_delta_rule_fwd_kkt_solve_kernel( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( - chunk_indices + i_t * 2 + 1 - ).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( - cu_seqlens + i_n + 1 - ).to(tl.int32) + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) T = eos - bos else: bos, eos = i_b * T, i_b * T + T diff --git a/python/sglang/kernels/ops/attention/fla/chunk_intra.py b/python/sglang/kernels/ops/attention/fla/chunk_intra.py index 454eba022..908384d9d 100644 --- a/python/sglang/kernels/ops/attention/fla/chunk_intra.py +++ b/python/sglang/kernels/ops/attention/fla/chunk_intra.py @@ -88,12 +88,14 @@ def chunk_kda_fwd_kernel_inter_solve_fused( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( - chunk_indices + i_t * 2 + 1 - ).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( - cu_seqlens + i_n + 1 - ).to(tl.int32) + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -819,12 +821,14 @@ def chunk_kda_fwd_kernel_intra_sub_chunk( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( - chunk_indices + i_t * 2 + 1 - ).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( - cu_seqlens + i_n + 1 - ).to(tl.int32) + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) T = eos - bos else: bos, eos = i_b * T, i_b * T + T diff --git a/python/sglang/kernels/ops/attention/fla/chunk_intra_token_parallel.py b/python/sglang/kernels/ops/attention/fla/chunk_intra_token_parallel.py index 7481b9923..feea3c53f 100644 --- a/python/sglang/kernels/ops/attention/fla/chunk_intra_token_parallel.py +++ b/python/sglang/kernels/ops/attention/fla/chunk_intra_token_parallel.py @@ -62,9 +62,10 @@ def chunk_kda_fwd_kernel_intra_token_parallel( left = mid + 1 i_n = left - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( - cu_seqlens + i_n + 1 - ).to(tl.int32) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) T = eos - bos i_t = i_tg - bos else: diff --git a/python/sglang/kernels/ops/attention/fla/chunk_o.py b/python/sglang/kernels/ops/attention/fla/chunk_o.py index c2c04312a..f644f15fc 100644 --- a/python/sglang/kernels/ops/attention/fla/chunk_o.py +++ b/python/sglang/kernels/ops/attention/fla/chunk_o.py @@ -53,12 +53,14 @@ def chunk_fwd_kernel_o( if IS_VARLEN: i_tg = i_t - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( - chunk_indices + i_t * 2 + 1 - ).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( - cu_seqlens + i_n + 1 - ).to(tl.int32) + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) T = eos - bos NT = tl.cdiv(T, BT) else: diff --git a/python/sglang/kernels/ops/attention/fla/cumsum.py b/python/sglang/kernels/ops/attention/fla/cumsum.py index 182211e55..b31331c72 100644 --- a/python/sglang/kernels/ops/attention/fla/cumsum.py +++ b/python/sglang/kernels/ops/attention/fla/cumsum.py @@ -37,12 +37,14 @@ def chunk_local_cumsum_scalar_kernel( i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( - chunk_indices + i_t * 2 + 1 - ).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( - cu_seqlens + i_n + 1 - ).to(tl.int32) + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -97,12 +99,14 @@ def chunk_local_cumsum_vector_kernel( i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( - chunk_indices + i_t * 2 + 1 - ).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( - cu_seqlens + i_n + 1 - ).to(tl.int32) + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -169,9 +173,9 @@ def chunk_local_cumsum_scalar( B, H, T = g.shape else: B, T, H = g.shape - assert chunk_size == 2 ** ( - chunk_size.bit_length() - 1 - ), "chunk_size must be a power of 2" + assert chunk_size == 2 ** (chunk_size.bit_length() - 1), ( + "chunk_size must be a power of 2" + ) BT = chunk_size if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, BT) @@ -216,9 +220,9 @@ def chunk_local_cumsum_vector( if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, BT) NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - assert chunk_size == 2 ** ( - chunk_size.bit_length() - 1 - ), "chunk_size must be a power of 2" + assert chunk_size == 2 ** (chunk_size.bit_length() - 1), ( + "chunk_size must be a power of 2" + ) g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) @@ -260,9 +264,9 @@ def chunk_local_cumsum( **kwargs, ) -> torch.Tensor: if cu_seqlens is not None: - assert ( - g.shape[0] == 1 - ), "Only batch size 1 is supported when cu_seqlens are provided" + assert g.shape[0] == 1, ( + "Only batch size 1 is supported when cu_seqlens are provided" + ) if len(g.shape) == 3: return chunk_local_cumsum_scalar( g=g, diff --git a/python/sglang/kernels/ops/attention/fla/fused_norm_gate.py b/python/sglang/kernels/ops/attention/fla/fused_norm_gate.py index 6514db459..cdc3df1cc 100644 --- a/python/sglang/kernels/ops/attention/fla/fused_norm_gate.py +++ b/python/sglang/kernels/ops/attention/fla/fused_norm_gate.py @@ -390,9 +390,9 @@ class FusedRMSNormGated(nn.Module): residual_in_fp32: bool = False, ) -> torch.Tensor: if _use_cpu: - assert ( - self.activation == "silu" - ), "CPU rmsnorm_gated currently only supports activation silu" + assert self.activation == "silu", ( + "CPU rmsnorm_gated currently only supports activation silu" + ) return torch.ops.sgl_kernel.fused_rmsnorm_gated_cpu( x, self.weight, g, self.eps ) diff --git a/python/sglang/kernels/ops/attention/fla/fused_recurrent.py b/python/sglang/kernels/ops/attention/fla/fused_recurrent.py index 457535586..1a59ad064 100644 --- a/python/sglang/kernels/ops/attention/fla/fused_recurrent.py +++ b/python/sglang/kernels/ops/attention/fla/fused_recurrent.py @@ -43,9 +43,10 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( i_n, i_hv = i_nh // HV, i_nh % HV i_h = i_hv // (HV // H) if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load( - cu_seqlens + i_n + 1 - ).to(tl.int64) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), + ) all = T T = eos - bos else: @@ -708,7 +709,6 @@ def fused_recurrent_kda_packed_decode( class FusedRecurrentFunction(torch.autograd.Function): - @staticmethod @input_guard def forward( @@ -907,9 +907,10 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel( i_n, i_hv = i_nh // HV, i_nh % HV i_h = i_hv // (HV // H) if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load( - cu_seqlens + i_n + 1 - ).to(tl.int64) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), + ) all = T T = eos - bos else: @@ -1144,7 +1145,6 @@ def fused_recurrent_gated_delta_rule_update_fwd( class FusedRecurrentUpdateFunction(torch.autograd.Function): - @staticmethod @input_guard def forward( diff --git a/python/sglang/kernels/ops/attention/fla/gdn_replayssm_spec_decode.py b/python/sglang/kernels/ops/attention/fla/gdn_replayssm_spec_decode.py index 0b4ef1a69..aa1e53122 100644 --- a/python/sglang/kernels/ops/attention/fla/gdn_replayssm_spec_decode.py +++ b/python/sglang/kernels/ops/attention/fla/gdn_replayssm_spec_decode.py @@ -678,9 +678,9 @@ def _launch_gdn_spec( num_slots, HV, V, K = checkpoint_state.shape H = k.shape[1] B = query_start_loc.shape[0] - 1 - assert ( - max_cache_len & (max_cache_len - 1) == 0 - ), "circular cache requires power-of-two max_cache_len" + assert max_cache_len & (max_cache_len - 1) == 0, ( + "circular cache requires power-of-two max_cache_len" + ) assert d_cache.shape[2] == max_cache_len BK = triton.next_power_of_2(K) diff --git a/python/sglang/kernels/ops/attention/fla/kda.py b/python/sglang/kernels/ops/attention/fla/kda.py index ad9720c53..605583a31 100644 --- a/python/sglang/kernels/ops/attention/fla/kda.py +++ b/python/sglang/kernels/ops/attention/fla/kda.py @@ -1046,18 +1046,18 @@ def kda_gate_chunk_cumsum( Cumulative-summed gated tensor of shape [B, T, H, K]. """ if cu_seqlens is not None: - assert ( - g.shape[0] == 1 - ), "Only batch size 1 is supported when cu_seqlens are provided" + assert g.shape[0] == 1, ( + "Only batch size 1 is supported when cu_seqlens are provided" + ) assert len(g.shape) == 4 B, T, H, S = g.shape BT = chunk_size if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, BT) NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - assert chunk_size == 2 ** ( - chunk_size.bit_length() - 1 - ), "chunk_size must be a power of 2" + assert chunk_size == 2 ** (chunk_size.bit_length() - 1), ( + "chunk_size must be a power of 2" + ) g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) diff --git a/python/sglang/kernels/ops/attention/fla/l2norm.py b/python/sglang/kernels/ops/attention/fla/l2norm.py index 0527324a8..a55e322c3 100644 --- a/python/sglang/kernels/ops/attention/fla/l2norm.py +++ b/python/sglang/kernels/ops/attention/fla/l2norm.py @@ -120,7 +120,6 @@ def l2norm_fwd( class L2NormFunction(torch.autograd.Function): - @staticmethod @input_guard def forward(ctx, x, eps=1e-6, output_dtype=None): @@ -137,7 +136,6 @@ l2_norm = l2norm class L2Norm(nn.Module): - def __init__(self, eps: float = 1e-6, output_dtype: Optional[torch.dtype] = None): super().__init__() self.eps = eps diff --git a/python/sglang/kernels/ops/attention/fla/layernorm_gated.py b/python/sglang/kernels/ops/attention/fla/layernorm_gated.py index a599e761f..05631f006 100644 --- a/python/sglang/kernels/ops/attention/fla/layernorm_gated.py +++ b/python/sglang/kernels/ops/attention/fla/layernorm_gated.py @@ -345,7 +345,6 @@ def rms_norm_gated( class LayerNormFn(torch.autograd.Function): - @staticmethod def forward( ctx, @@ -389,7 +388,6 @@ def layernorm_fn( class LayerNorm(torch.nn.Module): - def __init__( self, hidden_size, @@ -431,7 +429,6 @@ class LayerNorm(torch.nn.Module): class RMSNorm(torch.nn.Module): - def __init__( self, hidden_size, @@ -465,7 +462,9 @@ class RMSNorm(torch.nn.Module): self.norm_before_gate and self.group_size is None and self.activation == "swish" - ), "CPU rmsnorm_gated currently only supports norm before gate without group size or activation other than swish" + ), ( + "CPU rmsnorm_gated currently only supports norm before gate without group size or activation other than swish" + ) return torch.ops.sgl_kernel.fused_rmsnorm_gated_cpu( x, self.weight, z, self.eps ) diff --git a/python/sglang/kernels/ops/attention/fla/utils.py b/python/sglang/kernels/ops/attention/fla/utils.py index 4154a3c52..7c0b89006 100644 --- a/python/sglang/kernels/ops/attention/fla/utils.py +++ b/python/sglang/kernels/ops/attention/fla/utils.py @@ -326,9 +326,9 @@ if torch_release >= (2, 4): return device_torch_lib.device(index) else: - assert ( - device == "cuda" - ), "Only cuda device is supported for PyTorch version < 2.4.0." + assert device == "cuda", ( + "Only cuda device is supported for PyTorch version < 2.4.0." + ) autocast_custom_fwd = device_torch_lib.amp.custom_fwd autocast_custom_bwd = device_torch_lib.amp.custom_bwd diff --git a/python/sglang/kernels/ops/attention/fla/wy_fast.py b/python/sglang/kernels/ops/attention/fla/wy_fast.py index 9956b6a92..762e09030 100644 --- a/python/sglang/kernels/ops/attention/fla/wy_fast.py +++ b/python/sglang/kernels/ops/attention/fla/wy_fast.py @@ -43,12 +43,14 @@ def recompute_w_u_fwd_kernel( i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( - chunk_indices + i_t * 2 + 1 - ).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( - cu_seqlens + i_n + 1 - ).to(tl.int32) + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) T = eos - bos else: bos, eos = i_b * T, i_b * T + T diff --git a/python/sglang/kernels/ops/attention/flash_mla_sm120.py b/python/sglang/kernels/ops/attention/flash_mla_sm120.py index bf8dc234d..906d4aa22 100644 --- a/python/sglang/kernels/ops/attention/flash_mla_sm120.py +++ b/python/sglang/kernels/ops/attention/flash_mla_sm120.py @@ -77,9 +77,7 @@ def _gather_and_dequant(k_cache, indices, page_size): raw_pages = k_cache.as_strided( (num_pages, page_bytes), (page_bytes, 1), - ).view( - torch.uint8 - ) # (num_pages, page_bytes) uint8 + ).view(torch.uint8) # (num_pages, page_bytes) uint8 # Note: float8_e4m3fn and uint8 are both 1 byte, view is safe # Compute byte offsets within each page diff --git a/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/chunk_fwd.py b/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/chunk_fwd.py index de399d33e..480b3b621 100644 --- a/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/chunk_fwd.py +++ b/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/chunk_fwd.py @@ -922,8 +922,7 @@ def chunk_kda_fwd( if needs_eqlen_pad: if B != 1 and T % BT != 0: raise NotImplementedError( - f"eqlen with B>1 and T % {BT} != 0 not supported " - f"(got B={B}, T={T})." + f"eqlen with B>1 and T % {BT} != 0 not supported (got B={B}, T={T})." ) T_padded = ((T + CPB_BT - 1) // CPB_BT) * CPB_BT # Pre-allocated padded scratch buffers (per (B,T_padded,H,K,dtype) cache diff --git a/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/fuse_k4_only_persistent.py b/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/fuse_k4_only_persistent.py index 289049e7e..526cd6170 100644 --- a/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/fuse_k4_only_persistent.py +++ b/python/sglang/kernels/ops/attention/linear/kda_nvidia_prefill/fuse_k4_only_persistent.py @@ -407,7 +407,6 @@ def k4_persistent_kernel( # ==== WG1: State readout + decay ==== if warpgroup_idx == STATE_WG: - cId_128 = cute.make_identity_tensor((M6, N6)) tCtState_mn = transform_partitioned_tensor_layout(tCtState) @@ -698,7 +697,6 @@ def k4_persistent_kernel( # ==== TMA warp (warp 2) ==== elif warp_idx == TMA_WARP: - cta_layout = cute.make_layout(1) scheduler = GDNTileScheduler.create( @@ -995,7 +993,6 @@ def k4_persistent_kernel( # ==== WG2: W/NV/O readout (SS-mode, no Phase 2) ==== elif warpgroup_idx == READOUT_WG: - tCtW_mn = transform_partitioned_tensor_layout(tCtW) tCtNV_mn = transform_partitioned_tensor_layout(tCtNV) tCtO_mn = transform_partitioned_tensor_layout(tCtO) diff --git a/python/sglang/kernels/ops/attention/linear/kda_ptx_prefill/__init__.py b/python/sglang/kernels/ops/attention/linear/kda_ptx_prefill/__init__.py index 1a24cc0a7..bf87c27d6 100644 --- a/python/sglang/kernels/ops/attention/linear/kda_ptx_prefill/__init__.py +++ b/python/sglang/kernels/ops/attention/linear/kda_ptx_prefill/__init__.py @@ -107,9 +107,9 @@ def chunk_kda_fwd( Returns the fla-shaped 12-tuple: (o [B,T,H,128] bf16, final_state [N,H,128,128] fp32 or None, then Nones, ..., h, initial_state). """ - assert ( - chunk_size == CHUNK - ), f"kda_prefill supports chunk_size={CHUNK} only, got {chunk_size}" + assert chunk_size == CHUNK, ( + f"kda_prefill supports chunk_size={CHUNK} only, got {chunk_size}" + ) if cp_context is not None or disable_recompute: raise NotImplementedError( "kda_prefill is the inference forward path: cp_context, " @@ -124,9 +124,9 @@ def chunk_kda_fwd( if state_v_first and initial_state is not None: # [V,K]-layout state: pure transpose (K==V==128), exact, ~us/call initial_state = initial_state.transpose(-1, -2).contiguous() - assert ( - q.dim() == 4 and q.shape[-1] == K and v.shape[-1] == K - ), f"expected [B,T,H,{K}] q/k/v, got q={tuple(q.shape)} v={tuple(v.shape)}" + assert q.dim() == 4 and q.shape[-1] == K and v.shape[-1] == K, ( + f"expected [B,T,H,{K}] q/k/v, got q={tuple(q.shape)} v={tuple(v.shape)}" + ) B, T, H, _ = q.shape cu_cpu = None @@ -145,9 +145,9 @@ def chunk_kda_fwd( betaf = beta.reshape(Tt, H).contiguous() if use_gate_in_kernel: - assert ( - A_log is not None and dt_bias is not None - ), "use_gate_in_kernel=True requires A_log and dt_bias" + assert A_log is not None and dt_bias is not None, ( + "use_gate_in_kernel=True requires A_log and dt_bias" + ) assert g.dtype == torch.bfloat16, f"raw gate input must be bf16, got {g.dtype}" gf = g.reshape(Tt, H, K).contiguous() sg = lower_bound is not None # fla: lb presence selects safe-gate diff --git a/python/sglang/kernels/ops/attention/linear/lightning_attn.py b/python/sglang/kernels/ops/attention/linear/lightning_attn.py index d415eefc9..1a968dd88 100644 --- a/python/sglang/kernels/ops/attention/linear/lightning_attn.py +++ b/python/sglang/kernels/ops/attention/linear/lightning_attn.py @@ -392,7 +392,6 @@ def _fwd_none_diag_kernel( class _attention(torch.autograd.Function): - @staticmethod def forward(ctx, q, k, v, s, kv_history): # Forward pass of the lightning attention algorithm diff --git a/python/sglang/kernels/ops/attention/linear/seg_la.py b/python/sglang/kernels/ops/attention/linear/seg_la.py index abda8b529..9f5f6a44c 100644 --- a/python/sglang/kernels/ops/attention/linear/seg_la.py +++ b/python/sglang/kernels/ops/attention/linear/seg_la.py @@ -158,7 +158,6 @@ def seg_la_kernel( state = state * block_decay + tl.dot(k, v) else: - qk = tl.dot(q, k) * softmax_scale decays = tl.exp(decay_scale * (offs_b[:, None] - offs_b[None, :])) decays = tl.where(offs_b[None, :] <= offs_b[:, None], decays, 0.0) diff --git a/python/sglang/kernels/ops/attention/metadata.py b/python/sglang/kernels/ops/attention/metadata.py index 85426d045..88230f5e3 100644 --- a/python/sglang/kernels/ops/attention/metadata.py +++ b/python/sglang/kernels/ops/attention/metadata.py @@ -496,9 +496,9 @@ def draft_extend_set_metadata( row tails keep stale values that attention kernels never read past cache_seqlens, matching the eager replay path's bounded writes. """ - assert ( - page_size > 0 and (page_size & (page_size - 1)) == 0 - ), f"page_size must be a power of two, got {page_size}" + assert page_size > 0 and (page_size & (page_size - 1)) == 0, ( + f"page_size must be a power of two, got {page_size}" + ) batch_size = cache_seqlens_int32.shape[0] max_seq_pages = page_table.shape[1] @@ -588,9 +588,9 @@ def normal_decode_set_metadata( page_table / swa_page_table row is (re)written; the tail keeps stale values across CUDA-graph replays, so consumers must bound reads by cache_seqlens. """ - assert ( - page_size > 0 and (page_size & (page_size - 1)) == 0 - ), f"page_size must be a power of two, got {page_size}" + assert page_size > 0 and (page_size & (page_size - 1)) == 0, ( + f"page_size must be a power of two, got {page_size}" + ) batch_size = cache_seqlens_int32.shape[0] device = seq_lens.device diff --git a/python/sglang/kernels/ops/attention/minimax_qknorm_rope.py b/python/sglang/kernels/ops/attention/minimax_qknorm_rope.py index a67d25928..201186779 100644 --- a/python/sglang/kernels/ops/attention/minimax_qknorm_rope.py +++ b/python/sglang/kernels/ops/attention/minimax_qknorm_rope.py @@ -113,9 +113,9 @@ def minimax_qknorm_rope_grouped( """ groups = [(w, off, cnt) for (w, off, cnt) in groups if cnt > 0] num_groups = len(groups) - assert ( - 1 <= num_groups <= _MAX_GROUPS - ), f"need 1..{_MAX_GROUPS} groups, got {num_groups}" + assert 1 <= num_groups <= _MAX_GROUPS, ( + f"need 1..{_MAX_GROUPS} groups, got {num_groups}" + ) weights: List[torch.Tensor] = [g[0] for g in groups] offsets: List[int] = [int(g[1]) for g in groups] diff --git a/python/sglang/kernels/ops/attention/minimax_sparse/decode/topk_sparse.py b/python/sglang/kernels/ops/attention/minimax_sparse/decode/topk_sparse.py index c981585a7..c27cebe74 100644 --- a/python/sglang/kernels/ops/attention/minimax_sparse/decode/topk_sparse.py +++ b/python/sglang/kernels/ops/attention/minimax_sparse/decode/topk_sparse.py @@ -331,9 +331,9 @@ def flash_decode_with_gqa_share_sparse( max_slots, num_kv_heads, _ = k_cache.shape assert slot_ids.shape[0] == batch_size and seq_lens.shape[0] == batch_size assert topk_idx.shape[0] == num_kv_heads - assert ( - triton.next_power_of_2(block_size) == block_size - ), f"block_size must be a power of 2, but got {block_size}" + assert triton.next_power_of_2(block_size) == block_size, ( + f"block_size must be a power of 2, but got {block_size}" + ) # assert slot_ids.max() < max_slots, f"get slot_ids {slot_ids}, but kv_cache shape is {kv_cache.shape}" max_kv_len = req_to_token.shape[1] # gqa diff --git a/python/sglang/kernels/ops/attention/minimax_sparse/prefill/flash_with_topk_idx.py b/python/sglang/kernels/ops/attention/minimax_sparse/prefill/flash_with_topk_idx.py index 98e370d22..932a08e0e 100644 --- a/python/sglang/kernels/ops/attention/minimax_sparse/prefill/flash_with_topk_idx.py +++ b/python/sglang/kernels/ops/attention/minimax_sparse/prefill/flash_with_topk_idx.py @@ -495,9 +495,9 @@ def flash_prefill_with_topk_index( assert qk_head_dim <= 256 and v_head_dim <= 256, "head_dim must be less than 256" if sink is not None: assert sink.shape[0] == num_heads and sink.shape[1] == qk_head_dim - assert ( - init_blocks + local_blocks <= topk - ), "init_blocks + local_blocks must be less than topk" + assert init_blocks + local_blocks <= topk, ( + "init_blocks + local_blocks must be less than topk" + ) if sm_scale is None: sm_scale = qk_head_dim**-0.5 # q_scale multiplies every Q-side logit (QK dot and sink), so it folds into diff --git a/python/sglang/kernels/ops/attention/mla_kv_pack_quantize_fp8.py b/python/sglang/kernels/ops/attention/mla_kv_pack_quantize_fp8.py index 2d221af4d..e78729537 100644 --- a/python/sglang/kernels/ops/attention/mla_kv_pack_quantize_fp8.py +++ b/python/sglang/kernels/ops/attention/mla_kv_pack_quantize_fp8.py @@ -189,21 +189,21 @@ def mla_kv_pack_quantize_fp8( torch.bfloat16, torch.float16, ), f"k_nope must be bf16/fp16, got {k_nope.dtype}" - assert ( - k_pe.dtype == k_nope.dtype and v.dtype == k_nope.dtype - ), "k_nope, k_pe, v must share dtype" + assert k_pe.dtype == k_nope.dtype and v.dtype == k_nope.dtype, ( + "k_nope, k_pe, v must share dtype" + ) assert fp8_dtype in (torch.float8_e4m3fn, torch.float8_e5m2) s, num_heads, qk_nope = k_nope.shape qk_rope = k_pe.shape[-1] v_head = v.shape[-1] - assert ( - v.shape[0] == s and v.shape[1] == num_heads - ), f"v shape {tuple(v.shape)} mismatches k_nope {tuple(k_nope.shape)}" - assert ( - k_pe.shape[0] == s - ), f"k_pe first dim {k_pe.shape[0]} mismatches k_nope first dim {s}" + assert v.shape[0] == s and v.shape[1] == num_heads, ( + f"v shape {tuple(v.shape)} mismatches k_nope {tuple(k_nope.shape)}" + ) + assert k_pe.shape[0] == s, ( + f"k_pe first dim {k_pe.shape[0]} mismatches k_nope first dim {s}" + ) assert k_nope.stride(-1) == 1, "k_nope must have stride-1 inner dim" assert v.stride(-1) == 1, "v must have stride-1 inner dim" assert k_pe.stride(-1) == 1, "k_pe must have stride-1 inner dim" diff --git a/python/sglang/kernels/ops/attention/rocm_mla_decode_rope.py b/python/sglang/kernels/ops/attention/rocm_mla_decode_rope.py index 837dacee5..51786fc66 100644 --- a/python/sglang/kernels/ops/attention/rocm_mla_decode_rope.py +++ b/python/sglang/kernels/ops/attention/rocm_mla_decode_rope.py @@ -326,9 +326,9 @@ def _decode_grouped_att_m_fwd_rope( is_neox_style=True, ): if use_rope: - assert ( - k_pe_tokens_out is not None - ), "We must output the k_pe tokens with rope applied if rope fusion enabled." + assert k_pe_tokens_out is not None, ( + "We must output the k_pe tokens with rope applied if rope fusion enabled." + ) BLOCK = 32 diff --git a/python/sglang/kernels/ops/attention/score_mod.py b/python/sglang/kernels/ops/attention/score_mod.py index d39c20fbb..98540169d 100644 --- a/python/sglang/kernels/ops/attention/score_mod.py +++ b/python/sglang/kernels/ops/attention/score_mod.py @@ -30,9 +30,9 @@ import triton.language as tl def unpack_aux_tensors(score_mod, aux_tensors): if score_mod is None: return None, 0, 0, 0 - assert ( - aux_tensors is not None and len(aux_tensors) == 1 - ), "Triton score_mod currently requires exactly one aux tensor" + assert aux_tensors is not None and len(aux_tensors) == 1, ( + "Triton score_mod currently requires exactly one aux tensor" + ) aux0 = aux_tensors[0] assert aux0.dim() == 3 and aux0.stride(2) == 1, ( f"aux_tensors[0] must be 3D with a contiguous last dim, " diff --git a/python/sglang/kernels/ops/attention/sparse_mla_q8kv8_prefill_sm90.py b/python/sglang/kernels/ops/attention/sparse_mla_q8kv8_prefill_sm90.py index bd014c0f8..312febb67 100644 --- a/python/sglang/kernels/ops/attention/sparse_mla_q8kv8_prefill_sm90.py +++ b/python/sglang/kernels/ops/attention/sparse_mla_q8kv8_prefill_sm90.py @@ -303,7 +303,7 @@ def sparse_mla_q8kv8_prefill_fwd( ) if indices.ndim != 3: raise ValueError( - "indices must have shape (s_q, h_kv, topk), " f"got {tuple(indices.shape)}" + f"indices must have shape (s_q, h_kv, topk), got {tuple(indices.shape)}" ) s_q, h_q, d_qk = q.shape @@ -362,8 +362,7 @@ def sparse_mla_q8kv8_prefill_fwd( if indices.shape[:2] != (s_q, h_kv): raise ValueError( - "indices must have shape " - f"({s_q}, {h_kv}, topk), got {tuple(indices.shape)}" + f"indices must have shape ({s_q}, {h_kv}, topk), got {tuple(indices.shape)}" ) if indices.dtype != torch.int32: @@ -385,14 +384,13 @@ def sparse_mla_q8kv8_prefill_fwd( raise ValueError("topk_length must be a CUDA tensor") if topk_length.device != device: raise ValueError( - "topk_length must be on q's device " - f"{device}, got {topk_length.device}" + f"topk_length must be on q's device {device}, got {topk_length.device}" ) if not topk_length.is_contiguous(): raise ValueError("topk_length must be contiguous") if torch.any(topk_length < 0).item() or torch.any(topk_length > topk).item(): raise ValueError( - "topk_length values must satisfy " f"0 <= topk_length <= topk ({topk})" + f"topk_length values must satisfy 0 <= topk_length <= topk ({topk})" ) if d_v != 512: diff --git a/python/sglang/kernels/ops/communication/inkling_all_reduce.py b/python/sglang/kernels/ops/communication/inkling_all_reduce.py index 050d57ef4..e7b7af2e7 100644 --- a/python/sglang/kernels/ops/communication/inkling_all_reduce.py +++ b/python/sglang/kernels/ops/communication/inkling_all_reduce.py @@ -121,9 +121,9 @@ _AR_TUNED_TP8 = { } _AR_TUNED = {4: _AR_TUNED_TP4, 8: _AR_TUNED_TP8} _AR_TUNED_TOKENS = sorted(_AR_TUNED_TP4) # same token grid for every table -assert all( - set(t) == set(_AR_TUNED_TP4) for t in _AR_TUNED.values() -), "all tuned tables must share the same token grid" +assert all(set(t) == set(_AR_TUNED_TP4) for t in _AR_TUNED.values()), ( + "all tuned tables must share the same token grid" +) def select_ar_config(num_tokens: int, world_size: int = 4): diff --git a/python/sglang/kernels/ops/communication/mp.py b/python/sglang/kernels/ops/communication/mp.py index 957ac0647..817d158f1 100644 --- a/python/sglang/kernels/ops/communication/mp.py +++ b/python/sglang/kernels/ops/communication/mp.py @@ -161,8 +161,7 @@ def multigpu_launch( for N in num_gpus: if N <= 1 or N > num_devices: raise ValueError( - f"Invalid number of GPUs requested: {N} " - f"(available: {num_devices})" + f"Invalid number of GPUs requested: {N} (available: {num_devices})" ) os.environ[env_key] = "1" os.environ[pid_key] = str(os.getpid()) diff --git a/python/sglang/kernels/ops/diffusion/attention/sana_wm_gdn_chunkwise_triton.py b/python/sglang/kernels/ops/diffusion/attention/sana_wm_gdn_chunkwise_triton.py index cf697d07e..1272fae4a 100644 --- a/python/sglang/kernels/ops/diffusion/attention/sana_wm_gdn_chunkwise_triton.py +++ b/python/sglang/kernels/ops/diffusion/attention/sana_wm_gdn_chunkwise_triton.py @@ -1666,14 +1666,14 @@ def cam_scan_bidi_chunkwise( q, k, v: camera-prepared ``(B, H, D, N)`` fp32; beta: ``(B, H, F, S)`` fp32; decay: ``(B, H, F)`` fp32. Returns ``(B, H, D, N)`` fp32. """ - assert ( - q.shape == k.shape == v.shape - ), f"q/k/v shape mismatch: {q.shape} {k.shape} {v.shape}" + assert q.shape == k.shape == v.shape, ( + f"q/k/v shape mismatch: {q.shape} {k.shape} {v.shape}" + ) assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous() assert beta.is_contiguous() and decay.is_contiguous() - assert ( - q.dtype == torch.float32 - ), f"cam_scan_bidi_chunkwise requires fp32 q/k/v, got {q.dtype}" + assert q.dtype == torch.float32, ( + f"cam_scan_bidi_chunkwise requires fp32 q/k/v, got {q.dtype}" + ) B, H, D, N = q.shape F = beta.shape[2] diff --git a/python/sglang/kernels/ops/diffusion/attention/sana_wm_gdn_triton.py b/python/sglang/kernels/ops/diffusion/attention/sana_wm_gdn_triton.py index 07760f61c..e7873d852 100644 --- a/python/sglang/kernels/ops/diffusion/attention/sana_wm_gdn_triton.py +++ b/python/sglang/kernels/ops/diffusion/attention/sana_wm_gdn_triton.py @@ -175,9 +175,9 @@ def fused_qk_inv_rms( qkv: (B, N, 3, H, D) contiguous. Returns (q_inv_rms, k_inv_rms), each (B, N) float32. """ assert qkv.is_contiguous(), "qkv must be contiguous (B, N, 3, H, D)" - assert ( - qkv.dim() == 5 and qkv.shape[2] == 3 - ), f"expected (B, N, 3, H, D), got {tuple(qkv.shape)}" + assert qkv.dim() == 5 and qkv.shape[2] == 3, ( + f"expected (B, N, 3, H, D), got {tuple(qkv.shape)}" + ) B, N, _, H, D = qkv.shape C = H * D q_inv_rms = torch.empty((B, N), dtype=torch.float32, device=qkv.device) diff --git a/python/sglang/kernels/ops/diffusion/layout/usp_relayout_jit.py b/python/sglang/kernels/ops/diffusion/layout/usp_relayout_jit.py index 38321b4f1..816706a9b 100644 --- a/python/sglang/kernels/ops/diffusion/layout/usp_relayout_jit.py +++ b/python/sglang/kernels/ops/diffusion/layout/usp_relayout_jit.py @@ -26,7 +26,7 @@ def _jit_usp_relayout_module(dtype: torch.dtype) -> Module: cuda_wrappers=[ ( "usp_merge_heads", - "usp_relayout::" f"UspMergeHeadsKernel<{args}>::run", + f"usp_relayout::UspMergeHeadsKernel<{args}>::run", ), ], ) diff --git a/python/sglang/kernels/ops/diffusion/modulate/scale_shift_triton.py b/python/sglang/kernels/ops/diffusion/modulate/scale_shift_triton.py index 6741c1a2e..7a9a3b61f 100644 --- a/python/sglang/kernels/ops/diffusion/modulate/scale_shift_triton.py +++ b/python/sglang/kernels/ops/diffusion/modulate/scale_shift_triton.py @@ -413,9 +413,9 @@ def fuse_scale_shift_kernel( num_warps = 2 if block_n == 64 else 4 grid = (rows, triton.cdiv(C, block_n)) num_frames = scale.shape[1] - assert ( - L % num_frames == 0 - ), "seq_len must be divisible by num_frames for 4D scale/shift" + assert L % num_frames == 0, ( + "seq_len must be divisible by num_frames for 4D scale/shift" + ) frame_seqlen = L // num_frames # Compact scale [B, F, 1, C] -> [B*F, C] (per-frame) diff --git a/python/sglang/kernels/ops/diffusion/norm/fused_residual_norm_flydsl.py b/python/sglang/kernels/ops/diffusion/norm/fused_residual_norm_flydsl.py index cf5e45428..ddf5ffc28 100644 --- a/python/sglang/kernels/ops/diffusion/norm/fused_residual_norm_flydsl.py +++ b/python/sglang/kernels/ops/diffusion/norm/fused_residual_norm_flydsl.py @@ -51,9 +51,9 @@ def _build_fused_norm_module(D: int, is_rms: bool, has_gate: bool, has_weight: b VEC = _VEC NUM_WAVES = _NUM_WAVES BLOCK = NUM_WAVES * WARP_SIZE - assert ( - D % FLYDSL_NORM_MIN_ALIGNED_DIM == 0 - ), f"FlyDSL fused_residual_norm requires D % {FLYDSL_NORM_MIN_ALIGNED_DIM} == 0, got D={D}" + assert D % FLYDSL_NORM_MIN_ALIGNED_DIM == 0, ( + f"FlyDSL fused_residual_norm requires D % {FLYDSL_NORM_MIN_ALIGNED_DIM} == 0, got D={D}" + ) NUM_ITERS = D // (BLOCK * VEC) @flyc.kernel(known_block_size=[BLOCK, 1, 1]) @@ -543,9 +543,9 @@ def _build_norm_scale_shift_module(D: int, is_rms: bool, has_weight: bool): VEC = _VEC NUM_WAVES = _NUM_WAVES BLOCK = NUM_WAVES * WARP_SIZE - assert ( - D % FLYDSL_NORM_MIN_ALIGNED_DIM == 0 - ), f"FlyDSL norm_scale_shift requires D % {FLYDSL_NORM_MIN_ALIGNED_DIM} == 0, got D={D}" + assert D % FLYDSL_NORM_MIN_ALIGNED_DIM == 0, ( + f"FlyDSL norm_scale_shift requires D % {FLYDSL_NORM_MIN_ALIGNED_DIM} == 0, got D={D}" + ) NUM_ITERS = D // (BLOCK * VEC) @flyc.kernel(known_block_size=[BLOCK, 1, 1]) diff --git a/python/sglang/kernels/ops/elementwise/elementwise.py b/python/sglang/kernels/ops/elementwise/elementwise.py index 326c89ccf..85ba16c84 100644 --- a/python/sglang/kernels/ops/elementwise/elementwise.py +++ b/python/sglang/kernels/ops/elementwise/elementwise.py @@ -103,9 +103,9 @@ fused_dual_residual_rmsnorm_kernel_autotune = rmsnorm_autotune( def fused_dual_residual_rmsnorm(x, residual, weight1, weight2, eps, autotune=False): assert len(x.shape) == 2 - assert ( - x.shape == residual.shape and x.dtype == residual.dtype - ), f"{x.shape=} {residual.shape=} {x.dtype=} {residual.dtype=}" + assert x.shape == residual.shape and x.dtype == residual.dtype, ( + f"{x.shape=} {residual.shape=} {x.dtype=} {residual.dtype=}" + ) output, mid = torch.empty_like(x), torch.empty_like(x) bs, hidden_dim = x.shape if autotune: @@ -434,9 +434,9 @@ def fused_sigmoid_mul( gate_stride_head = gate.stride(1) else: # Flat path: both tensors have the same shape - assert ( - attn_output.shape == gate.shape - ), "attn_output and gate must have the same shape" + assert attn_output.shape == gate.shape, ( + "attn_output and gate must have the same shape" + ) hidden_dim = attn_output.shape[-1] num_tokens = attn_output.numel() // hidden_dim head_dim = hidden_dim diff --git a/python/sglang/kernels/ops/gemm/cutedsl_dsv3_fused_a_gemm.py b/python/sglang/kernels/ops/gemm/cutedsl_dsv3_fused_a_gemm.py index 0df7e9470..9cd83c793 100644 --- a/python/sglang/kernels/ops/gemm/cutedsl_dsv3_fused_a_gemm.py +++ b/python/sglang/kernels/ops/gemm/cutedsl_dsv3_fused_a_gemm.py @@ -338,9 +338,9 @@ def _dsv3_fused_a_gemm_run(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Te assert mat_a.dtype == torch.bfloat16 and mat_b.dtype == torch.bfloat16 assert K % 1024 == 0, f"K must be a multiple of 1024, got {K}" assert N % TILE_M == 0, f"N must be a multiple of {TILE_M}, got {N}" - assert ( - tuple(mat_b.shape) == (K, N) and mat_b.stride(0) == 1 - ), "mat_b must be [K, N] column-major" + assert tuple(mat_b.shape) == (K, N) and mat_b.stride(0) == 1, ( + "mat_b must be [K, N] column-major" + ) assert 1 <= M <= 16, "num_tokens must be in [1, 16]" assert mat_a.stride(1) == 1, "mat_a must be row-major [M, K]" diff --git a/python/sglang/kernels/ops/gemm/flashinfer_pr4266_dense_bf16_gemm_sm100_splitk.py b/python/sglang/kernels/ops/gemm/flashinfer_pr4266_dense_bf16_gemm_sm100_splitk.py index 9597f836d..92ec82470 100644 --- a/python/sglang/kernels/ops/gemm/flashinfer_pr4266_dense_bf16_gemm_sm100_splitk.py +++ b/python/sglang/kernels/ops/gemm/flashinfer_pr4266_dense_bf16_gemm_sm100_splitk.py @@ -229,7 +229,9 @@ def default_tactic(m: int, n: int, k: int) -> SplitKTactic: ab_stages=( _MIN_AB_STAGES if k <= 2 * _CTA_K and m > 8 - else min(max_stages, 6) if k <= 4 * _CTA_K else max_stages + else min(max_stages, 6) + if k <= 4 * _CTA_K + else max_stages ), ) validate_tactic(tactic, m, n, k) diff --git a/python/sglang/kernels/ops/grammar/bitmask_ops.py b/python/sglang/kernels/ops/grammar/bitmask_ops.py index 9a195c006..e439effa1 100644 --- a/python/sglang/kernels/ops/grammar/bitmask_ops.py +++ b/python/sglang/kernels/ops/grammar/bitmask_ops.py @@ -114,9 +114,9 @@ def apply_token_bitmask_inplace_triton( indices = torch.tensor(indices, dtype=torch.int32, device=logits.device) num_rows = indices.shape[0] else: - assert ( - logits_shape[0] == bitmask_shape[0] - ), f"batch size mismatch: logits {logits_shape[0]} vs bitmask {bitmask_shape[0]}" + assert logits_shape[0] == bitmask_shape[0], ( + f"batch size mismatch: logits {logits_shape[0]} vs bitmask {bitmask_shape[0]}" + ) num_rows = logits_shape[0] if NUM_SMS > 0: diff --git a/python/sglang/kernels/ops/kimi_k3/sp_collective.py b/python/sglang/kernels/ops/kimi_k3/sp_collective.py index abb93b474..2d87f84e5 100644 --- a/python/sglang/kernels/ops/kimi_k3/sp_collective.py +++ b/python/sglang/kernels/ops/kimi_k3/sp_collective.py @@ -51,10 +51,7 @@ def _device_name(device: torch.device) -> str: def _table(world_size: int, hidden_size: int, device: torch.device) -> Optional[dict]: path = os.path.join( _CONFIG_DIR, - ( - f"world={world_size},H={hidden_size}," - f"device_name={_device_name(device)}.json" - ), + (f"world={world_size},H={hidden_size},device_name={_device_name(device)}.json"), ) if path not in _TABLES: if os.path.exists(path): diff --git a/python/sglang/kernels/ops/kvcache/cache_ops.py b/python/sglang/kernels/ops/kvcache/cache_ops.py index 6bc1450d5..bc4f71034 100644 --- a/python/sglang/kernels/ops/kvcache/cache_ops.py +++ b/python/sglang/kernels/ops/kvcache/cache_ops.py @@ -100,18 +100,18 @@ def concat_and_cast_mha_k_triton( k_rope: torch.Tensor, ): # The source data type will be implicitly converted to the target data type. - assert ( - len(k.shape) == 3 and len(k_nope.shape) == 3 and len(k_rope.shape) == 3 - ), f"shape should be 3d, but got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}" - assert ( - k.shape[0] == k_nope.shape[0] and k.shape[0] == k_rope.shape[0] - ), f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}" - assert ( - k.shape[1] == k_nope.shape[1] and 1 == k_rope.shape[1] - ), f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}" - assert ( - k.shape[-1] == k_nope.shape[-1] + k_rope.shape[-1] - ), f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}" + assert len(k.shape) == 3 and len(k_nope.shape) == 3 and len(k_rope.shape) == 3, ( + f"shape should be 3d, but got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}" + ) + assert k.shape[0] == k_nope.shape[0] and k.shape[0] == k_rope.shape[0], ( + f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}" + ) + assert k.shape[1] == k_nope.shape[1] and 1 == k_rope.shape[1], ( + f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}" + ) + assert k.shape[-1] == k_nope.shape[-1] + k_rope.shape[-1], ( + f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}" + ) nope_dim = k_nope.shape[-1] rope_dim = k_rope.shape[-1] @@ -638,9 +638,9 @@ def absorbed_bmm_concat_cast_q_fp8( assert q_fp8_pad.shape[0] >= num_tokens and q_fp8_pad.shape[1] >= num_heads assert q_fp8_pad.shape[2] == n_dim + rope_dim # tl.arange / tl.dot constraints - assert ( - k_dim % 16 == 0 and 16 <= k_dim <= 256 - ), "K must be a multiple of 16 in [16, 256]" + assert k_dim % 16 == 0 and 16 <= k_dim <= 256, ( + "K must be a multiple of 16 in [16, 256]" + ) assert (rope_dim & (rope_dim - 1)) == 0, "ROPE must be a power of two" assert n_dim % block_n == 0, "N must be a multiple of block_n" assert q_nope.stride(2) == 1 and q_rope.stride(2) == 1 @@ -680,22 +680,22 @@ def absorbed_bmm_concat_cast_q_fp8( # Largest power-of-2 divisor of K, capped at 128 (K % 16 == 0 # makes this >= 16), unless the caller pinned block_k. blk_k = block_k or min(k_dim & -k_dim, 128) - assert ( - k_dim % blk_k == 0 and blk_k & (blk_k - 1) == 0 and blk_k >= 16 - ), "loop needs BLOCK_K a power-of-2 divisor of K >= 16" + assert k_dim % blk_k == 0 and blk_k & (blk_k - 1) == 0 and blk_k >= 16, ( + "loop needs BLOCK_K a power-of-2 divisor of K >= 16" + ) k_mode = 1 elif v == "two_dot": blk_k = 1 << (k_dim.bit_length() - 1) # largest power of 2 < K k1 = k_dim - blk_k - assert ( - k1 & (k1 - 1) == 0 and k1 >= 16 - ), "two_dot needs K = pow2 + pow2 with both halves >= 16" + assert k1 & (k1 - 1) == 0 and k1 >= 16, ( + "two_dot needs K = pow2 + pow2 with both halves >= 16" + ) k_mode = 2 elif v == "three_dot": blk_k = k_dim // 3 - assert ( - k_dim % 3 == 0 and blk_k & (blk_k - 1) == 0 and blk_k >= 16 - ), "three_dot needs K = 3 * pow2 with pow2 >= 16" + assert k_dim % 3 == 0 and blk_k & (blk_k - 1) == 0 and blk_k >= 16, ( + "three_dot needs K = 3 * pow2 with pow2 >= 16" + ) k_mode = 3 elif v == "pad": blk_k = 1 << k_dim.bit_length() # next power of 2 above K diff --git a/python/sglang/kernels/ops/kvcache/hisparse.py b/python/sglang/kernels/ops/kvcache/hisparse.py index cfebdcbf3..2e2281915 100644 --- a/python/sglang/kernels/ops/kvcache/hisparse.py +++ b/python/sglang/kernels/ops/kvcache/hisparse.py @@ -292,9 +292,9 @@ def _load_cache_to_device_buffer_mla( miss_count: torch.Tensor | None, skip_io: bool, ) -> None: - assert ( - hot_buffer_size >= num_top_k - ), f"hot_buffer_size ({hot_buffer_size}) must be >= num_top_k ({num_top_k})" + assert hot_buffer_size >= num_top_k, ( + f"hot_buffer_size ({hot_buffer_size}) must be >= num_top_k ({num_top_k})" + ) record_miss_plan = miss_src is not None module = _jit_sparse_module( diff --git a/python/sglang/kernels/ops/kvcache/kv_read_table.py b/python/sglang/kernels/ops/kvcache/kv_read_table.py index e19436ea6..9c3141a5e 100644 --- a/python/sglang/kernels/ops/kvcache/kv_read_table.py +++ b/python/sglang/kernels/ops/kvcache/kv_read_table.py @@ -192,9 +192,9 @@ def build_kv_read_table( region's live prefix is written -- never rebound, never tail-cleared. """ bs = int(req_pool_indices.numel()) - assert ( - out.dtype == torch.int32 - ), f"build_kv_read_table: out must be int32, got {out.dtype}" + assert out.dtype == torch.int32, ( + f"build_kv_read_table: out must be int32, got {out.dtype}" + ) assert out.dim() == 2 and out.shape[0] >= bs and out.shape[1] >= max_pages, ( f"build_kv_read_table: out {tuple(out.shape)} cannot hold " f"(bs={bs}, max_pages={max_pages})" @@ -262,7 +262,7 @@ def build_kv_read_table_packed( """ bs = int(req_pool_indices.numel()) assert out.dtype in (torch.int32, torch.int64), ( - f"build_kv_read_table_packed: out must be int32 or int64, got " f"{out.dtype}" + f"build_kv_read_table_packed: out must be int32 or int64, got {out.dtype}" ) assert out.dim() == 1 and out.numel() >= max_tokens, ( f"build_kv_read_table_packed: out {tuple(out.shape)} cannot hold " diff --git a/python/sglang/kernels/ops/kvcache/rope_cache.py b/python/sglang/kernels/ops/kvcache/rope_cache.py index 49b4b8ca9..49615558b 100644 --- a/python/sglang/kernels/ops/kvcache/rope_cache.py +++ b/python/sglang/kernels/ops/kvcache/rope_cache.py @@ -613,37 +613,37 @@ def fused_qk_rope_reshape_and_cache( value_shuffle_layout = False (t_slot,) = slot_mapping.shape - assert ( - t == tk == tv and t_slot <= tk - ), f"Number of tokens should be identical for q, kand v. The number of tokens of slot_mapping should no more than that of q, k and v, {t=} {tk=} {tv=} {t_slot=}" - assert ( - block_size == block_size_v - ), f"block size should be identical for key_cache, and value_cache {block_size} {block_size_v}" - assert ( - kh == vh == kh_cache == vh_cache - ), "KV head should be identical for k, v, key_cache, and value_cache" - assert ( - t_cache == t_cache_v - ), "Number of tokens should be identical for key_cache, and value_cache" + assert t == tk == tv and t_slot <= tk, ( + f"Number of tokens should be identical for q, kand v. The number of tokens of slot_mapping should no more than that of q, k and v, {t=} {tk=} {tv=} {t_slot=}" + ) + assert block_size == block_size_v, ( + f"block size should be identical for key_cache, and value_cache {block_size} {block_size_v}" + ) + assert kh == vh == kh_cache == vh_cache, ( + "KV head should be identical for k, v, key_cache, and value_cache" + ) + assert t_cache == t_cache_v, ( + "Number of tokens should be identical for key_cache, and value_cache" + ) if flash_layout: - assert ( - d == dk == dv == dk_cache == dv_cache - ), "D dimension should be identical for q, k, and v" + assert d == dk == dv == dk_cache == dv_cache, ( + "D dimension should be identical for q, k, and v" + ) else: - assert ( - d == dk == dv == dkx_cache * x_cache == dv_cache - ), "D dimension should be identical for q, k, and v" + assert d == dk == dv == dkx_cache * x_cache == dv_cache, ( + "D dimension should be identical for q, k, and v" + ) assert x_cache == triton.next_power_of_2(x_cache), "x_size should be power of 2" assert d == triton.next_power_of_2(d), "D dimension should be power of 2" - assert block_size == triton.next_power_of_2( - block_size - ), "block_size should be power of 2" + assert block_size == triton.next_power_of_2(block_size), ( + "block_size should be power of 2" + ) assert qh % kh == 0, "Q heads must be multiple of H heads" d_freq = cos_sin.shape[-1] // 2 - assert (d_freq == d // 2) or ( - d_freq == d - ), "cos/sin last dim should be the same or half of the qk last dim" + assert (d_freq == d // 2) or (d_freq == d), ( + "cos/sin last dim should be the same or half of the qk last dim" + ) reuse_freqs_front_part = d_freq == d // 2 if q_out is None: @@ -654,9 +654,9 @@ def fused_qk_rope_reshape_and_cache( if zeros_out is not None: tz, qhz, dz = zeros_out.shape - assert ( - t == tz and qh == qhz and d == dz - ), f"q and zeros shape mismatch {q.shape=} {zeros_out.shape=}" + assert t == tz and qh == qhz and d == dz, ( + f"q and zeros shape mismatch {q.shape=} {zeros_out.shape=}" + ) output_zeros = True elif output_zeros: zeros_out = torch.empty((t, qh, d), dtype=q.dtype, device=q.device) diff --git a/python/sglang/kernels/ops/kvcache/trtllm_mha_page_table.py b/python/sglang/kernels/ops/kvcache/trtllm_mha_page_table.py index 67d23a576..ec54f6fed 100644 --- a/python/sglang/kernels/ops/kvcache/trtllm_mha_page_table.py +++ b/python/sglang/kernels/ops/kvcache/trtllm_mha_page_table.py @@ -113,12 +113,12 @@ def build_trtllm_mha_page_table( ``full_to_swa`` is provided, which then also requires ``swa_page_table``. """ has_swa = full_to_swa is not None - assert has_swa == ( - swa_page_table is not None - ), "full_to_swa and swa_page_table must be provided together" - assert ( - _MHA_KV_INDEX_BLOCK_TOKENS % page_size == 0 - ), f"page_size={page_size} must divide _MHA_KV_INDEX_BLOCK_TOKENS={_MHA_KV_INDEX_BLOCK_TOKENS}" + assert has_swa == (swa_page_table is not None), ( + "full_to_swa and swa_page_table must be provided together" + ) + assert _MHA_KV_INDEX_BLOCK_TOKENS % page_size == 0, ( + f"page_size={page_size} must divide _MHA_KV_INDEX_BLOCK_TOKENS={_MHA_KV_INDEX_BLOCK_TOKENS}" + ) bs, num_pages = page_table.shape full_to_swa_numel = full_to_swa.numel() if has_swa else 0 create_trtllm_mha_kv_indices_triton[ diff --git a/python/sglang/kernels/ops/layernorm/__init__.py b/python/sglang/kernels/ops/layernorm/__init__.py index fa17572fa..00f857771 100644 --- a/python/sglang/kernels/ops/layernorm/__init__.py +++ b/python/sglang/kernels/ops/layernorm/__init__.py @@ -385,8 +385,7 @@ class GemmaFusedAddRMSNormOp(BaseFusedOp): "(rocm-triton, sglang.kernels.jit)." ), KernelBackend.TORCH: ( - "Gemma-style fused residual-add + RMS normalization " - "(pure-torch reference)." + "Gemma-style fused residual-add + RMS normalization (pure-torch reference)." ), } diff --git a/python/sglang/kernels/ops/layernorm/mhc.py b/python/sglang/kernels/ops/layernorm/mhc.py index d34cca1b4..4784ebf0d 100644 --- a/python/sglang/kernels/ops/layernorm/mhc.py +++ b/python/sglang/kernels/ops/layernorm/mhc.py @@ -1103,9 +1103,9 @@ def mhc_pre( gemm_out_sqrsum = torch.empty( n_splits, num_tokens, dtype=torch.float32, device=residual.device ) - assert ( - n_splits == 1 - ), "The simple TileLang version gemm_sqrsum doesn't support split-k" + assert n_splits == 1, ( + "The simple TileLang version gemm_sqrsum doesn't support split-k" + ) _mhc_pre_gemm_sqrsum_dispatch()( residual_flat.view(num_tokens, hc_mult * hidden_size), fn_flat, @@ -1119,9 +1119,9 @@ def mhc_pre( if norm_weight is not None: assert norm_eps is not None, "norm_eps required when norm_weight is provided" - assert norm_weight.shape == ( - hidden_size, - ), f"norm_weight shape {tuple(norm_weight.shape)} != (hidden_size={hidden_size},)" + assert norm_weight.shape == (hidden_size,), ( + f"norm_weight shape {tuple(norm_weight.shape)} != (hidden_size={hidden_size},)" + ) norm_weight_bf = ( norm_weight.bfloat16() if norm_weight.dtype != torch.bfloat16 diff --git a/python/sglang/kernels/ops/lplb/cuda_solver.py b/python/sglang/kernels/ops/lplb/cuda_solver.py index 61ab5cfe6..3a658cdf2 100644 --- a/python/sglang/kernels/ops/lplb/cuda_solver.py +++ b/python/sglang/kernels/ops/lplb/cuda_solver.py @@ -250,9 +250,9 @@ def dispatch_probability( if random_vals is None: random_vals = torch.rand(n, dtype=torch.float32, device=topk_ids.device) else: - assert random_vals.shape == ( - n, - ), f"random_vals must be shape ({n},), got {tuple(random_vals.shape)}" + assert random_vals.shape == (n,), ( + f"random_vals must be shape ({n},), got {tuple(random_vals.shape)}" + ) module = _dispatch_module(max_copies, DISPATCH_BLOCK_DIM) module.dispatch_probability(out, flat_ids, log2phy_prob, map32, random_vals) return out.view(original_shape).to(topk_ids.dtype) @@ -299,9 +299,9 @@ def dispatch_probability_torch_reference( n = flat_ids.shape[0] num_logical, max_copies = log2phy_prob.shape assert log2phy_map.shape == (num_logical, max_copies) - assert random_vals.shape == ( - n, - ), f"random_vals must be shape ({n},), got {tuple(random_vals.shape)}" + assert random_vals.shape == (n,), ( + f"random_vals must be shape ({n},), got {tuple(random_vals.shape)}" + ) # Gather per-row probabilities and physical maps. probs = log2phy_prob[flat_ids] # (N, max_copies), float32 diff --git a/python/sglang/kernels/ops/lplb/shmem_budget.py b/python/sglang/kernels/ops/lplb/shmem_budget.py index c91f04e19..99b3d4946 100644 --- a/python/sglang/kernels/ops/lplb/shmem_budget.py +++ b/python/sglang/kernels/ops/lplb/shmem_budget.py @@ -111,8 +111,8 @@ def assert_fits(nc: int, nv: int, gpu: str = "h100") -> None: cap = gpu_budget_bytes(gpu) if used > cap: raise ValueError( - f"fused IPM kernel needs {used/1024:.1f} KiB of shared memory for " - f"NC={nc}, NV={nv}, but {gpu} allows {cap/1024:.1f} KiB/block. " + f"fused IPM kernel needs {used / 1024:.1f} KiB of shared memory for " + f"NC={nc}, NV={nv}, but {gpu} allows {cap / 1024:.1f} KiB/block. " f"Either reduce problem size or switch to a tiled design." ) @@ -145,8 +145,8 @@ def report(nc: int, nv: int, gpu: str = "h100") -> str: status = "FITS" if bd.total_bytes <= cap else "OVER BUDGET" return ( f"[shmem] NC={nc} NV={nv} gpu={gpu} | " - f"A={bd.a_bytes/1024:.1f}K " - f"ata={bd.ata_bytes/1024:.1f}K " - f"rest={(bd.c_bytes+bd.x_bytes+bd.rhs_bytes+bd.d_bytes)/1024:.1f}K | " - f"total={bd.total_bytes/1024:.1f}K / {cap/1024:.1f}K {status}" + f"A={bd.a_bytes / 1024:.1f}K " + f"ata={bd.ata_bytes / 1024:.1f}K " + f"rest={(bd.c_bytes + bd.x_bytes + bd.rhs_bytes + bd.d_bytes) / 1024:.1f}K | " + f"total={bd.total_bytes / 1024:.1f}K / {cap / 1024:.1f}K {status}" ) diff --git a/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py b/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py index be5cca881..321a15d0d 100644 --- a/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py +++ b/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py @@ -332,7 +332,6 @@ def _causal_conv1d_fwd_kernel( # continuous batching matrix_w = w_col0 matrix_x = col0 for j in tl.static_range(KERNEL_WIDTH): - if KERNEL_WIDTH == 2: if j == 1: # KERNEL_WIDTH-1: matrix_w = w_col1 @@ -502,9 +501,9 @@ def causal_conv1d_fn( assert padded_batch == cache_indices.size(0) if has_initial_state is not None: assert has_initial_state.size() == (padded_batch,) - assert ( - conv_states is not None - ), "ERROR: `has_initial_state` is used, which needs also `conv_states`" + assert conv_states is not None, ( + "ERROR: `has_initial_state` is used, which needs also `conv_states`" + ) assert weight.stride(1) == 1 assert (dim, width) == weight.shape assert is_channel_last, "Need to run in channel-last layout" @@ -1053,9 +1052,9 @@ def causal_conv1d_update( if validate_data: assert dim == weight.size(0) - assert ( - conv_state.stride(-2) == 1 - ), f"ERROR: expect contiguous along feat-dim of conv_state (currently stride={conv_state.stride()})" + assert conv_state.stride(-2) == 1, ( + f"ERROR: expect contiguous along feat-dim of conv_state (currently stride={conv_state.stride()})" + ) assert state_len >= width - 1 # when above happens, we don't shift-left to keep any records in conv_state assert dim == conv_state.size(1) diff --git a/python/sglang/kernels/ops/mamba/triton_ops/mamba_ssm.py b/python/sglang/kernels/ops/mamba/triton_ops/mamba_ssm.py index 3ed64d131..d58c8798f 100644 --- a/python/sglang/kernels/ops/mamba/triton_ops/mamba_ssm.py +++ b/python/sglang/kernels/ops/mamba/triton_ops/mamba_ssm.py @@ -52,8 +52,9 @@ cvt.rs.f16x2.f32 $0, $2, $1, $3; @triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) @triton.heuristics( { - "HAS_STATE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] - is not None + "HAS_STATE_BATCH_INDICES": lambda args: ( + args["state_batch_indices_ptr"] is not None + ) } ) @triton.heuristics( @@ -61,24 +62,23 @@ cvt.rs.f16x2.f32 $0, $2, $1, $3; ) @triton.heuristics( { - "CACHE_INTERMEDIATE_STATES": lambda args: args["intermediate_states_buffer"] - is not None + "CACHE_INTERMEDIATE_STATES": lambda args: ( + args["intermediate_states_buffer"] is not None + ) } ) @triton.heuristics( { - "HAS_EAGLE_TREE_CUSTOM_ATTN_MASK": lambda args: args[ - "retrieve_parent_token_ptr" - ] - is not None + "HAS_EAGLE_TREE_CUSTOM_ATTN_MASK": lambda args: ( + args["retrieve_parent_token_ptr"] is not None + ) } ) @triton.heuristics( { - "HAS_INTERMEDIATE_STATE_INDICES": lambda args: args[ - "intermediate_state_indices_ptr" - ] - is not None + "HAS_INTERMEDIATE_STATE_INDICES": lambda args: ( + args["intermediate_state_indices_ptr"] is not None + ) } ) @triton.jit(do_not_specialize=["T"]) diff --git a/python/sglang/kernels/ops/mamba/triton_ops/ssd_chunk_scan.py b/python/sglang/kernels/ops/mamba/triton_ops/ssd_chunk_scan.py index 52b197139..5fd4fb138 100644 --- a/python/sglang/kernels/ops/mamba/triton_ops/ssd_chunk_scan.py +++ b/python/sglang/kernels/ops/mamba/triton_ops/ssd_chunk_scan.py @@ -176,7 +176,6 @@ def _chunk_scan_fwd_kernel( ) # if a seq is changed exactly on boundary or (c_off > 0) # implies a new example (pseudo chunk) ): - # - replace prev_states_ptr with init_states prev_states_ptr = ( initstates_ptr @@ -193,7 +192,6 @@ def _chunk_scan_fwd_kernel( # - handle chunk state limit if HAS_INITSTATES: - # have to split this if otherwise compilation will have problems dA_cs_m_boundary = 0.0 @@ -214,7 +212,6 @@ def _chunk_scan_fwd_kernel( # (logical) chunk indices. if (c_idx == c_idx_n) or c_off > 0: - # get the next offset c_off_n = tl.load( chunk_offsets_ptr + (pid_c + 1), @@ -265,7 +262,6 @@ def _chunk_scan_fwd_kernel( + offs_k_dstate[:, None] * prev_states_dstate ) if HAS_SEQ_IDX: - if not HAS_INITSTATES: # - this is for continuous batching where there is no init states scale_m = tl.where(seq_idx_m == seq_idx_prev, tl.exp(dA_cs_m), 0.0) @@ -455,9 +451,9 @@ def _chunk_scan_fwd( # with initial states, we need to take care of how # seq_idx crosses the boundaries assert batch == 1, "chunk scan only supports initial states with batch 1" - assert ( - chunk_indices is not None and chunk_offsets is not None - ), "chunk_indices and chunk_offsets should have been set" + assert chunk_indices is not None and chunk_offsets is not None, ( + "chunk_indices and chunk_offsets should have been set" + ) else: chunk_indices, chunk_offsets = None, None else: diff --git a/python/sglang/kernels/ops/mamba/triton_ops/ssd_chunk_state.py b/python/sglang/kernels/ops/mamba/triton_ops/ssd_chunk_state.py index 162d859d4..d33808a51 100644 --- a/python/sglang/kernels/ops/mamba/triton_ops/ssd_chunk_state.py +++ b/python/sglang/kernels/ops/mamba/triton_ops/ssd_chunk_state.py @@ -390,7 +390,6 @@ def _chunk_state_varlen_kernel( # - if start_idx < pid_c * chunk_size, then we need to take the past_states_ptrs # - if state_idx >= pid * chunk_size, then we need to insert initstates if (start_idx < pid_c * chunk_size) or (HAS_INITSTATES): # first chunk - dA_cs_boundary = 0.0 # default if not HAS_INITSTATES: @@ -399,7 +398,6 @@ def _chunk_state_varlen_kernel( + offs_n[None, :] * stride_chunk_states_dstate ) else: - # - this seems repetitive, buts its to help the compiler if start_idx < pid_c * chunk_size: past_states_ptrs = chunk_states_ptr + ( diff --git a/python/sglang/kernels/ops/mamba/triton_ops/ssd_combined.py b/python/sglang/kernels/ops/mamba/triton_ops/ssd_combined.py index c7f16e70e..b79ff2365 100644 --- a/python/sglang/kernels/ops/mamba/triton_ops/ssd_combined.py +++ b/python/sglang/kernels/ops/mamba/triton_ops/ssd_combined.py @@ -163,9 +163,9 @@ def _mamba_chunk_scan_combined_fwd( if cu_seqlens is None: return out_x, dt, dA_cumsum, states, final_states else: - assert ( - batch == 1 - ), "passing cu_seqlens to get the varlen states is only supported if batch dimension is 1" + assert batch == 1, ( + "passing cu_seqlens to get the varlen states is only supported if batch dimension is 1" + ) varlen_states = chunk_state_varlen( B.squeeze(0), x.squeeze(0), @@ -223,9 +223,9 @@ def mamba_chunk_scan_combined( if not return_varlen_states: cu_seqlens = None else: - assert ( - cu_seqlens is not None - ), "cu_seqlens must be provided if return_varlen_states is True" + assert cu_seqlens is not None, ( + "cu_seqlens must be provided if return_varlen_states is True" + ) out_x, dt_out, dA_cumsum, states, final_states, *rest = ( _mamba_chunk_scan_combined_fwd( x, diff --git a/python/sglang/kernels/ops/mamba/triton_ops/ssd_state_passing.py b/python/sglang/kernels/ops/mamba/triton_ops/ssd_state_passing.py index d448a1d5c..c25ea23a1 100644 --- a/python/sglang/kernels/ops/mamba/triton_ops/ssd_state_passing.py +++ b/python/sglang/kernels/ops/mamba/triton_ops/ssd_state_passing.py @@ -189,15 +189,15 @@ def _state_passing_fwd( # - if cu_seqlens is provided, then the initial states # are used for continuous batching. In which case we # require seq_idx to be provided - assert ( - seq_idx is not None - ), "seq_idx must be provided for continuous batching" + assert seq_idx is not None, ( + "seq_idx must be provided for continuous batching" + ) # - we also need chunk_offsets to be provided, to account # for computation of dA_cumsum from the start of the # sequence - assert ( - chunk_offsets is not None - ), "chunk_offsets must be provided for continuous batching" + assert chunk_offsets is not None, ( + "chunk_offsets must be provided for continuous batching" + ) else: # - this is the regular batching case, where initial # states are used are for each example of the batch. diff --git a/python/sglang/kernels/ops/moe/ep_moe_kernels.py b/python/sglang/kernels/ops/moe/ep_moe_kernels.py index 20263ae3a..98c6c78be 100644 --- a/python/sglang/kernels/ops/moe/ep_moe_kernels.py +++ b/python/sglang/kernels/ops/moe/ep_moe_kernels.py @@ -418,16 +418,16 @@ def silu_and_mul_masked_post_quant_fwd( if output_scale.dtype == torch.int32: assert scale_ue8m0, "packed int32 scales are UE8M0 by definition" - assert ( - num_real_tokens is not None and topk is not None - ), "the packed schedule sizes its grid from num_real_tokens * topk" + assert num_real_tokens is not None and topk is not None, ( + "the packed schedule sizes its grid from num_real_tokens * topk" + ) E, m_max, _ = input.shape G = size_n // quant_group_size assert G % 4 == 0, "packed UE8M0 path requires num_groups % 4 == 0" BLOCK_N = quant_group_size * 4 - assert ( - size_n % BLOCK_N == 0 - ), "packed UE8M0 path requires size_n % (4*group) == 0" + assert size_n % BLOCK_N == 0, ( + "packed UE8M0 path requires size_n % (4*group) == 0" + ) hidden_dim_split = size_n // BLOCK_N assert tuple(output_scale.shape) == (E, hidden_dim_split, m_max) @@ -1175,9 +1175,9 @@ def ep_scatter( is_fp8 = recv_x_scale is not None and recv_x.dtype != torch.bfloat16 if is_fp8: - assert ( - recv_x_scale.dtype == output_tensor_scale.dtype - ), f"recv_x_scale.dtype: {recv_x_scale.dtype}, output_tensor_scale.dtype: {output_tensor_scale.dtype}" + assert recv_x_scale.dtype == output_tensor_scale.dtype, ( + f"recv_x_scale.dtype: {recv_x_scale.dtype}, output_tensor_scale.dtype: {output_tensor_scale.dtype}" + ) assert ( recv_x_scale.shape[1] == output_tensor_scale.shape[1] == scale_hidden_size ) diff --git a/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py b/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py index 7fa26c7b3..e90e0fb08 100644 --- a/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py +++ b/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py @@ -866,9 +866,9 @@ def invoke_fused_moe_kernel( assert B_scale is not None if block_shape is None: # activation channel-wise int8 quantization - assert ( - per_channel_quant - ), "int8 quantization only supports channel-wise quantization except for block-wise quantization" + assert per_channel_quant, ( + "int8 quantization only supports channel-wise quantization except for block-wise quantization" + ) A, A_scale = per_token_quant_int8(A) else: # activation block-wise int8 quantization @@ -902,23 +902,23 @@ def invoke_fused_moe_kernel( if fuse_sum_all_reduce: assert not c_sorted, "fuse_sum_all_reduce only supports c_sorted=False" if fuse_add_to_output: - assert ( - not fuse_sum_all_reduce - ), "fuse_add_to_output and fuse_sum_all_reduce are mutually exclusive" - assert ( - add_output_mask is not None - ), "add_output_mask required when fuse_add_to_output=True" + assert not fuse_sum_all_reduce, ( + "fuse_add_to_output and fuse_sum_all_reduce are mutually exclusive" + ) + assert add_output_mask is not None, ( + "add_output_mask required when fuse_add_to_output=True" + ) # ===== TO BE REFACTORED ==== if mask_output: - assert ( - not fuse_add_to_output - ), "mask_output and fuse_add_to_output are mutually exclusive" - assert ( - not fuse_sum_all_reduce - ), "mask_output and fuse_sum_all_reduce are mutually exclusive" - assert ( - add_output_mask is not None - ), "add_output_mask required when mask_output=True" + assert not fuse_add_to_output, ( + "mask_output and fuse_add_to_output are mutually exclusive" + ) + assert not fuse_sum_all_reduce, ( + "mask_output and fuse_sum_all_reduce are mutually exclusive" + ) + assert add_output_mask is not None, ( + "add_output_mask required when mask_output=True" + ) # ===== END TO BE REFACTORED ==== if ( @@ -926,9 +926,9 @@ def invoke_fused_moe_kernel( and block_shape is not None and block_shape[1] > 0 ): - assert ( - not fuse_sum_all_reduce - ), "fuse_sum_all_reduce is not supported for GPTQ/AWQ kernels" + assert not fuse_sum_all_reduce, ( + "fuse_sum_all_reduce is not supported for GPTQ/AWQ kernels" + ) assert B_scale is not None and B_scale.ndim == 3 assert B_zp is None or B_zp.ndim == 3 assert bias is None @@ -1559,9 +1559,9 @@ def fused_append_shared_experts_with_weights( ``apply_sigmoid`` (the sigmoid is intrinsic), so the two are mutually exclusive. """ - assert not ( - fuse_gate and apply_sigmoid - ), "fuse_gate already applies sigmoid in-kernel; do not also set apply_sigmoid" + assert not (fuse_gate and apply_sigmoid), ( + "fuse_gate already applies sigmoid in-kernel; do not also set apply_sigmoid" + ) assert N is not None, "N (shared expert base id) must be provided" m, k = topk_ids.shape s = int(num_fused_shared_experts) @@ -1569,9 +1569,9 @@ def fused_append_shared_experts_with_weights( return topk_ids, topk_weights if fuse_gate: - assert ( - hidden_states is not None and gate_weight is not None - ), "fuse_gate=True requires hidden_states and gate_weight" + assert hidden_states is not None and gate_weight is not None, ( + "fuse_gate=True requires hidden_states and gate_weight" + ) hidden_arg = hidden_states.contiguous() wgate_arg = gate_weight.reshape(-1).contiguous() hidden_dim = hidden_arg.shape[1] diff --git a/python/sglang/kernels/ops/moe/inkling_gate_topk_renorm.py b/python/sglang/kernels/ops/moe/inkling_gate_topk_renorm.py index 5601a2859..452275c85 100644 --- a/python/sglang/kernels/ops/moe/inkling_gate_topk_renorm.py +++ b/python/sglang/kernels/ops/moe/inkling_gate_topk_renorm.py @@ -244,9 +244,9 @@ def _get_fused_scratch(device: torch.device) -> tuple[torch.Tensor, torch.Tensor # buffers in the capture pool, where other graphs' replays can reuse # (clobber) them. Call ensure_gate_gemv_fused_scratch() eagerly first # (InklingGate.__init__ does). - assert ( - not torch.cuda.is_current_stream_capturing() - ), "fused gate scratch must be allocated before CUDA graph capture" + assert not torch.cuda.is_current_stream_capturing(), ( + "fused gate scratch must be allocated before CUDA graph capture" + ) workspace = torch.empty( (_FUSED_MAX_TOKENS, _LOGITS_PAD), dtype=torch.float32, device=device ) diff --git a/python/sglang/kernels/ops/moe/inkling_moe.py b/python/sglang/kernels/ops/moe/inkling_moe.py index 04496548a..8769c917d 100644 --- a/python/sglang/kernels/ops/moe/inkling_moe.py +++ b/python/sglang/kernels/ops/moe/inkling_moe.py @@ -331,14 +331,14 @@ def silu_and_mul_triton( Adapted from ``inkling_kernels.activation.silu_and_mul_fwd`` (without MXFP). """ - assert ( - gateup_output.is_contiguous() - ), f"{gateup_output.shape=} {gateup_output.stride()=}" + assert gateup_output.is_contiguous(), ( + f"{gateup_output.shape=} {gateup_output.stride()=}" + ) assert gateup_output.ndim == 2, f"{gateup_output.shape=}" if topk_weights is not None: - assert ( - topk_weights.is_contiguous() - ), f"{topk_weights.shape=} {topk_weights.stride()=}" + assert topk_weights.is_contiguous(), ( + f"{topk_weights.shape=} {topk_weights.stride()=}" + ) assert topk_weights.ndim == 1, f"{topk_weights.shape=}" M = gateup_output.shape[0] @@ -502,9 +502,9 @@ def compute_expert_block_metadata( block_size_m: int = BLOCK_SIZE_M, ): assert num_tokens_per_expert.ndim == 1, f"{num_tokens_per_expert.shape=}" - assert ( - num_tokens_per_expert.is_contiguous() - ), f"{num_tokens_per_expert.shape=} {num_tokens_per_expert.stride()=}" + assert num_tokens_per_expert.is_contiguous(), ( + f"{num_tokens_per_expert.shape=} {num_tokens_per_expert.stride()=}" + ) num_experts = num_tokens_per_expert.numel() max_num_blocks = _get_max_num_blocks(num_routed_tokens, [block_size_m], num_experts) diff --git a/python/sglang/kernels/ops/moe/moe_fused_gate.py b/python/sglang/kernels/ops/moe/moe_fused_gate.py index 640a6610b..b60ecce2e 100644 --- a/python/sglang/kernels/ops/moe/moe_fused_gate.py +++ b/python/sglang/kernels/ops/moe/moe_fused_gate.py @@ -52,9 +52,9 @@ def moe_fused_gate_jit( apply_routed_scaling_factor_on_output: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor]: scoring_func_int = _SCORING_FUNC_MAP.get(scoring_func.lower()) - assert ( - scoring_func_int is not None - ), f"Unknown scoring_func '{scoring_func}', must be one of {list(_SCORING_FUNC_MAP.keys())}" + assert scoring_func_int is not None, ( + f"Unknown scoring_func '{scoring_func}', must be one of {list(_SCORING_FUNC_MAP.keys())}" + ) assert input.dtype == torch.float32, "input must be float32" assert bias.dtype == torch.float32, "bias must be float32" @@ -214,9 +214,7 @@ def _router_triton_kernel( win_lane = tl.min(lane_id, axis=1)[:, None].to(tl.int32) # [BLOCK_M, 1] win_activated = tl.sum( tl.where(offs_n[None, :] == win_lane, activated, 0.0), axis=1 - )[ - :, None - ] # [BLOCK_M, 1] + )[:, None] # [BLOCK_M, 1] slot = offs_k[None, :] == k # [1, BLOCK_K] selected_vals = tl.where(slot, win_activated, selected_vals) selected_idx = tl.where(slot, win_lane, selected_idx) @@ -278,9 +276,9 @@ def moe_fused_gate( the existing call sites. """ scoring_func_int = _SCORING_FUNC_MAP.get(scoring_func.lower()) - assert ( - scoring_func_int is not None - ), f"Unknown scoring_func '{scoring_func}', must be one of {list(_SCORING_FUNC_MAP.keys())}" + assert scoring_func_int is not None, ( + f"Unknown scoring_func '{scoring_func}', must be one of {list(_SCORING_FUNC_MAP.keys())}" + ) assert scores.dtype in ( torch.float32, torch.float16, @@ -288,9 +286,9 @@ def moe_fused_gate( ), "scores must be float32/float16/bfloat16" assert scores.ndim == 2, "scores must be 2D" if bias is None: - assert ( - scoring_func.lower() == "softmax" - ), "bias is required for non-softmax routing" + assert scoring_func.lower() == "softmax", ( + "bias is required for non-softmax routing" + ) else: # The kernel loads the bias and upcasts it to fp32 in-register (see # _router_triton_kernel), so a non-fp32 bias (DeepSeek-V4 stores the @@ -301,9 +299,9 @@ def moe_fused_gate( torch.bfloat16, ), "bias must be float32/float16/bfloat16" assert bias.ndim == 1, "bias must be 1D" - assert scores.size(1) == bias.size( - 0 - ), "scores and bias must have same num_experts" + assert scores.size(1) == bias.size(0), ( + "scores and bias must have same num_experts" + ) assert topk > num_fused_shared_experts, "topk must be > num_fused_shared_experts" if routed_scaling_factor is None: routed_scaling_factor = 1.0 diff --git a/python/sglang/kernels/ops/moe/pack_topk_ids.py b/python/sglang/kernels/ops/moe/pack_topk_ids.py index 3c548c952..437aef580 100644 --- a/python/sglang/kernels/ops/moe/pack_topk_ids.py +++ b/python/sglang/kernels/ops/moe/pack_topk_ids.py @@ -15,7 +15,6 @@ from sglang.kernels.jit.utils import is_arch_support_pdl class PackTopkIds: - @classmethod def execute( cls, topk_ids: torch.Tensor, topk_weights: torch.Tensor @@ -33,18 +32,18 @@ class PackTopkIds: @classmethod def triton(cls, topk_ids: torch.Tensor, topk_weights: torch.Tensor) -> torch.Tensor: - assert ( - topk_ids.shape == topk_weights.shape - ), f"shape mismatch: {topk_ids.shape=} vs {topk_weights.shape=}" + assert topk_ids.shape == topk_weights.shape, ( + f"shape mismatch: {topk_ids.shape=} vs {topk_weights.shape=}" + ) assert topk_ids.ndim >= 1, f"expected >=1D, got {topk_ids.shape=}" assert topk_ids.dtype in ( torch.int32, torch.int64, ), f"topk_ids must be int32 or int64, got {topk_ids.dtype}" - assert ( - topk_weights.dtype == torch.float32 - ), f"topk_weights must be float32, got {topk_weights.dtype}" + assert topk_weights.dtype == torch.float32, ( + f"topk_weights must be float32, got {topk_weights.dtype}" + ) assert topk_ids.is_contiguous(), "topk_ids must be contiguous" assert topk_weights.is_contiguous(), "topk_weights must be contiguous" diff --git a/python/sglang/kernels/ops/moe/rocm_moe_utils.py b/python/sglang/kernels/ops/moe/rocm_moe_utils.py index 43382610f..92f60e1cc 100644 --- a/python/sglang/kernels/ops/moe/rocm_moe_utils.py +++ b/python/sglang/kernels/ops/moe/rocm_moe_utils.py @@ -92,11 +92,11 @@ def rocm_fused_experts_tkw1( # AITER tkw1 kernel for FP8 models with `apply_router_weight_on_input` # This applies topk_weights on the GEMM output of the first FC layer # rather than the second FC. - assert ( - topk_weights.dim() == 2 - ), "`topk_weights` should be in shape (num_tokens, topk)" + assert topk_weights.dim() == 2, ( + "`topk_weights` should be in shape (num_tokens, topk)" + ) assert topk_weights.shape[-1] == 1, ( - "Only support topk=1 when" " `apply_router_weight_on_input` is True" + "Only support topk=1 when `apply_router_weight_on_input` is True" ) return rocm_aiter_asm_moe_tkw1( @@ -307,7 +307,9 @@ def upscale_mxfp4(hidden_state, hidden_state_scale, recv_token_num, output_dtype OUT_TL = ( tl.float16 if output_dtype == torch.float16 - else tl.bfloat16 if output_dtype == torch.bfloat16 else tl.float32 + else tl.bfloat16 + if output_dtype == torch.bfloat16 + else tl.float32 ) upscale_fp4x2_block32_kernel[grid]( diff --git a/python/sglang/kernels/ops/moe/shuffle_rows_with_scales.py b/python/sglang/kernels/ops/moe/shuffle_rows_with_scales.py index 5bc616699..fb5abc7e0 100644 --- a/python/sglang/kernels/ops/moe/shuffle_rows_with_scales.py +++ b/python/sglang/kernels/ops/moe/shuffle_rows_with_scales.py @@ -95,12 +95,12 @@ def shuffle_rows_with_scales( assert q.dim() == 2 and scale.dim() == 2, "q and scale must be 2D" assert q.is_contiguous() and scale.is_contiguous(), "q and scale must be contiguous" assert q.element_size() == 1, f"q must be a 1-byte dtype, got {q.dtype}" - assert ( - q.shape[0] == scale.shape[0] - ), f"row count mismatch: q {q.shape[0]} vs scale {scale.shape[0]}" - assert ( - dst2src_map.numel() >= num_dst_rows - ), f"map holds {dst2src_map.numel()} rows, need {num_dst_rows}" + assert q.shape[0] == scale.shape[0], ( + f"row count mismatch: q {q.shape[0]} vs scale {scale.shape[0]}" + ) + assert dst2src_map.numel() >= num_dst_rows, ( + f"map holds {dst2src_map.numel()} rows, need {num_dst_rows}" + ) # The kernel reads the map as whatever dtype it carries and casts to int64, # so a float map would truncate into a plausible-looking row id instead of # failing. diff --git a/python/sglang/kernels/ops/moe/sigmoid_gate_topk_renorm.py b/python/sglang/kernels/ops/moe/sigmoid_gate_topk_renorm.py index 9221086e9..0f55f5047 100644 --- a/python/sglang/kernels/ops/moe/sigmoid_gate_topk_renorm.py +++ b/python/sglang/kernels/ops/moe/sigmoid_gate_topk_renorm.py @@ -164,16 +164,16 @@ def sigmoid_gate_topk_renorm( # Only column-stride-1 is required (the kernel reads rows via stride_lm). In # InklingGate the gate logits are a [t,258] slice of a padded [t,264] tensor, so # they are NOT contiguous but are column-contiguous -- no copy needed. - assert ( - logits.ndim == 2 and logits.stride(1) == 1 - ), f"{logits.shape=} {logits.stride()=}" - assert ( - logits.shape[0] * logits.stride(0) <= 2**31 - ), f"assumes int32 indexing: {logits.stride()=}" + assert logits.ndim == 2 and logits.stride(1) == 1, ( + f"{logits.shape=} {logits.stride()=}" + ) + assert logits.shape[0] * logits.stride(0) <= 2**31, ( + f"assumes int32 indexing: {logits.stride()=}" + ) assert k <= 32, f"topk kernels only support k <= 32: {k=}" - assert ( - n_shared_experts >= 0 - ), f"expected non-negative shared experts: {n_shared_experts=}" + assert n_shared_experts >= 0, ( + f"expected non-negative shared experts: {n_shared_experts=}" + ) M, G = logits.shape N = G - n_shared_experts A = k + n_shared_experts diff --git a/python/sglang/kernels/ops/moe/trtllm_lora_temp/kimi_k2_moe_fused_gate.py b/python/sglang/kernels/ops/moe/trtllm_lora_temp/kimi_k2_moe_fused_gate.py index a78bdd18d..7ad2476f4 100644 --- a/python/sglang/kernels/ops/moe/trtllm_lora_temp/kimi_k2_moe_fused_gate.py +++ b/python/sglang/kernels/ops/moe/trtllm_lora_temp/kimi_k2_moe_fused_gate.py @@ -39,12 +39,12 @@ def kimi_k2_moe_fused_gate( (output_weights, expert_indices). """ _supported = (torch.float32, torch.bfloat16, torch.float16) - assert ( - input.dtype in _supported - ), f"input must be float32/bfloat16/float16, got {input.dtype}" - assert ( - bias.dtype in _supported - ), f"bias must be float32/bfloat16/float16, got {bias.dtype}" + assert input.dtype in _supported, ( + f"input must be float32/bfloat16/float16, got {input.dtype}" + ) + assert bias.dtype in _supported, ( + f"bias must be float32/bfloat16/float16, got {bias.dtype}" + ) assert input.ndim == 2, "input must be 2D" assert bias.ndim == 1, "bias must be 1D" assert input.size(1) == bias.size(0), "input and bias must have same num_experts" diff --git a/python/sglang/kernels/ops/quantization/fp8_kernel.py b/python/sglang/kernels/ops/quantization/fp8_kernel.py index 6c8c805f0..5cfe48e09 100644 --- a/python/sglang/kernels/ops/quantization/fp8_kernel.py +++ b/python/sglang/kernels/ops/quantization/fp8_kernel.py @@ -269,9 +269,9 @@ def _per_token_group_quant_8bit_raw( Returns: Tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the scaling factor for quantization. """ - assert ( - x.shape[-1] % group_size == 0 - ), "the last dimension of `x` cannot be divisible by `group_size`" + assert x.shape[-1] % group_size == 0, ( + "the last dimension of `x` cannot be divisible by `group_size`" + ) assert x.is_contiguous(), "`x` is not contiguous" if _is_hip: @@ -585,9 +585,9 @@ def _run_per_token_group_quant_8bit_kernel( ) return - assert ( - eps == 1e-10 - ), f"per_token_group_quant bakes the absmax floor in at 1e-10, got {eps}" + assert eps == 1e-10, ( + f"per_token_group_quant bakes the absmax floor in at 1e-10, got {eps}" + ) expected_range = (-448.0, 448.0) if x_q.dtype == fp8_dtype else (-128.0, 127.0) assert (fp8_min, fp8_max) == expected_range, ( f"per_token_group_quant bakes the {x_q.dtype} quant range in at {expected_range}, " @@ -614,9 +614,9 @@ def sglang_per_token_group_quant_fp8( fuse_silu_and_mul: bool = False, masked_m: Optional[torch.Tensor] = None, ): - assert ( - x.shape[-1] % group_size == 0 - ), "the last dimension of `x` cannot be divisible by `group_size`" + assert x.shape[-1] % group_size == 0, ( + "the last dimension of `x` cannot be divisible by `group_size`" + ) assert x.is_contiguous(), "`x` is not contiguous" if ( @@ -676,9 +676,9 @@ def sglang_per_token_group_quant_fp8_row_padded( bit-exact; the caller still slices the GEMM output back to m. """ assert x.dim() == 2, "row-padded quant expects a 2D input" - assert ( - x.shape[-1] % group_size == 0 - ), "the last dimension of `x` must be divisible by `group_size`" + assert x.shape[-1] % group_size == 0, ( + "the last dimension of `x` must be divisible by `group_size`" + ) assert x.is_contiguous(), "`x` is not contiguous" supported_group_sizes = ( @@ -726,9 +726,9 @@ def sglang_per_token_group_quant_fp8_ue8m0( group_size: int, eps: float = 1e-10, ) -> Tuple[torch.Tensor, torch.Tensor]: - assert ( - x.shape[-1] % group_size == 0 - ), f"hidden ({x.shape[-1]}) must be divisible by group_size ({group_size})" + assert x.shape[-1] % group_size == 0, ( + f"hidden ({x.shape[-1]}) must be divisible by group_size ({group_size})" + ) assert x.is_contiguous(), "x must be contiguous" *x_batch, x_q_mn, x_q_k = x.shape @@ -1466,9 +1466,9 @@ def prepare_block_fp8_matmul_inputs( if As.dtype == torch.float: assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1] elif As.dtype == torch.int: - assert ( - triton.cdiv(triton.cdiv(A.shape[-1], block_k), 4) == As.shape[-1] - ), f"{A.shape=} {As.shape=} {block_size=}" + assert triton.cdiv(triton.cdiv(A.shape[-1], block_k), 4) == As.shape[-1], ( + f"{A.shape=} {As.shape=} {block_size=}" + ) else: raise NotImplementedError @@ -1484,9 +1484,9 @@ def prepare_block_fp8_matmul_inputs( assert triton.cdiv(K, block_k) == Bs.shape[1] elif Bs.dtype == torch.int: assert N == Bs.shape[0], f"{B.shape=} {Bs.shape=} {block_size=}" - assert ( - triton.cdiv(triton.cdiv(K, block_k), 4) == Bs.shape[1] - ), f"{B.shape=} {Bs.shape=} {block_size=}" + assert triton.cdiv(triton.cdiv(K, block_k), 4) == Bs.shape[1], ( + f"{B.shape=} {Bs.shape=} {block_size=}" + ) else: raise NotImplementedError @@ -1932,9 +1932,9 @@ if _is_hip: _native_dynamic_per_tensor_quant_fp8(output, input, scale) else: # Static scaling - assert ( - scale.numel() == 1 - ), f"Expected scalar scale, got numel={scale.numel()}" + assert scale.numel() == 1, ( + f"Expected scalar scale, got numel={scale.numel()}" + ) if _use_aiter: static_per_tensor_quant(output, input, scale) elif _has_vllm: @@ -1973,9 +1973,9 @@ else: ) # False for dynamic else: # Static scaling - assert ( - scale.numel() == 1 - ), f"Expected scalar scale, got numel={scale.numel()}" + assert scale.numel() == 1, ( + f"Expected scalar scale, got numel={scale.numel()}" + ) sgl_per_tensor_quant_fp8( input, output, scale, is_static=True ) # True for static @@ -2054,9 +2054,9 @@ def per_token_group_quant_fp8_hopper_moe_mn_major( ) -> Tuple[torch.Tensor, torch.Tensor]: assert A.dim() == 2 assert A.is_contiguous(), "`A` is not contiguous" - assert ( - A.shape[-1] % group_size == 0 - ), "the last dimension of `A` cannot be divisible by `group_size`" + assert A.shape[-1] % group_size == 0, ( + "the last dimension of `A` cannot be divisible by `group_size`" + ) a_q = torch.empty_like(A, device=A.device, dtype=fp8_dtype) M, K = A.shape[0], A.shape[1] diff --git a/python/sglang/kernels/ops/quantization/fp8_quantize.py b/python/sglang/kernels/ops/quantization/fp8_quantize.py index 6eadb78a0..53c4c6636 100644 --- a/python/sglang/kernels/ops/quantization/fp8_quantize.py +++ b/python/sglang/kernels/ops/quantization/fp8_quantize.py @@ -81,7 +81,7 @@ def _flatten_to_2d(x: torch.Tensor): if x.stride(d) != expected: raise ValueError( f"cannot flatten dim {d}: stride={x.stride(d)} but expected " - f"shape[{d+1}]*stride[{d+1}]={expected}. Tensor shape={tuple(x.shape)}, " + f"shape[{d + 1}]*stride[{d + 1}]={expected}. Tensor shape={tuple(x.shape)}, " f"stride={tuple(x.stride())}." ) return M, N, row_stride diff --git a/python/sglang/kernels/ops/quantization/int8_kernel.py b/python/sglang/kernels/ops/quantization/int8_kernel.py index 3c73744f4..2748db0d8 100644 --- a/python/sglang/kernels/ops/quantization/int8_kernel.py +++ b/python/sglang/kernels/ops/quantization/int8_kernel.py @@ -152,9 +152,9 @@ def per_token_group_quant_int8( Returns: Tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the scaling factor for quantization. """ - assert ( - x.shape[-1] % group_size == 0 - ), "the last dimension of `x` cannot be divisible by `group_size`" + assert x.shape[-1] % group_size == 0, ( + "the last dimension of `x` cannot be divisible by `group_size`" + ) assert x.is_contiguous(), "`x` is not contiguous" iinfo = torch.iinfo(dtype) @@ -197,15 +197,15 @@ def sglang_per_token_group_quant_int8( eps: float = 1e-10, dtype: torch.dtype = torch.int8, ): - assert ( - x.shape[-1] % group_size == 0 - ), "the last dimension of `x` cannot be divisible by `group_size`" + assert x.shape[-1] % group_size == 0, ( + "the last dimension of `x` cannot be divisible by `group_size`" + ) assert x.is_contiguous(), "`x` is not contiguous" assert dtype == torch.int8 # per_token_group_quant bakes the int8 constants in ([-128, 127], eps 1e-10). - assert ( - eps == 1e-10 - ), f"per_token_group_quant bakes the absmax floor in at 1e-10, got {eps}" + assert eps == 1e-10, ( + f"per_token_group_quant bakes the absmax floor in at 1e-10, got {eps}" + ) return per_token_group_quant(x, group_size=group_size, out_dtype=dtype) diff --git a/python/sglang/kernels/ops/quantization/mxfp8_amd_gfx95.py b/python/sglang/kernels/ops/quantization/mxfp8_amd_gfx95.py index 9c90ce8ed..773ab1511 100644 --- a/python/sglang/kernels/ops/quantization/mxfp8_amd_gfx95.py +++ b/python/sglang/kernels/ops/quantization/mxfp8_amd_gfx95.py @@ -309,9 +309,9 @@ def dot_scaled_mxfp8_blockscaled_linear( kernel_out_dtype = input_2d.dtype else: # Activations already MXFP8-quantized by a fused upstream op. - assert ( - input_2d.dtype == MXFP8_VALUE_DTYPE - ), "pre-quantized input must be FP8 E4M3 when input_scale is given." + assert input_2d.dtype == MXFP8_VALUE_DTYPE, ( + "pre-quantized input must be FP8 E4M3 when input_scale is given." + ) assert input_scale.dtype == torch.uint8 and input_scale.shape == ( m, k // 32, diff --git a/python/sglang/kernels/ops/quantization/mxfp8_interleave_sf.py b/python/sglang/kernels/ops/quantization/mxfp8_interleave_sf.py index ede6ce7a2..8a9ce307d 100644 --- a/python/sglang/kernels/ops/quantization/mxfp8_interleave_sf.py +++ b/python/sglang/kernels/ops/quantization/mxfp8_interleave_sf.py @@ -69,9 +69,9 @@ def store_sf_interleaved( page_size: int = 128, ): """Scatter-write per-token scale factors into interleaved page layout.""" - assert ( - page_size == 128 - ), f"Interleaved SF layout requires page_size=128, got {page_size}" + assert page_size == 128, ( + f"Interleaved SF layout requires page_size=128, got {page_size}" + ) num_tokens, nheads, sf_dim = sf_in.shape assert sf_dim == 4, f"Expected sf_dim=4 (hdim=128, sf_vec_size=32), got {sf_dim}" diff --git a/python/sglang/kernels/ops/sampling/murmur_hash.py b/python/sglang/kernels/ops/sampling/murmur_hash.py index 2a090a4ca..3028136b8 100644 --- a/python/sglang/kernels/ops/sampling/murmur_hash.py +++ b/python/sglang/kernels/ops/sampling/murmur_hash.py @@ -102,12 +102,12 @@ def murmur_hash32_kernel( def murmur_hash32(seed, positions, col_indices): - assert ( - seed.shape == positions.shape - ), "Seed and positions must have the same shape (n,)" - assert ( - len(seed.shape) == 1 and len(col_indices.shape) == 1 - ), f"Inputs must be 1D tensors {seed.shape=} {col_indices.shape=}" + assert seed.shape == positions.shape, ( + "Seed and positions must have the same shape (n,)" + ) + assert len(seed.shape) == 1 and len(col_indices.shape) == 1, ( + f"Inputs must be 1D tensors {seed.shape=} {col_indices.shape=}" + ) n = seed.shape[0] m = col_indices.shape[0] device = seed.device diff --git a/python/sglang/kernels/ops/speculative/dspark/dspark_accept.py b/python/sglang/kernels/ops/speculative/dspark/dspark_accept.py index 948a9d60a..a63748c0d 100644 --- a/python/sglang/kernels/ops/speculative/dspark/dspark_accept.py +++ b/python/sglang/kernels/ops/speculative/dspark/dspark_accept.py @@ -306,9 +306,9 @@ def softmax_temp( ) -> torch.Tensor: num_rows = logits.shape[0] bs = num_rows // rows_per_request - assert ( - bs * rows_per_request == num_rows - ), f"num_rows {num_rows} not divisible by rows_per_request {rows_per_request}" + assert bs * rows_per_request == num_rows, ( + f"num_rows {num_rows} not divisible by rows_per_request {rows_per_request}" + ) temp_per_row = torch.repeat_interleave( temperatures.reshape(bs).to(torch.float32), rows_per_request, dim=0 ) @@ -366,9 +366,9 @@ def softmax_temp_triton( ) -> torch.Tensor: num_rows, vocab = logits.shape[0], logits.shape[-1] bs = num_rows // rows_per_request - assert ( - bs * rows_per_request == num_rows - ), f"num_rows {num_rows} not divisible by rows_per_request {rows_per_request}" + assert bs * rows_per_request == num_rows, ( + f"num_rows {num_rows} not divisible by rows_per_request {rows_per_request}" + ) temperatures = temperatures.reshape(bs).to(torch.float32).contiguous() out = torch.empty((num_rows, vocab), dtype=torch.float32, device=logits.device) BLOCK_V = 4096 @@ -397,9 +397,9 @@ def softmax_temp_flashinfer( ) num_rows, vocab = logits.shape[0], logits.shape[-1] bs = num_rows // rows_per_request - assert ( - bs * rows_per_request == num_rows - ), f"num_rows {num_rows} not divisible by rows_per_request {rows_per_request}" + assert bs * rows_per_request == num_rows, ( + f"num_rows {num_rows} not divisible by rows_per_request {rows_per_request}" + ) temp_per_row = torch.repeat_interleave( temperatures.reshape(bs).to(torch.float32), rows_per_request, dim=0 ).contiguous() diff --git a/python/sglang/kernels/ops/speculative/dspark/dspark_verify_window.py b/python/sglang/kernels/ops/speculative/dspark/dspark_verify_window.py index 0de315543..f733ba2bd 100644 --- a/python/sglang/kernels/ops/speculative/dspark/dspark_verify_window.py +++ b/python/sglang/kernels/ops/speculative/dspark/dspark_verify_window.py @@ -390,9 +390,9 @@ def compact_row_index_triton( verify_lens = verify_lens.to(device=device, dtype=torch.int64).contiguous() bs = verify_lens.shape[0] # The search converges only for bs <= 2**(NBITS-1); beyond it silently mismaps. - assert bs <= 1 << ( - _SEARCH_NBITS - 1 - ), f"bs={bs} exceeds row-index search capacity {1 << (_SEARCH_NBITS - 1)}" + assert bs <= 1 << (_SEARCH_NBITS - 1), ( + f"bs={bs} exceeds row-index search capacity {1 << (_SEARCH_NBITS - 1)}" + ) incl = torch.cumsum(verify_lens, dim=0).contiguous() req = torch.empty(padded_total, dtype=torch.int64, device=device) within = torch.empty(padded_total, dtype=torch.int64, device=device) diff --git a/python/sglang/lang/backend/openai.py b/python/sglang/lang/backend/openai.py index a2d006bb7..3ee2f4fb1 100644 --- a/python/sglang/lang/backend/openai.py +++ b/python/sglang/lang/backend/openai.py @@ -129,9 +129,9 @@ class OpenAI(BaseBackend): if key not in self.spec_kwargs: self.spec_kwargs[key] = value else: - assert ( - value == self.spec_kwargs[key] - ), "sampling parameters should be consistent if turn on api speculative execution." + assert value == self.spec_kwargs[key], ( + "sampling parameters should be consistent if turn on api speculative execution." + ) self.spec_format.append( {"text": "", "stop": params["stop"], "name": spec_var_name} ) @@ -180,9 +180,9 @@ class OpenAI(BaseBackend): ) # Keep the returned list (or string) as is. elif sampling_params.dtype in [str, "str", "string"]: - assert ( - not self.is_chat_model - ), "constrained type not supported on chat model" + assert not self.is_chat_model, ( + "constrained type not supported on chat model" + ) kwargs = sampling_params.to_openai_kwargs() kwargs.pop("stop") comp = openai_completion( @@ -200,9 +200,9 @@ class OpenAI(BaseBackend): else: comp = '"' + comp + '"' elif sampling_params.dtype in [int, "int"]: - assert ( - not self.is_chat_model - ), "constrained type not supported on chat model" + assert not self.is_chat_model, ( + "constrained type not supported on chat model" + ) kwargs = sampling_params.to_openai_kwargs() kwargs.pop("stop") comp = openai_completion( diff --git a/python/sglang/lang/backend/runtime_endpoint.py b/python/sglang/lang/backend/runtime_endpoint.py index db61e431f..84849d4ea 100644 --- a/python/sglang/lang/backend/runtime_endpoint.py +++ b/python/sglang/lang/backend/runtime_endpoint.py @@ -133,18 +133,14 @@ class RuntimeEndpoint(BaseBackend): dtype_regex = None if sampling_params.dtype in ["int", int]: - dtype_regex = REGEX_INT sampling_params.stop.extend([" ", "\n"]) elif sampling_params.dtype in ["float", float]: - dtype_regex = REGEX_FLOAT sampling_params.stop.extend([" ", "\n"]) elif sampling_params.dtype in ["str", str]: - dtype_regex = REGEX_STR elif sampling_params.dtype in ["bool", bool]: - dtype_regex = REGEX_BOOL else: raise RuntimeError(f"Invalid dtype: {sampling_params.dtype}") diff --git a/python/sglang/lang/choices.py b/python/sglang/lang/choices.py index e52c6b362..ceac01b69 100644 --- a/python/sglang/lang/choices.py +++ b/python/sglang/lang/choices.py @@ -12,7 +12,6 @@ class ChoicesDecision: class ChoicesSamplingMethod(ABC): - @property def requires_unconditional_logprobs(self) -> bool: return False @@ -30,7 +29,6 @@ class ChoicesSamplingMethod(ABC): class TokenLengthNormalized(ChoicesSamplingMethod): - def __call__( self, *, @@ -54,7 +52,6 @@ token_length_normalized = TokenLengthNormalized() class GreedyTokenSelection(ChoicesSamplingMethod): - def __call__( self, *, @@ -108,7 +105,6 @@ greedy_token_selection = GreedyTokenSelection() class UnconditionalLikelihoodNormalized(ChoicesSamplingMethod): - @property def requires_unconditional_logprobs(self) -> bool: return True diff --git a/python/sglang/lang/interpreter.py b/python/sglang/lang/interpreter.py index 90dd41857..c07a84bd2 100644 --- a/python/sglang/lang/interpreter.py +++ b/python/sglang/lang/interpreter.py @@ -624,9 +624,9 @@ class StreamExecutor: self.meta_info[name] = meta_info self.variable_event[name].set() else: - assert ( - self.num_api_spec_tokens is None - ), "stream is not supported with api speculative execution" + assert self.num_api_spec_tokens is None, ( + "stream is not supported with api speculative execution" + ) generator = self.backend.generate_stream( self, sampling_params=sampling_params ) diff --git a/python/sglang/lang/ir.py b/python/sglang/lang/ir.py index 45cb6c859..771e929a4 100644 --- a/python/sglang/lang/ir.py +++ b/python/sglang/lang/ir.py @@ -531,7 +531,6 @@ class SglRoleEnd(SglExpr): class SglSelect(SglExpr): - def __init__( self, name: str, diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py index 6133d5f97..36527369a 100755 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py @@ -2356,8 +2356,7 @@ def main(): args.breakable_cuda_graph or args.quality_bcg_matrix ): parser.error( - "--bcg-text-buckets requires --breakable-cuda-graph or " - "--quality-bcg-matrix" + "--bcg-text-buckets requires --breakable-cuda-graph or --quality-bcg-matrix" ) if args.cleanup_model_cache and not args.model_cache_root: parser.error("--cleanup-model-cache requires --model-cache-root") diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_flux_pipeline.py b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_flux_pipeline.py index 63c099682..59147a6c3 100644 --- a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_flux_pipeline.py +++ b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_flux_pipeline.py @@ -135,12 +135,12 @@ def test_comfyui_flux_pipeline_direct() -> None: assert noise_pred is not None, "noise_pred should not be None in OutputBatch" assert isinstance(noise_pred, torch.Tensor), "noise_pred should be a torch.Tensor" - assert ( - noise_pred.device.type == "cuda" - ), f"noise_pred should be on cuda, got {noise_pred.device}" - assert ( - noise_pred.dtype == torch.bfloat16 - ), f"noise_pred should be bfloat16, got {noise_pred.dtype}" + assert noise_pred.device.type == "cuda", ( + f"noise_pred should be on cuda, got {noise_pred.device}" + ) + assert noise_pred.dtype == torch.bfloat16, ( + f"noise_pred should be bfloat16, got {noise_pred.dtype}" + ) print("✓ Successfully retrieved noise_pred from OutputBatch!") print(f" noise_pred shape: {noise_pred.shape}") diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_h3_request.py b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_h3_request.py index b46cf3b5e..6b5b8e231 100644 --- a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_h3_request.py +++ b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_h3_request.py @@ -145,15 +145,19 @@ def _run_node(**node_kwargs): ) node = SGLDiffusionGenerateH3() - with mock.patch( - f"{PKG}.core.server_api.requests.post", - side_effect=fake_post, - ), mock.patch( - f"{PKG}.core.server_api.requests.get", - side_effect=fake_get, - ), mock.patch( - f"{PKG}.nodes.get_image_path", - side_effect=lambda image: "/tmp/frame.png", + with ( + mock.patch( + f"{PKG}.core.server_api.requests.post", + side_effect=fake_post, + ), + mock.patch( + f"{PKG}.core.server_api.requests.get", + side_effect=fake_get, + ), + mock.patch( + f"{PKG}.nodes.get_image_path", + side_effect=lambda image: "/tmp/frame.png", + ), ): result = node.generate(sgld_client=client, **node_kwargs) return captured, result @@ -243,13 +247,16 @@ def test_extra_fields_win_over_generic_defaults(): captured.update(json) return _Response({"id": "job-1"}) - with mock.patch( - f"{PKG}.core.server_api.requests.post", - side_effect=fake_post, - ), mock.patch( - f"{PKG}.core.server_api.requests.get", - side_effect=lambda *a, **k: _Response( - {"id": "job-1", "status": "completed", "size": RESOLVED_SIZE} + with ( + mock.patch( + f"{PKG}.core.server_api.requests.post", + side_effect=fake_post, + ), + mock.patch( + f"{PKG}.core.server_api.requests.get", + side_effect=lambda *a, **k: _Response( + {"id": "job-1", "status": "completed", "size": RESOLVED_SIZE} + ), ), ): client.generate_video( diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_qwen_image_edit_pipeline.py b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_qwen_image_edit_pipeline.py index 609f8df04..422441d56 100644 --- a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_qwen_image_edit_pipeline.py +++ b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_qwen_image_edit_pipeline.py @@ -116,12 +116,12 @@ def test_comfyui_qwen_image_edit_pipeline_direct() -> None: assert noise_pred is not None, "noise_pred should not be None in OutputBatch" assert isinstance(noise_pred, torch.Tensor), "noise_pred should be a torch.Tensor" - assert ( - noise_pred.device.type == "cuda" - ), f"noise_pred should be on cuda, got {noise_pred.device}" - assert ( - noise_pred.dtype == torch.bfloat16 - ), f"noise_pred should be bfloat16, got {noise_pred.dtype}" + assert noise_pred.device.type == "cuda", ( + f"noise_pred should be on cuda, got {noise_pred.device}" + ) + assert noise_pred.dtype == torch.bfloat16, ( + f"noise_pred should be bfloat16, got {noise_pred.dtype}" + ) print("✓ Successfully retrieved noise_pred from OutputBatch (Edit Mode)!") print(f" noise_pred shape: {noise_pred.shape}") diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_qwen_image_pipeline.py b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_qwen_image_pipeline.py index 88d73f663..bb7b070f0 100644 --- a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_qwen_image_pipeline.py +++ b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_qwen_image_pipeline.py @@ -100,12 +100,12 @@ def test_comfyui_qwen_image_pipeline_direct() -> None: assert noise_pred is not None, "noise_pred should not be None in OutputBatch" assert isinstance(noise_pred, torch.Tensor), "noise_pred should be a torch.Tensor" - assert ( - noise_pred.device.type == "cuda" - ), f"noise_pred should be on cuda, got {noise_pred.device}" - assert ( - noise_pred.dtype == torch.bfloat16 - ), f"noise_pred should be bfloat16, got {noise_pred.dtype}" + assert noise_pred.device.type == "cuda", ( + f"noise_pred should be on cuda, got {noise_pred.device}" + ) + assert noise_pred.dtype == torch.bfloat16, ( + f"noise_pred should be bfloat16, got {noise_pred.dtype}" + ) print("✓ Successfully retrieved noise_pred from OutputBatch!") print(f" noise_pred shape: {noise_pred.shape}") diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_zimage_pipeline.py b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_zimage_pipeline.py index 77e6811bf..4b053f1db 100644 --- a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_zimage_pipeline.py +++ b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_zimage_pipeline.py @@ -102,12 +102,12 @@ def test_comfyui_zimage_pipeline_direct() -> None: assert noise_pred is not None, "noise_pred should not be None in OutputBatch" assert isinstance(noise_pred, torch.Tensor), "noise_pred should be a torch.Tensor" - assert ( - noise_pred.device.type == "cuda" - ), f"noise_pred should be on cuda, got {noise_pred.device}" - assert ( - noise_pred.dtype == torch.bfloat16 - ), f"noise_pred should be bfloat16, got {noise_pred.dtype}" + assert noise_pred.device.type == "cuda", ( + f"noise_pred should be on cuda, got {noise_pred.device}" + ) + assert noise_pred.dtype == torch.bfloat16, ( + f"noise_pred should be bfloat16, got {noise_pred.dtype}" + ) print("✓ Successfully retrieved noise_pred from OutputBatch!") print(f" noise_pred shape: {noise_pred.shape}") diff --git a/python/sglang/multimodal_gen/benchmarks/bench_serving.py b/python/sglang/multimodal_gen/benchmarks/bench_serving.py index dca9a4c28..b5a1dffee 100644 --- a/python/sglang/multimodal_gen/benchmarks/bench_serving.py +++ b/python/sglang/multimodal_gen/benchmarks/bench_serving.py @@ -631,7 +631,7 @@ async def benchmark(args): warm_out = await limited_request_func(warm_req, session, None) warmup_pairs.append((warm_req, warm_out)) logger.info( - f"Warmup {i+1}/{args.warmup_requests}: " + f"Warmup {i + 1}/{args.warmup_requests}: " f"latency={warm_out.latency:.2f}s, success={warm_out.success}" ) diff --git a/python/sglang/multimodal_gen/configs/models/adapter/ltx_2_connector.py b/python/sglang/multimodal_gen/configs/models/adapter/ltx_2_connector.py index 6caefb45e..03b18fe08 100644 --- a/python/sglang/multimodal_gen/configs/models/adapter/ltx_2_connector.py +++ b/python/sglang/multimodal_gen/configs/models/adapter/ltx_2_connector.py @@ -67,7 +67,6 @@ class LTX2ConnectorArchConfig(AdapterArchConfig): @dataclass class LTX2ConnectorConfig(AdapterConfig): - arch_config: AdapterArchConfig = field(default_factory=LTX2ConnectorArchConfig) prefix: str = "LTX2" diff --git a/python/sglang/multimodal_gen/configs/models/base.py b/python/sglang/multimodal_gen/configs/models/base.py index 3bca5cc67..a789b25af 100644 --- a/python/sglang/multimodal_gen/configs/models/base.py +++ b/python/sglang/multimodal_gen/configs/models/base.py @@ -83,9 +83,9 @@ class ModelConfig: arch_config.__post_init__() def update_model_config(self, source_model_dict: dict[str, Any]) -> None: - assert ( - "arch_config" not in source_model_dict - ), "Source model config shouldn't contain arch_config." + assert "arch_config" not in source_model_dict, ( + "Source model config shouldn't contain arch_config." + ) valid_fields = {f.name for f in fields(self)} diff --git a/python/sglang/multimodal_gen/configs/models/dits/flux.py b/python/sglang/multimodal_gen/configs/models/dits/flux.py index 97adc01e8..d1a8a6885 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/flux.py +++ b/python/sglang/multimodal_gen/configs/models/dits/flux.py @@ -110,7 +110,6 @@ class FluxArchConfig(DiTArchConfig): @dataclass class FluxConfig(DiTConfig): - arch_config: DiTArchConfig = field(default_factory=FluxArchConfig) prefix: str = "Flux" diff --git a/python/sglang/multimodal_gen/configs/models/dits/krea2.py b/python/sglang/multimodal_gen/configs/models/dits/krea2.py index 99a83f97e..b330b5364 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/krea2.py +++ b/python/sglang/multimodal_gen/configs/models/dits/krea2.py @@ -57,9 +57,9 @@ class Krea2ArchConfig(DiTArchConfig): self.num_attention_heads = self.heads self.num_channels_latents = self.channels assert self.features % self.heads == 0 - assert ( - sum(self.axes_dims) == self.features // self.heads - ), f"sum(axes_dims)={sum(self.axes_dims)} != head_dim={self.features // self.heads}" + assert sum(self.axes_dims) == self.features // self.heads, ( + f"sum(axes_dims)={sum(self.axes_dims)} != head_dim={self.features // self.heads}" + ) @property def head_dim(self) -> int: diff --git a/python/sglang/multimodal_gen/configs/models/dits/mova_audio.py b/python/sglang/multimodal_gen/configs/models/dits/mova_audio.py index ad2ee5d93..637025a4f 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/mova_audio.py +++ b/python/sglang/multimodal_gen/configs/models/dits/mova_audio.py @@ -50,9 +50,9 @@ class MOVAAudioArchConfig(DiTArchConfig): self.hidden_size = self.dim self.num_attention_heads = self.num_heads self.num_channels_latents = self.out_dim - assert ( - not self.has_image_input - ), "has_image_input must be False; it's a config from Diffsynth Studio, which means the model uses CLIP for image encoding (we don't)." + assert not self.has_image_input, ( + "has_image_input must be False; it's a config from Diffsynth Studio, which means the model uses CLIP for image encoding (we don't)." + ) @dataclass diff --git a/python/sglang/multimodal_gen/configs/models/dits/mova_video.py b/python/sglang/multimodal_gen/configs/models/dits/mova_video.py index b0c2e882f..e3f010e88 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/mova_video.py +++ b/python/sglang/multimodal_gen/configs/models/dits/mova_video.py @@ -49,9 +49,9 @@ class MOVAVideoArchConfig(DiTArchConfig): self.hidden_size = self.dim self.num_attention_heads = self.num_heads self.num_channels_latents = self.out_dim - assert ( - not self.has_image_input - ), "has_image_input must be False; it's a config from Diffsynth Studio, which means the model uses CLIP for image encoding (we don't)." + assert not self.has_image_input, ( + "has_image_input must be False; it's a config from Diffsynth Studio, which means the model uses CLIP for image encoding (we don't)." + ) @dataclass diff --git a/python/sglang/multimodal_gen/configs/models/dits/wanvideo.py b/python/sglang/multimodal_gen/configs/models/dits/wanvideo.py index 85a3acda8..daad209ee 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/wanvideo.py +++ b/python/sglang/multimodal_gen/configs/models/dits/wanvideo.py @@ -100,9 +100,7 @@ class WanVideoArchConfig(DiTArchConfig): local_attn_size: int = ( -1 ) # Window size for temporal local attention (-1 indicates global attention) - sink_size: int = ( - 0 # Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache - ) + sink_size: int = 0 # Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache num_frames_per_block: int = 3 sliding_window_num_frames: int = 21 attention_type: str = "original" diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index a15e1a5cf..e86ad7bde 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -1180,9 +1180,9 @@ class PipelineConfig: elif isinstance(current_value, tuple) and all( isinstance(v, ModelConfig) for v in current_value ): - assert len(current_value) == len( - new_value - ), "Users shouldn't delete or add text encoder config objects in your json" + assert len(current_value) == len(new_value), ( + "Users shouldn't delete or add text encoder config objects in your json" + ) for target_config, source_config in zip( current_value, new_value, strict=True ): diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/longlive2.py b/python/sglang/multimodal_gen/configs/pipeline_configs/longlive2.py index 209a484ab..feccd04a2 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/longlive2.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/longlive2.py @@ -17,7 +17,6 @@ logger = init_logger(__name__) @dataclass class LongLive2T2VConfig(Wan2_2_TI2V_5B_Config): - is_causal: bool = True task_type: ModelTaskType = ModelTaskType.TI2V vae_precision: str = "bf16" diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/sana.py b/python/sglang/multimodal_gen/configs/pipeline_configs/sana.py index e564947d5..c3967b40e 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/sana.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/sana.py @@ -41,7 +41,6 @@ def sana_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Ten @dataclass class SanaPipelineConfig(SpatialImagePipelineConfig): - task_type: ModelTaskType = ModelTaskType.T2I # should_use_guidance=False disables *embedded* guidance (timestep-conditioned diff --git a/python/sglang/multimodal_gen/configs/post_training/pipeline_configs/qwen_image_rollout_pipeline_mixin.py b/python/sglang/multimodal_gen/configs/post_training/pipeline_configs/qwen_image_rollout_pipeline_mixin.py index fa5323f47..a3c1bd62e 100644 --- a/python/sglang/multimodal_gen/configs/post_training/pipeline_configs/qwen_image_rollout_pipeline_mixin.py +++ b/python/sglang/multimodal_gen/configs/post_training/pipeline_configs/qwen_image_rollout_pipeline_mixin.py @@ -12,7 +12,6 @@ from sglang.multimodal_gen.runtime.post_training.sp_utils import ( class QwenImageRolloutPipelineMixin: - def gather_denoising_env_static_for_sp(self, batch, cond_kwargs: dict | None): if cond_kwargs is None: return None diff --git a/python/sglang/multimodal_gen/configs/post_training/pipeline_configs/zimage_rollout_pipeline_mixin.py b/python/sglang/multimodal_gen/configs/post_training/pipeline_configs/zimage_rollout_pipeline_mixin.py index 5c4e60ad5..f49e63ac3 100644 --- a/python/sglang/multimodal_gen/configs/post_training/pipeline_configs/zimage_rollout_pipeline_mixin.py +++ b/python/sglang/multimodal_gen/configs/post_training/pipeline_configs/zimage_rollout_pipeline_mixin.py @@ -12,7 +12,6 @@ from sglang.multimodal_gen.runtime.post_training.sp_utils import ( class ZImageRolloutPipelineMixin: - def gather_denoising_env_static_for_sp(self, batch, cond_kwargs: dict | None): if cond_kwargs is None: return None diff --git a/python/sglang/multimodal_gen/configs/sample/cosmos3.py b/python/sglang/multimodal_gen/configs/sample/cosmos3.py index e75e9e71b..75512d53f 100644 --- a/python/sglang/multimodal_gen/configs/sample/cosmos3.py +++ b/python/sglang/multimodal_gen/configs/sample/cosmos3.py @@ -501,8 +501,7 @@ class Cosmos3SamplingParams(SamplingParams): raise ValueError("num_conditional_frames must be non-negative") if self.num_conditional_frames >= self.num_video_frames_per_chunk: raise ValueError( - "num_conditional_frames must be smaller than " - "num_video_frames_per_chunk" + "num_conditional_frames must be smaller than num_video_frames_per_chunk" ) if self.num_first_chunk_conditional_frames < 0: raise ValueError("num_first_chunk_conditional_frames must be non-negative") diff --git a/python/sglang/multimodal_gen/configs/sample/minimax_h3.py b/python/sglang/multimodal_gen/configs/sample/minimax_h3.py index 85934ec90..4c4d4546b 100644 --- a/python/sglang/multimodal_gen/configs/sample/minimax_h3.py +++ b/python/sglang/multimodal_gen/configs/sample/minimax_h3.py @@ -267,7 +267,9 @@ class MiniMaxH3SamplingParams(SamplingParams): seed=( _seed_override if _seed_override is not None - else self.seed if isinstance(self.seed, int) else None + else self.seed + if isinstance(self.seed, int) + else None ), ) ) diff --git a/python/sglang/multimodal_gen/configs/sample/qwenimage.py b/python/sglang/multimodal_gen/configs/sample/qwenimage.py index 127dfa5fc..8e0b09033 100644 --- a/python/sglang/multimodal_gen/configs/sample/qwenimage.py +++ b/python/sglang/multimodal_gen/configs/sample/qwenimage.py @@ -17,9 +17,7 @@ class QwenImageSamplingParams(SamplingParams): @dataclass class QwenImage2512SamplingParams(QwenImageSamplingParams): - negative_prompt: str = ( - "低分辨率,低画质,肢体畸形,手指畸形,画面过饱和,蜡像感,人脸无细节,过度光滑,画面具有AI感。构图混乱。文字模糊,扭曲。" - ) + negative_prompt: str = "低分辨率,低画质,肢体畸形,手指畸形,画面过饱和,蜡像感,人脸无细节,过度光滑,画面具有AI感。构图混乱。文字模糊,扭曲。" @dataclass diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index d169d7ddc..9dfc3462e 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -133,9 +133,7 @@ class SamplingParams: prompt: str | list[str] | None = field( default=None, metadata={"batch_sig_exclude": True} ) - negative_prompt: str = ( - "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" - ) + negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" prompt_path: str | None = field(default=None, metadata={"batch_sig_exclude": True}) output_path: str | None = field(default=None, metadata={"batch_sig_exclude": True}) output_file_name: str | None = field( @@ -263,12 +261,8 @@ class SamplingParams: ) return_trajectory_latents: bool = False # returns all latents for each timestep return_trajectory_decoded: bool = False # returns decoded latents for each timestep - rollout_return_denoising_env: bool = ( - False # populate ``denoising_env`` (image/pos/neg kwargs, guidance) for RL replay - ) - rollout_return_dit_trajectory: bool = ( - False # per-step noisy latents + final latent + timesteps (RolloutDitTrajectory) - ) + rollout_return_denoising_env: bool = False # populate ``denoising_env`` (image/pos/neg kwargs, guidance) for RL replay + rollout_return_dit_trajectory: bool = False # per-step noisy latents + final latent + timesteps (RolloutDitTrajectory) # 0-indexed denoising-loop step filters; None = all steps. rollout_sde_step_indices: list[int] | None = None rollout_return_step_indices: list[int] | None = None @@ -488,8 +482,7 @@ class SamplingParams: if self.quality not in QUALITY_LEVELS: raise ValueError( - f"quality must be one of {list(QUALITY_LEVELS)}, " - f"got {self.quality!r}" + f"quality must be one of {list(QUALITY_LEVELS)}, got {self.quality!r}" ) # These are always required to be sane regardless of pipeline. diff --git a/python/sglang/multimodal_gen/configs/sample/wan.py b/python/sglang/multimodal_gen/configs/sample/wan.py index 0464ed60e..dfea73974 100644 --- a/python/sglang/multimodal_gen/configs/sample/wan.py +++ b/python/sglang/multimodal_gen/configs/sample/wan.py @@ -51,9 +51,7 @@ class WanT2V_1_3B_SamplingParams(SamplingParams): # Denoising stage guidance_scale: float = 3.0 - negative_prompt: str = ( - "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" - ) + negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" num_inference_steps: int = 50 # Wan T2V 1.3B supported resolutions @@ -85,9 +83,7 @@ class WanT2V_14B_SamplingParams(SamplingParams): # Denoising stage guidance_scale: float = 5.0 - negative_prompt: str = ( - "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" - ) + negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" num_inference_steps: int = 50 # Wan T2V 14B supported resolutions diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index cf76d3422..a31842f16 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -784,9 +784,11 @@ def _register_configs(): hf_model_paths=["Lightricks/LTX-2"], model_detectors=[ lambda path: "ltx" in path.lower() and "video" in path.lower(), - lambda path: "ltx-2" in path.lower() - and "ltx-2.3" not in path.lower() - and "ltx-2.5" not in path.lower(), + lambda path: ( + "ltx-2" in path.lower() + and "ltx-2.3" not in path.lower() + and "ltx-2.5" not in path.lower() + ), ], ) register_configs( @@ -973,8 +975,9 @@ def _register_configs(): "MiniMax/MiniMax-H3", ], model_detectors=[ - lambda model_id: "minimaxh3" - in model_id.lower().replace("-", "").replace("_", "") + lambda model_id: ( + "minimaxh3" in model_id.lower().replace("-", "").replace("_", "") + ) ], ) register_configs( @@ -984,8 +987,9 @@ def _register_configs(): "FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree", ], model_detectors=[ - lambda model_id: "fasth3" - in model_id.lower().replace("-", "").replace("_", "") + lambda model_id: ( + "fasth3" in model_id.lower().replace("-", "").replace("_", "") + ) ], ) # FLUX @@ -1191,7 +1195,7 @@ def _register_configs(): ], model_detectors=[ # Match "sana-wm" or "sana_wm" but NOT plain T2I "sana" checkpoints. - lambda hf_id: ("sana-wm" in hf_id.lower() or "sana_wm" in hf_id.lower()), + lambda hf_id: "sana-wm" in hf_id.lower() or "sana_wm" in hf_id.lower(), ], ) @@ -1203,9 +1207,7 @@ def _register_configs(): "Efficient-Large-Model/SANA-Video_2B_480p_diffusers", ], model_detectors=[ - lambda hf_id: ( - "sana-video" in hf_id.lower() or "sana_video" in hf_id.lower() - ) + lambda hf_id: "sana-video" in hf_id.lower() or "sana_video" in hf_id.lower() ], ) @@ -1293,8 +1295,10 @@ def _register_configs(): "jdopensource/JoyAI-Echo", ], model_detectors=[ - lambda hf_id: ("joy-echo" in hf_id.lower() or "joyai-echo" in hf_id.lower()) - and "image-edit" not in hf_id.lower(), + lambda hf_id: ( + ("joy-echo" in hf_id.lower() or "joyai-echo" in hf_id.lower()) + and "image-edit" not in hf_id.lower() + ), ], ) @@ -1354,9 +1358,11 @@ def _register_configs(): "meituan-longcat/LongCat-Image-Edit-Turbo", ], model_detectors=[ - lambda hf_id: "longcat" in hf_id.lower() - and "edit" in hf_id.lower() - and "turbo" in hf_id.lower(), + lambda hf_id: ( + "longcat" in hf_id.lower() + and "edit" in hf_id.lower() + and "turbo" in hf_id.lower() + ), ], ) @@ -1368,9 +1374,11 @@ def _register_configs(): "meituan-longcat/LongCat-Image-Edit", ], model_detectors=[ - lambda hf_id: "longcat" in hf_id.lower() - and "edit" in hf_id.lower() - and "turbo" not in hf_id.lower(), + lambda hf_id: ( + "longcat" in hf_id.lower() + and "edit" in hf_id.lower() + and "turbo" not in hf_id.lower() + ), ], ) diff --git a/python/sglang/multimodal_gen/runtime/disaggregation/scheduler_mixin.py b/python/sglang/multimodal_gen/runtime/disaggregation/scheduler_mixin.py index 3f868280f..998241307 100644 --- a/python/sglang/multimodal_gen/runtime/disaggregation/scheduler_mixin.py +++ b/python/sglang/multimodal_gen/runtime/disaggregation/scheduler_mixin.py @@ -1083,7 +1083,7 @@ class SchedulerDisaggMixin: ) use_prefetch = self._compute_ready_queue is not None logger.info( - "Pool mode %s rank %d event loop started " "(multi_rank=%s, prefetch=%s)", + "Pool mode %s rank %d event loop started (multi_rank=%s, prefetch=%s)", role_name, self.gpu_id, is_multi_rank, diff --git a/python/sglang/multimodal_gen/runtime/disaggregation/transport/buffer.py b/python/sglang/multimodal_gen/runtime/disaggregation/transport/buffer.py index 5afba92e2..cd44a0a49 100644 --- a/python/sglang/multimodal_gen/runtime/disaggregation/transport/buffer.py +++ b/python/sglang/multimodal_gen/runtime/disaggregation/transport/buffer.py @@ -54,8 +54,7 @@ class TransferTensorBuffer: pool_location = "pinned CPU" if device == "cpu" else f"GPU ({device})" logger.info( - "TransferTensorBuffer[%s]: allocated %d MiB %s memory " - "(min_block=%d KiB)", + "TransferTensorBuffer[%s]: allocated %d MiB %s memory (min_block=%d KiB)", role_name, actual_size >> 20, pool_location, diff --git a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/base_device_communicator.py b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/base_device_communicator.py index f76d53937..be3dd2552 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/base_device_communicator.py +++ b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/base_device_communicator.py @@ -165,9 +165,9 @@ class DistributedAutograd: if world_size == 1: return input_ - assert ( - input_.dim() == 4 - ), f"input must be 4D tensor, got {input_.dim()} and shape {input_.shape}" + assert input_.dim() == 4, ( + f"input must be 4D tensor, got {input_.dim()} and shape {input_.shape}" + ) if world_size == 2 and scatter_dim in (1, 2): fast = _ipc_all_to_all_4d(group, input_, scatter_dim) @@ -303,9 +303,9 @@ class DeviceCommunicatorBase: NOTE: `dst` is the local rank of the destination rank. """ world_size = self.world_size - assert ( - -input_.dim() <= dim < input_.dim() - ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + assert -input_.dim() <= dim < input_.dim(), ( + f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + ) if dim < 0: # Convert negative dim to positive. dim += input_.dim() diff --git a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/cpu_communicator.py b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/cpu_communicator.py index 3345cc379..fd9b13070 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/cpu_communicator.py +++ b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/cpu_communicator.py @@ -15,7 +15,6 @@ from .base_device_communicator import DeviceCommunicatorBase class CpuCommunicator(DeviceCommunicatorBase): - def __init__( self, cpu_group: ProcessGroup, @@ -55,9 +54,9 @@ class CpuCommunicator(DeviceCommunicatorBase): NOTE: `dst` is the local rank of the destination rank. """ world_size = self.world_size - assert ( - -input_.dim() <= dim < input_.dim() - ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + assert -input_.dim() <= dim < input_.dim(), ( + f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + ) if dim < 0: # Convert negative dim to positive. dim += input_.dim() @@ -107,7 +106,6 @@ class CpuCommunicator(DeviceCommunicatorBase): class _CPUSHMDistributed: - def __init__(self, communicator: CpuCommunicator): instance_identifier = os.environ["VLLM_DIST_IDENT"] unique_name = communicator.unique_name diff --git a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/cuda_communicator.py b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/cuda_communicator.py index a9a0625fb..307f381ef 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/cuda_communicator.py +++ b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/cuda_communicator.py @@ -13,7 +13,6 @@ from sglang.multimodal_gen.runtime.distributed.device_communicators.base_device_ class CudaCommunicator(DeviceCommunicatorBase): - def __init__( self, cpu_group: ProcessGroup, diff --git a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/pynccl.py b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/pynccl.py index 45a9ff40f..64d551b5c 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/pynccl.py +++ b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/pynccl.py @@ -28,7 +28,6 @@ logger = init_logger(__name__) class PyNcclCommunicator: - def __init__( self, group: ProcessGroup | StatelessProcessGroup, @@ -48,9 +47,9 @@ class PyNcclCommunicator: """ if not isinstance(group, StatelessProcessGroup): assert dist.is_initialized() - assert ( - dist.get_backend(group) != dist.Backend.NCCL - ), "PyNcclCommunicator should be attached to a non-NCCL group." + assert dist.get_backend(group) != dist.Backend.NCCL, ( + "PyNcclCommunicator should be attached to a non-NCCL group." + ) # note: this rank is the rank in the group self.rank = dist.get_rank(group) self.world_size = dist.get_world_size(group) diff --git a/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py b/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py index 9e06b97fb..d4d502ac4 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py +++ b/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py @@ -87,8 +87,7 @@ def _split_tensor_dict( tensor_list = [] for key, value in tensor_dict.items(): assert "%" not in key, ( - "Avoid having '%' in key " - "as it is used as a separator for nested entries." + "Avoid having '%' in key as it is used as a separator for nested entries." ) if isinstance(value, torch.Tensor): # Note: we cannot use `value.device` here, @@ -425,9 +424,9 @@ class GroupCoordinator: # Bypass the function if we are using only 1 GPU. if world_size == 1: return input_ - assert ( - -input_.dim() <= dim < input_.dim() - ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + assert -input_.dim() <= dim < input_.dim(), ( + f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + ) if dim < 0: # Convert negative dim to positive. dim += input_.dim() @@ -481,9 +480,9 @@ class GroupCoordinator: # Bypass the function if we are using only 1 GPU. if world_size == 1: return input_ - assert ( - -input_.dim() <= dim < input_.dim() - ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + assert -input_.dim() <= dim < input_.dim(), ( + f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + ) if dim < 0: # Convert negative dim to positive. dim += input_.dim() @@ -599,9 +598,9 @@ class GroupCoordinator: assert src < self.world_size, f"Invalid src rank ({src})" - assert ( - src != self.rank - ), "Invalid source rank. Source rank is the same as the current rank." + assert src != self.rank, ( + "Invalid source rank. Source rank is the same as the current rank." + ) size_tensor = torch.empty(1, dtype=torch.long, device="cpu") @@ -621,9 +620,9 @@ class GroupCoordinator: object_tensor, src=self.ranks[src], group=self.cpu_group ) - assert ( - rank_object == rank_size - ), "Received object sender rank does not match the size sender rank." + assert rank_object == rank_size, ( + "Received object sender rank does not match the size sender rank." + ) obj = pickle.loads(object_tensor.numpy().tobytes()) @@ -652,9 +651,9 @@ class GroupCoordinator: rank = self.rank if rank == src_global_rank: metadata_list: List[Tuple[Any, Any]] = [] - assert isinstance( - tensor_dict, dict - ), f"Expecting a dictionary, got {type(tensor_dict)}" + assert isinstance(tensor_dict, dict), ( + f"Expecting a dictionary, got {type(tensor_dict)}" + ) metadata_list, tensor_list = _split_tensor_dict(tensor_dict) # `metadata_list` lives in CPU memory. # `broadcast_object_list` has serialization & deserialization, @@ -736,9 +735,9 @@ class GroupCoordinator: assert dst < self.world_size, f"Invalid dst rank ({dst})" metadata_list: List[Tuple[Any, Any]] = [] - assert isinstance( - tensor_dict, dict - ), f"Expecting a dictionary, got {type(tensor_dict)}" + assert isinstance(tensor_dict, dict), ( + f"Expecting a dictionary, got {type(tensor_dict)}" + ) metadata_list, tensor_list = _split_tensor_dict(tensor_dict) # `metadata_list` lives in CPU memory. # `send_object_list` has serialization & deserialization, @@ -1193,14 +1192,14 @@ class PipelineGroupCoordinator(GroupCoordinator): def get_pipeline_recv_data( self, idx: int = -1, name: str = "latent" ) -> torch.Tensor: - assert ( - len(self.receiving_tasks) > 0 - ), "No tasks to receive, call add_pipeline_recv_task first" + assert len(self.receiving_tasks) > 0, ( + "No tasks to receive, call add_pipeline_recv_task first" + ) receiving_task = self.receiving_tasks.pop(0) receiving_task[0].wait() - assert ( - receiving_task[1] == name and receiving_task[2] == idx - ), "Received tensor does not match the requested" + assert receiving_task[1] == name and receiving_task[2] == idx, ( + "Received tensor does not match the requested" + ) return self.recv_buffer[name][idx] def _pipeline_irecv(self, tensor: torch.tensor): @@ -1255,14 +1254,14 @@ class PipelineGroupCoordinator(GroupCoordinator): self.recv_skip_tasks_queue.append(idx) def get_pipeline_recv_skip_data(self, idx: int = -1) -> torch.Tensor: - assert ( - len(self.receiving_skip_tasks) > 0 - ), "No tasks to receive, call add_pipeline_recv_skip_task first" + assert len(self.receiving_skip_tasks) > 0, ( + "No tasks to receive, call add_pipeline_recv_skip_task first" + ) receiving_skip_task = self.receiving_skip_tasks.pop(0) receiving_skip_task[0].wait() - assert ( - receiving_skip_task[2] == idx - ), "Received tensor does not match the requested" + assert receiving_skip_task[2] == idx, ( + "Received tensor does not match the requested" + ) return self.skip_tensor_recv_buffer[idx] def recv_skip_next(self): diff --git a/python/sglang/multimodal_gen/runtime/distributed/parallel_groups.py b/python/sglang/multimodal_gen/runtime/distributed/parallel_groups.py index d1d3bc2f6..0c5476bb8 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/parallel_groups.py +++ b/python/sglang/multimodal_gen/runtime/distributed/parallel_groups.py @@ -43,9 +43,9 @@ def set_seq_parallel_pg_by_sp_groups( """ sp_degree = sp_ring_degree * sp_ulysses_degree assert sp_degree > 0 - assert all( - len(g) == sp_degree for g in sp_groups - ), f"Each SP group must have size {sp_degree}, got sizes {[len(g) for g in sp_groups]}" + assert all(len(g) == sp_degree for g in sp_groups), ( + f"Each SP group must have size {sp_degree}, got sizes {[len(g) for g in sp_groups]}" + ) ulyssess_pg = None ring_pg = None diff --git a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py index eb642a83b..a22b54dfc 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py +++ b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py @@ -298,7 +298,6 @@ def init_distributed_environment( ) if timeout is not None: - extra_args["timeout"] = datetime.timedelta(seconds=timeout) logger.info(f"Setting distributed timeout to {timeout} seconds") @@ -325,9 +324,9 @@ def init_distributed_environment( ranks = list(range(torch.distributed.get_world_size())) _WORLD = init_world_group(ranks, local_rank, backend) else: - assert ( - _WORLD.world_size == torch.distributed.get_world_size() - ), "world group already initialized with a different world size" + assert _WORLD.world_size == torch.distributed.get_world_size(), ( + "world group already initialized with a different world size" + ) _sync_srt_world_group() @@ -648,12 +647,12 @@ def maybe_init_distributed_environment_and_model_parallel( if _WORLD is not None and model_parallel_is_initialized(): # make sure the tp and sp sizes are correct - assert ( - get_tp_world_size() == tp_size - ), f"You are trying to initialize model parallel groups with size {tp_size}, but they are already initialized with size {get_tp_world_size()}" - assert ( - get_sp_world_size() == sp_size - ), f"You are trying to initialize model parallel groups with size {sp_size}, but they are already initialized with size {get_sp_world_size()}" + assert get_tp_world_size() == tp_size, ( + f"You are trying to initialize model parallel groups with size {tp_size}, but they are already initialized with size {get_tp_world_size()}" + ) + assert get_sp_world_size() == sp_size, ( + f"You are trying to initialize model parallel groups with size {sp_size}, but they are already initialized with size {get_sp_world_size()}" + ) return local_rank = int(os.environ.get("LOCAL_RANK", 0)) world_size = int(os.environ.get("WORLD_SIZE", 1)) @@ -773,9 +772,9 @@ def is_the_same_node_as( memory system (shared access to shared memory). """ if isinstance(pg, ProcessGroup): - assert ( - torch.distributed.get_backend(pg) != torch.distributed.Backend.NCCL - ), "in_the_same_node_as should be tested with a non-NCCL group." + assert torch.distributed.get_backend(pg) != torch.distributed.Backend.NCCL, ( + "in_the_same_node_as should be tested with a non-NCCL group." + ) # local rank inside the group rank = torch.distributed.get_rank(group=pg) world_size = torch.distributed.get_world_size(group=pg) @@ -933,9 +932,9 @@ def is_pipeline_last_stage() -> bool: # CFG def get_cfg_group() -> GroupCoordinator: - assert ( - _CFG is not None - ), "classifier_free_guidance parallel group is not initialized" + assert _CFG is not None, ( + "classifier_free_guidance parallel group is not initialized" + ) return _CFG diff --git a/python/sglang/multimodal_gen/runtime/distributed/utils.py b/python/sglang/multimodal_gen/runtime/distributed/utils.py index ef481c9e3..a3577610d 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/utils.py +++ b/python/sglang/multimodal_gen/runtime/distributed/utils.py @@ -134,13 +134,13 @@ class StatelessProcessGroup: """ if self.rank == src: self.expire_data() - key = f"broadcast_from/{src}/" f"{self.broadcast_send_counter}" + key = f"broadcast_from/{src}/{self.broadcast_send_counter}" self.store.set(key, pickle.dumps(obj)) self.broadcast_send_counter += 1 self.entries.append((key, time.perf_counter())) return obj else: - key = f"broadcast_from/{src}/" f"{self.broadcast_recv_src_counter[src]}" + key = f"broadcast_from/{src}/{self.broadcast_recv_src_counter[src]}" recv_obj = pickle.loads(self.store.get(key)) self.broadcast_recv_src_counter[src] += 1 return recv_obj diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/cli/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/cli/utils.py index 1699fca94..955df8f3a 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/cli/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/cli/utils.py @@ -14,7 +14,6 @@ logger = init_logger(__name__) class RaiseNotImplementedAction(argparse.Action): - def __call__(self, parser, namespace, values, option_string=None): raise NotImplementedError(f"The {option_string} option is not yet implemented") diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py index 3b600bbba..71492523d 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py @@ -209,9 +209,9 @@ def _build_response( peak_memory_mb = result.peak_memory_mb if result.peak_memory_mb > 0 else None rollout_trajectory_data = result.rollout_trajectory_data if rollout: - assert ( - rollout_trajectory_data is not None - ), "rollout_trajectory_data must be present when rollout=True" + assert rollout_trajectory_data is not None, ( + "rollout_trajectory_data must be present when rollout=True" + ) serialized_dit_timesteps = None serialized_dit_sigmas = None diff --git a/python/sglang/multimodal_gen/runtime/launch_server.py b/python/sglang/multimodal_gen/runtime/launch_server.py index df6fed265..f4b655ae5 100644 --- a/python/sglang/multimodal_gen/runtime/launch_server.py +++ b/python/sglang/multimodal_gen/runtime/launch_server.py @@ -662,7 +662,7 @@ def launch_disagg_role(server_args: ServerArgs): role_type = server_args.disagg_role if server_args.disagg_server_addr is None: raise ValueError( - "--disagg-server-addr is required for --disagg-role " f"{role_type.value}" + f"--disagg-server-addr is required for --disagg-role {role_type.value}" ) # Derive endpoints diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/STA_configuration.py b/python/sglang/multimodal_gen/runtime/layers/attention/STA_configuration.py index 9635a6740..5c17d0405 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/STA_configuration.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/STA_configuration.py @@ -138,7 +138,9 @@ def configure_sta( print("\nStrategy usage counts:") total_heads = time_step_num * layer_num * head_num # Fixed dimensions for strategy, count in strategy_counts.items(): - print(f"Strategy {strategy}: {count} heads ({count/total_heads*100:.2f}%)") + print( + f"Strategy {strategy}: {count} heads ({count / total_heads * 100:.2f}%)" + ) # Convert dictionary to 3D list with fixed dimensions mask_strategy_3d = dict_to_3d_list( @@ -221,7 +223,9 @@ def configure_sta( print("\nStrategy usage counts:") total_heads = time_step_num * layer_num * head_num # Fixed dimensions for strategy, count in strategy_counts.items(): - print(f"Strategy {strategy}: {count} heads ({count/total_heads*100:.2f}%)") + print( + f"Strategy {strategy}: {count} heads ({count / total_heads * 100:.2f}%)" + ) # Convert dictionary to 3D list with fixed dimensions mask_strategy_3d = dict_to_3d_list( diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter_sage.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter_sage.py index 56c274c42..17ad922a0 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter_sage.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter_sage.py @@ -13,7 +13,6 @@ from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum class AITERSageBackend(AttentionBackend): - @staticmethod def get_enum() -> AttentionBackendEnum: return AttentionBackendEnum.AITER_SAGE @@ -35,7 +34,6 @@ class AITERSageBackend(AttentionBackend): class AITERSageImpl(AttentionImpl): - def __init__( self, num_heads: int, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/ascend_fa.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/ascend_fa.py index b25388ae4..98e74fada 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/ascend_fa.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/ascend_fa.py @@ -173,7 +173,6 @@ class AscendFAMetadataBuilder(AttentionMetadataBuilder): class AscendFABackend(AttentionBackend): - @staticmethod def get_enum() -> AttentionBackendEnum: return AttentionBackendEnum.FA @@ -198,7 +197,6 @@ class AscendFABackend(AttentionBackend): class AscendFAImpl(AttentionImpl): - def __init__( self, num_heads: int, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py index 733b3bbe4..7f50900f3 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py @@ -125,7 +125,6 @@ class AttentionMetadataBuilder(ABC, Generic[T]): class AttentionLayer(Protocol): - _k_scale: torch.Tensor _v_scale: torch.Tensor _k_scale_float: float @@ -142,7 +141,6 @@ class AttentionLayer(Protocol): class AttentionImpl(ABC, Generic[T]): - @abstractmethod def __init__( self, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/block_sparse_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/block_sparse_attn.py index fa6df1ef2..f51ffa9ac 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/block_sparse_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/block_sparse_attn.py @@ -21,7 +21,6 @@ BSA_BLOCK_SIZE = 128 class BlockSparseAttentionBackend(AttentionBackend): - accept_output_buffer: bool = True @staticmethod @@ -126,7 +125,6 @@ class BlockSparseAttentionMetadataBuilder(AttentionMetadataBuilder): class BlockSparseAttentionImpl(AttentionImpl): - def __init__( self, num_heads: int, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/cube_sparse_attn/backend.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/cube_sparse_attn/backend.py index 26f2be21d..50c4d94ac 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/cube_sparse_attn/backend.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/cube_sparse_attn/backend.py @@ -36,7 +36,6 @@ from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum class CubeSparseAttentionBackend(AttentionBackend): - @staticmethod def get_enum() -> AttentionBackendEnum: return AttentionBackendEnum.CUBE_SPARSE_ATTN @@ -63,7 +62,6 @@ class CubeSparseAttentionMetadata(AttentionMetadata): class CubeSparseAttentionMetadataBuilder(AttentionMetadataBuilder): - def __init__(self): pass @@ -279,7 +277,6 @@ def cube_sparse_attention( class CubeSparseAttentionImpl(AttentionImpl): - def __init__( self, num_heads: int, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py index 5710e321d..b0bb625c9 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py @@ -67,9 +67,9 @@ def flash_attn_varlen_func_fake_out( head_dim_v = v.shape[-1] if cu_seqlens_q is not None: - assert cu_seqlens_q.shape == ( - batch_size + 1, - ), "cu_seqlens_q must have shape (batch_size + 1,)" + assert cu_seqlens_q.shape == (batch_size + 1,), ( + "cu_seqlens_q must have shape (batch_size + 1,)" + ) assert cu_seqlens_q.dtype == torch.int32, "cu_seqlens_q must be int32" assert cu_seqlens_q.stride(0) == 1, "cu_seqlens_q must be contiguous" @@ -129,9 +129,9 @@ def flash_attn_varlen_func_fake_out_lse( head_dim_v = v.shape[-1] if cu_seqlens_q is not None: - assert cu_seqlens_q.shape == ( - batch_size + 1, - ), "cu_seqlens_q must have shape (batch_size + 1,)" + assert cu_seqlens_q.shape == (batch_size + 1,), ( + "cu_seqlens_q must have shape (batch_size + 1,)" + ) assert cu_seqlens_q.dtype == torch.int32, "cu_seqlens_q must be int32" assert cu_seqlens_q.stride(0) == 1, "cu_seqlens_q must be contiguous" @@ -329,7 +329,6 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder): class FlashAttentionBackend(AttentionBackend): - @classmethod def supports_ring_rotation(cls) -> bool: return True diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn_2.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn_2.py index 62a1974ad..7b6b8df43 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn_2.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn_2.py @@ -44,7 +44,6 @@ class FlashAttention2Backend(AttentionBackend): class FlashAttention2Impl(AttentionImpl): - def __init__( self, num_heads: int, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/laser_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/laser_attn.py index 487a4f936..5afb82b95 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/laser_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/laser_attn.py @@ -32,7 +32,6 @@ _BF16_LASER_SCALE = 256.0 class LaserAttentionBackend(AttentionBackend): - accept_output_buffer: bool = True @staticmethod @@ -49,7 +48,6 @@ class LaserAttentionBackend(AttentionBackend): class LaserAttentionImpl(AttentionImpl): - def __init__( self, num_heads: int, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/rain_fusion_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/rain_fusion_attn.py index 95ea1374b..48245874e 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/rain_fusion_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/rain_fusion_attn.py @@ -22,7 +22,6 @@ logger = init_logger(__name__) class RainFusionAttentionBackend(AttentionBackend): - accept_output_buffer: bool = True @staticmethod @@ -99,7 +98,6 @@ class RainFusionAttentionMetadataBuilder(AttentionMetadataBuilder): class RainFusionAttentionImpl(AttentionImpl): - def __init__( self, num_heads: int, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py index c4ffbaea5..a137d8017 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py @@ -33,7 +33,6 @@ def _trailing_padding_used_len( class SageAttentionBackend(AttentionBackend): - @classmethod def supports_ring_rotation(cls) -> bool: return True @@ -54,7 +53,6 @@ class SageAttentionBackend(AttentionBackend): class SageAttentionImpl(AttentionImpl): - def __init__( self, num_heads: int, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py index 02bbd5e6b..73f55158c 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py @@ -28,7 +28,6 @@ _MPS_VARLEN_QUERY_CHUNK_SIZE = 128 class SDPABackend(AttentionBackend): - accept_output_buffer: bool = True @staticmethod @@ -49,7 +48,6 @@ class SDPABackend(AttentionBackend): class SDPAImpl(AttentionImpl): - def __init__( self, num_heads: int, @@ -268,7 +266,7 @@ class DynamicCudnnSDPAImpl(SDPAImpl): # cuDNN raises "No available kernel" for some shapes; pin the # FA fail-safe path for this layer and keep going. logger.warning( - "cuDNN SDPA failed (%s); falling back to FlashAttention " "for %s.", + "cuDNN SDPA failed (%s); falling back to FlashAttention for %s.", e, type(self).__name__, ) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sliding_tile_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sliding_tile_attn.py index 37a1acf30..0db59e292 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sliding_tile_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sliding_tile_attn.py @@ -35,7 +35,6 @@ logger = init_logger(__name__) class RangeDict(dict): - def __getitem__(self, item: int) -> str: for key in self.keys(): if isinstance(key, tuple): @@ -81,7 +80,6 @@ class SlidingTileAttentionMetadata(AttentionMetadata): class SlidingTileAttentionMetadataBuilder(AttentionMetadataBuilder): - def __init__(self): pass @@ -105,7 +103,6 @@ class SlidingTileAttentionMetadataBuilder(AttentionMetadataBuilder): class SlidingTileAttentionImpl(AttentionImpl): - def __init__( self, num_heads: int, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sol_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sol_attn.py index b7be98848..5f984631a 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sol_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sol_attn.py @@ -99,7 +99,6 @@ class SolAttnBackend(AttentionBackend): class SolAttnImpl(AttentionImpl): - def __init__( self, num_heads: int, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_linear_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_linear_attn.py index 3acae28e6..b3dcd6dee 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_linear_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_linear_attn.py @@ -356,9 +356,9 @@ class SageSparseLinearAttentionImpl(AttentionImpl, nn.Module): ) -> None: nn.Module.__init__(self) - assert ( - SAGESLA_ENABLED - ), "Install spas_sage_attn(pip install git+https://github.com/thu-ml/SpargeAttn.git --no-build-isolation) first to enable SageSLA." + assert SAGESLA_ENABLED, ( + "Install spas_sage_attn(pip install git+https://github.com/thu-ml/SpargeAttn.git --no-build-isolation) first to enable SageSLA." + ) self.num_heads = num_heads self.head_size = head_size @@ -454,7 +454,9 @@ class SageSparseLinearAttentionImpl(AttentionImpl, nn.Module): assert headdim in [ 64, 128, - ], "headdim should be in [64, 128]. For other headdim, you can use padding and specify the softmax scale." + ], ( + "headdim should be in [64, 128]. For other headdim, you can use padding and specify the softmax scale." + ) # Quantize Q, K to INT8 q_int8, q_scale, k_int8, k_scale = get_vanilla_qk_quant(q, k, km, BLKQ, BLKK) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_video_gen_2_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_video_gen_2_attn.py index 0d07259c0..f5d03072c 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_video_gen_2_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_video_gen_2_attn.py @@ -42,7 +42,6 @@ logger = init_logger(__name__) class SparseVideoGen2AttentionBackend(AttentionBackend): - accept_output_buffer: bool = True @staticmethod @@ -116,7 +115,6 @@ def _require_kwarg(kwargs: dict[str, Any], name: str) -> Any: class SparseVideoGen2AttentionMetadataBuilder(AttentionMetadataBuilder): - def __init__(self) -> None: pass @@ -180,7 +178,6 @@ class SparseVideoGen2AttentionMetadataBuilder(AttentionMetadataBuilder): class SparseVideoGen2AttentionImpl(AttentionImpl): - def __init__( self, num_heads: int, @@ -461,9 +458,9 @@ class SparseVideoGen2AttentionImpl(AttentionImpl): if prompt_length is None: prompt_length = context_length - assert ( - seq_len == context_length + num_frame * frame_size - ), f"Query Shape: {seq_len} is not equivalent to {context_length} + {num_frame} * {frame_size}" + assert seq_len == context_length + num_frame * frame_size, ( + f"Query Shape: {seq_len} is not equivalent to {context_length} + {num_frame} * {frame_size}" + ) # Determine if we use Full Attention to calculate full_attention_flag = False diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn.py index 4f72dcc1f..6607321a8 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn.py @@ -98,9 +98,7 @@ def construct_variable_block_sizes( t_sizes[:, None, None] # [n_t, 1, 1] * h_sizes[None, :, None] # [1, n_h, 1] * w_sizes[None, None, :] # [1, 1, n_w] - ).reshape( - -1 - ) # [n_t * n_h * n_w] + ).reshape(-1) # [n_t * n_h * n_w] return block_sizes diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn_h3.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn_h3.py index 8367dce0f..2a0812284 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn_h3.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn_h3.py @@ -393,7 +393,7 @@ class VideoSparseAttentionH3Impl(AttentionImpl): attn_metadata: AttentionMetadata, ) -> torch.Tensor: raise NotImplementedError( - "VSA-H3 serves MiniMax-H3's packed varlen attention; use " "forward_varlen." + "VSA-H3 serves MiniMax-H3's packed varlen attention; use forward_varlen." ) def forward_varlen( diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/vmoba.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/vmoba.py index e07c74336..614c38307 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/vmoba.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/vmoba.py @@ -26,7 +26,6 @@ logger = init_logger(__name__) class VMOBAAttentionBackend(AttentionBackend): - accept_output_buffer: bool = True @staticmethod @@ -89,7 +88,6 @@ def pad_input(hidden_states, indices, batch, seqlen): class VideoMobaAttentionMetadataBuilder(AttentionMetadataBuilder): - def __init__(self): pass @@ -124,7 +122,9 @@ class VideoMobaAttentionMetadataBuilder(AttentionMetadataBuilder): raw_latent_shape[0] % patch_size[0] == 0 and raw_latent_shape[1] % patch_size[1] == 0 and raw_latent_shape[2] % patch_size[2] == 0 - ), f"spatial patch_resolution {raw_latent_shape} should be divisible by patch_size {patch_size}" + ), ( + f"spatial patch_resolution {raw_latent_shape} should be divisible by patch_size {patch_size}" + ) patch_resolution = [ t // pt for t, pt in zip(raw_latent_shape, patch_size, strict=False) ] @@ -150,7 +150,6 @@ class VideoMobaAttentionMetadataBuilder(AttentionMetadataBuilder): class VMOBAAttentionImpl(AttentionImpl): - def __init__( self, num_heads, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/xpu_backend.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/xpu_backend.py index 613f2df66..783010317 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/xpu_backend.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/xpu_backend.py @@ -62,7 +62,6 @@ def _get_cu_seqlens(device_index: int, bsz: int, seqlen: int) -> torch.Tensor: class XPUAttentionImpl(AttentionImpl): - def __init__( self, num_heads: int, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index f0776fa5d..800d04400 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py @@ -562,9 +562,9 @@ class UlyssesAttention_VSA(UlyssesAttention): "K/V-gather SP does not support video sparse attention." ) # Check text tokens are not supported for VSA now - assert ( - replicated_q is None and replicated_k is None and replicated_v is None - ), "Replicated QKV is not supported for VSA now" + assert replicated_q is None and replicated_k is None and replicated_v is None, ( + "Replicated QKV is not supported for VSA now" + ) # Check input shapes assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "Expected 4D tensors" @@ -1042,9 +1042,9 @@ class USPAttention(nn.Module): inv_indices = attn_mask_meta["inv_indices"] # Guard against a caller passing meta from a different # mask shape (silent corruption otherwise). - assert ( - inv_indices.shape[0] == bs * seq - ), "attn_mask_meta shape does not match attn_mask" + assert inv_indices.shape[0] == bs * seq, ( + "attn_mask_meta shape does not match attn_mask" + ) # All-False mask: FA varlen rejects zero-length input. # Fall through to SDPA which handles it via broadcast. # (Joint attention with an image side is always non-empty @@ -1169,9 +1169,9 @@ class USPAttention(nn.Module): # Zero-copy tail path: run varlen FA straight over the # padded layout, each row split into [valid | pad] segments # (contiguous reshapes only, no repacking). - assert ( - cu_tail.numel() == 2 * bs + 1 - ), "cu_seqlens_tail does not match the batch size" + assert cu_tail.numel() == 2 * bs + 1, ( + "cu_seqlens_tail does not match the batch size" + ) out = flash_attn_varlen_func( q=q.reshape(bs * seq, *q.shape[2:]), k=k.reshape(bs * seq, *k.shape[2:]), @@ -1259,9 +1259,9 @@ class USPAttention(nn.Module): gathered_mask_meta = build_varlen_mask_meta(gathered_mask) indices = gathered_mask_meta["indices"] inv_indices = gathered_mask_meta["inv_indices"] - assert ( - inv_indices.shape[0] == bs * seq - ), "gathered attn_mask shape does not match q/k/v" + assert inv_indices.shape[0] == bs * seq, ( + "gathered attn_mask shape does not match q/k/v" + ) if indices.shape[0] > 0: q_unpad, k_unpad, v_unpad = fused_pack_qkv(q, k, v, indices) out_unpad = flash_attn_varlen_func( diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/selector.py b/python/sglang/multimodal_gen/runtime/layers/attention/selector.py index 8b78133fc..f86cb9b94 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/selector.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/selector.py @@ -173,8 +173,7 @@ def _log_component_attn_backend_summary( backend_parts.append(backend_name) logger.info_once( - f"Attention backends for {context.component_name}: " - f"{', '.join(backend_parts)}" + f"Attention backends for {context.component_name}: {', '.join(backend_parts)}" ) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py index c40948d68..31e2c8103 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py @@ -49,9 +49,9 @@ def single_all_to_all(input, local_seq_2_local_head, group, async_op=False): # b, s, n, h if local_seq_2_local_head: bs, local_seq_len, num_total_head, head_dim = input.shape - assert ( - num_total_head % seq_world_size == 0 - ), f"Number of heads ({num_total_head}) must be divisible by the sequence parallel size ({seq_world_size})!" + assert num_total_head % seq_world_size == 0, ( + f"Number of heads ({num_total_head}) must be divisible by the sequence parallel size ({seq_world_size})!" + ) input_t = rearrange( input, "bs seq_len (w h) d -> w bs seq_len h d", diff --git a/python/sglang/multimodal_gen/runtime/layers/kvcache/causal_attention_cache.py b/python/sglang/multimodal_gen/runtime/layers/kvcache/causal_attention_cache.py index c8722b31c..e42284adc 100644 --- a/python/sglang/multimodal_gen/runtime/layers/kvcache/causal_attention_cache.py +++ b/python/sglang/multimodal_gen/runtime/layers/kvcache/causal_attention_cache.py @@ -221,8 +221,7 @@ class CausalSelfAttentionKVCache: self.k[:, sink_tokens : sink_tokens + num_rolled_tokens] = ( self.k[ :, - sink_tokens - + num_evicted_tokens : sink_tokens + sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens, ].clone() @@ -230,8 +229,7 @@ class CausalSelfAttentionKVCache: self.v[:, sink_tokens : sink_tokens + num_rolled_tokens] = ( self.v[ :, - sink_tokens - + num_evicted_tokens : sink_tokens + sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens, ].clone() @@ -244,8 +242,7 @@ class CausalSelfAttentionKVCache: :, ] = self.k[ :, - sink_tokens - + num_evicted_tokens : sink_tokens + sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens, cache_head_slice, @@ -258,8 +255,7 @@ class CausalSelfAttentionKVCache: :, ] = self.v[ :, - sink_tokens - + num_evicted_tokens : sink_tokens + sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens, cache_head_slice, diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 356d25ffa..c22de734a 100755 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -446,7 +446,6 @@ class LayerNorm(CustomOp): # FSDP's MixedPrecisionPolicy @CustomOp.register("fp32_layer_norm") class FP32LayerNorm(CustomOp, nn.LayerNorm): - def __init__( self, normalized_shape, diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py b/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py index 7dd00c526..ef6f0621d 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py @@ -80,7 +80,7 @@ def register_quantization_config(quantization: str): ) if not issubclass(quant_config_cls, QuantizationConfig): raise ValueError( - "The quantization config must be a subclass of " "`QuantizationConfig`." + "The quantization config must be a subclass of `QuantizationConfig`." ) _CUSTOMIZED_METHOD_TO_QUANT_CONFIG[quantization] = quant_config_cls QUANTIZATION_METHODS.append(quantization) diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_w4a4_config.py b/python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_w4a4_config.py index cde021f64..357a35e4a 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_w4a4_config.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_w4a4_config.py @@ -64,8 +64,7 @@ class KitchenW4A4Config(QuantizationConfig): continue if marker_format != "convrot_w4a4": raise ValueError( - f"Unsupported Comfy W4A4 format for {prefix!r}: " - f"{marker_format!r}" + f"Unsupported Comfy W4A4 format for {prefix!r}: {marker_format!r}" ) self._parse_marker(prefix, marker) diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_w4a8_config.py b/python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_w4a8_config.py index 38a12f667..f80d538ba 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_w4a8_config.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_w4a8_config.py @@ -55,8 +55,7 @@ class KitchenW4A8Config(QuantizationConfig): continue if marker_format != "asym_w4a8_int8": raise ValueError( - f"Unsupported Comfy W4A8 format for {prefix!r}: " - f"{marker_format!r}" + f"Unsupported Comfy W4A8 format for {prefix!r}: {marker_format!r}" ) if marker.get("convrot") is not True: raise ValueError( diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/fp8.py b/python/sglang/multimodal_gen/runtime/layers/quantization/fp8.py index ca39b2f0e..78e37c64a 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/fp8.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/fp8.py @@ -254,9 +254,9 @@ class Fp8LinearMethod(LinearMethodBase): ) layer.input_scale = None elif _is_cpu: - assert ( - _is_cpu_amx_available - ), "Fp8LinearMethod on CPU requires that CPU has AMX support" + assert _is_cpu_amx_available, ( + "Fp8LinearMethod on CPU requires that CPU has AMX support" + ) _amx_process_weight_after_loading(layer, ["weight"]) layer.weight_scale_inv = torch.nn.Parameter( layer.weight_scale_inv.data, requires_grad=False diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py index eebcc1c7a..c1b36a427 100755 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py @@ -122,7 +122,7 @@ def _prepare_nvfp4_swiglu_fusion_weights( ) if weight.shape[0] % 128 != 0: raise ValueError( - "Fused NVFP4 SwiGLU requires FC1 N % 128 == 0, " f"got N={weight.shape[0]}." + f"Fused NVFP4 SwiGLU requires FC1 N % 128 == 0, got N={weight.shape[0]}." ) # FLUX.2 stores [gate; up]. The kernel consumes 64-row groups in diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim.py b/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim.py index 62b84ee17..cd11a5539 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim.py @@ -201,7 +201,6 @@ class ModelSlimConfig(QuantizationConfig): class ModelSlimLinearMethod(LinearMethodBase): - def __init__(self, quantization_config: ModelSlimConfig): self.quantization_config = quantization_config diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp4_scheme.py b/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp4_scheme.py index 00bb925b9..e2a20a773 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp4_scheme.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp4_scheme.py @@ -35,7 +35,6 @@ MXFP4_DUAL_LEVEL_RATIO = 16 class ModelSlimMXFP4Scheme(ModelSlimLinearScheme): - def create_weights( self, layer: torch.nn.Module, diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp8_scheme.py b/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp8_scheme.py index c4c5dfb7b..44be13f47 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp8_scheme.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp8_scheme.py @@ -25,7 +25,6 @@ MXFP8_BLOCK_SIZE = 32 class ModelSlimMXFP8Scheme(ModelSlimLinearScheme): - def create_weights( self, layer: torch.nn.Module, diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp4.py b/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp4.py index e2ddc79b7..629a2f9a1 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp4.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp4.py @@ -185,8 +185,7 @@ class Mxfp4LinearMethod(LinearMethodBase): if any(fn is None for fn in (dynamic_mxfp4_quant, shuffle_weight, gemm_a4w4)): raise RuntimeError( - "aiter MXFP4 kernels not available. " - "Install aiter with MXFP4 support." + "aiter MXFP4 kernels not available. Install aiter with MXFP4 support." ) weight_data = layer.weight.data diff --git a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/factory.py b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/factory.py index 807660ea0..f03978264 100644 --- a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/factory.py +++ b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/factory.py @@ -118,9 +118,9 @@ def get_rotary_pos_embed( if rope_dim_list is None: rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)] - assert ( - sum(rope_dim_list) == head_dim - ), "sum(rope_dim_list) should equal to head_dim of attention layer" + assert sum(rope_dim_list) == head_dim, ( + "sum(rope_dim_list) should equal to head_dim of attention layer" + ) # Get SP info - now handled within NDRotaryEmbedding # sp_group = get_sp_group() diff --git a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/mrope.py b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/mrope.py index bfe77777f..e284afab8 100644 --- a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/mrope.py +++ b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/mrope.py @@ -298,9 +298,9 @@ class NDRotaryEmbedding(torch.nn.Module): self.theta_rescale_factor = [theta_rescale_factor[0]] * self.ndim else: self.theta_rescale_factor = theta_rescale_factor - assert ( - len(self.theta_rescale_factor) == self.ndim - ), "len(theta_rescale_factor) should equal to len(rope_dim_list)" + assert len(self.theta_rescale_factor) == self.ndim, ( + "len(theta_rescale_factor) should equal to len(rope_dim_list)" + ) if isinstance(interpolation_factor, (int, float)): self.interpolation_factor = [interpolation_factor] * self.ndim @@ -308,9 +308,9 @@ class NDRotaryEmbedding(torch.nn.Module): self.interpolation_factor = [interpolation_factor[0]] * self.ndim else: self.interpolation_factor = interpolation_factor - assert ( - len(self.interpolation_factor) == self.ndim - ), "len(interpolation_factor) should equal to len(rope_dim_list)" + assert len(self.interpolation_factor) == self.ndim, ( + "len(interpolation_factor) should equal to len(rope_dim_list)" + ) self.rope_generators: list[OneDRotaryEmbedding] = torch.nn.ModuleList() _config_to_gen_idx: dict[tuple, int] = {} diff --git a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py index 459a38620..57dc71211 100644 --- a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py +++ b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py @@ -195,7 +195,7 @@ def apply_flashinfer_rope_qk_inplace( raise ValueError("positions must be a 1D Tensor") if positions.numel() != bsz * seqlen: raise ValueError( - f"positions length must be bsz*seqlen={bsz*seqlen}, got {positions.numel()}" + f"positions length must be bsz*seqlen={bsz * seqlen}, got {positions.numel()}" ) positions = positions.to(device=q.device, dtype=torch.long) diff --git a/python/sglang/multimodal_gen/runtime/layers/usp.py b/python/sglang/multimodal_gen/runtime/layers/usp.py index 1ed06e785..8b2d66db1 100644 --- a/python/sglang/multimodal_gen/runtime/layers/usp.py +++ b/python/sglang/multimodal_gen/runtime/layers/usp.py @@ -312,9 +312,9 @@ def _usp_input_all_to_all(x: torch.Tensor, head_dim: int = 1) -> torch.Tensor: # Shape transition: [b, s_local, h_global, d] -> [h_global, b, s_local, d] permute_order = (2, 0, 1, 3) - assert ( - h_global % world_size == 0 - ), f"h_global ({h_global}) must be divisible by world_size ({world_size})" + assert h_global % world_size == 0, ( + f"h_global ({h_global}) must be divisible by world_size ({world_size})" + ) h_local, s_global = h_global // world_size, s_local * world_size @@ -488,9 +488,9 @@ def _usp_input_all_to_all_varlen( assert x.ndim == 4, f"x must have 4 dimensions, got {x.ndim}" assert head_dim in (1, 2), f"head_dim must be 1 or 2, got {head_dim}" - assert ( - len(seq_lens) == world_size - ), f"seq_lens must have length {world_size}, got {len(seq_lens)}" + assert len(seq_lens) == world_size, ( + f"seq_lens must have length {world_size}, got {len(seq_lens)}" + ) rank = get_ulysses_parallel_rank() @@ -504,12 +504,12 @@ def _usp_input_all_to_all_varlen( # Shape transition: [b, s_local, h_global, d] -> [h_global, b, s_local, d] permute_order = (2, 0, 1, 3) - assert ( - s_local == seq_lens[rank] - ), f"s_local ({s_local}) must equal seq_lens[{rank}] ({seq_lens[rank]})" - assert ( - h_global % world_size == 0 - ), f"h_global ({h_global}) must be divisible by world_size ({world_size})" + assert s_local == seq_lens[rank], ( + f"s_local ({s_local}) must equal seq_lens[{rank}] ({seq_lens[rank]})" + ) + assert h_global % world_size == 0, ( + f"h_global ({h_global}) must be divisible by world_size ({world_size})" + ) h_local = h_global // world_size @@ -578,9 +578,9 @@ def _usp_output_all_to_all(x: torch.Tensor, head_dim: int = 1) -> torch.Tensor: # Shape transition: [b, s_global, h_local, d] -> [s_global, b, h_local, d] permute_order = (1, 0, 2, 3) - assert ( - s_global % world_size == 0 - ), f"s_global ({s_global}) must be divisible by world_size ({world_size})" + assert s_global % world_size == 0, ( + f"s_global ({s_global}) must be divisible by world_size ({world_size})" + ) s_local, h_global = s_global // world_size, h_local * world_size @@ -632,9 +632,9 @@ def _usp_output_all_to_all_varlen( assert x.ndim == 4, f"x must have 4 dimensions, got {x.ndim}" assert head_dim in (1, 2), f"head_dim must be 1 or 2, got {head_dim}" - assert ( - len(seq_lens) == world_size - ), f"seq_lens must have length {world_size}, got {len(seq_lens)}" + assert len(seq_lens) == world_size, ( + f"seq_lens must have length {world_size}, got {len(seq_lens)}" + ) rank = get_ulysses_parallel_rank() @@ -648,9 +648,9 @@ def _usp_output_all_to_all_varlen( # Shape transition: [b, s_global, h_local, d] -> [h_local, b, s_global, d] permute_order = (2, 0, 1, 3) - assert s_global == sum( - seq_lens - ), f"s_global ({s_global}) must equal sum(seq_lens) ({sum(seq_lens)})" + assert s_global == sum(seq_lens), ( + f"s_global ({s_global}) must equal sum(seq_lens) ({sum(seq_lens)})" + ) s_local = seq_lens[rank] diff --git a/python/sglang/multimodal_gen/runtime/layers/utils.py b/python/sglang/multimodal_gen/runtime/layers/utils.py index 2454eae78..b4e115073 100644 --- a/python/sglang/multimodal_gen/runtime/layers/utils.py +++ b/python/sglang/multimodal_gen/runtime/layers/utils.py @@ -147,7 +147,6 @@ class CustomOpWrapper: def real_impl(self) -> Callable: if self._impl is None: if not hasattr(torch.ops.sglang, self.op_name): - # NOTE(dark): if torch compile fail here, mark the decorator as eager # lazy registration does not work with torch compile direct_register_custom_op( @@ -240,15 +239,15 @@ def register_custom_op( """ extra_kwarg_keys = set(extra_kwargs.keys()) expected_kwarg_keys = set({"out_shape", "fake_impl"}) - assert ( - expected_kwarg_keys >= extra_kwarg_keys - ), f"Unexpected extra kwargs: {extra_kwarg_keys - expected_kwarg_keys}" + assert expected_kwarg_keys >= extra_kwarg_keys, ( + f"Unexpected extra kwargs: {extra_kwarg_keys - expected_kwarg_keys}" + ) has_out_shape = "out_shape" in extra_kwargs has_fake_impl = "fake_impl" in extra_kwargs - assert not ( - has_out_shape and has_fake_impl - ), "Only one of `out_shape` or `fake_impl` should be provided." + assert not (has_out_shape and has_fake_impl), ( + "Only one of `out_shape` or `fake_impl` should be provided." + ) # Assume inplace if neither out_shape nor fake_impl is provided if not (has_out_shape or has_fake_impl): extra_kwargs["out_shape"] = None diff --git a/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py b/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py index 64efd15cb..2f3efe73e 100644 --- a/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py +++ b/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py @@ -258,9 +258,9 @@ class TimestepEmbedder(nn.Module): t, self.frequency_embedding_size, self.max_period, dtype=self.freq_dtype ).to(self.mlp.fc_in.weight.dtype) if timestep_seq_len is not None: - assert ( - t_freq.shape[0] % timestep_seq_len == 0 - ), "timestep length is not divisible by timestep_seq_len" + assert t_freq.shape[0] % timestep_seq_len == 0, ( + "timestep length is not divisible by timestep_seq_len" + ) batch_size = t_freq.shape[0] // timestep_seq_len t_freq = t_freq.unflatten(0, (batch_size, timestep_seq_len)) # t_freq = t_freq.to(self.mlp.fc_in.weight.dtype) diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py index 9890b8bae..e1116d4d5 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py @@ -655,9 +655,9 @@ class ComponentLoader(ABC): ] expected_library = loader_cls.expected_library # Assert that the library matches what's expected for this component type - assert ( - transformers_or_diffusers == expected_library - ), f"{loader_type} must be loaded from {expected_library}, got {transformers_or_diffusers}" + assert transformers_or_diffusers == expected_library, ( + f"{loader_type} must be loaded from {expected_library}, got {transformers_or_diffusers}" + ) loader = loader_cls() loader.component_type = structural_component_name loader.component_architecture = component_architecture diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/pe_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/pe_loader.py index c51114fb9..756ec3236 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/pe_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/pe_loader.py @@ -92,7 +92,6 @@ class PEModelWrapper(nn.Module, LayerwiseOffloadableModuleMixin): class SGLangPEModelWrapper: - def __init__(self, model_url): self.model_url = model_url.rstrip("/") # Tokenizer is initialized separately during pipeline setup diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/scheduler_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/scheduler_loader.py index 44b5b4139..1bc0ae537 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/scheduler_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/scheduler_loader.py @@ -50,9 +50,9 @@ class SchedulerLoader(ComponentLoader): getattr(server_args.pipeline_config, "scheduler_class_override", None) or checkpoint_class_name ) - assert ( - class_name is not None - ), "Model config does not contain a _class_name attribute. Only diffusers format is supported." + assert class_name is not None, ( + "Model config does not contain a _class_name attribute. Only diffusers format is supported." + ) if checkpoint_class_name is not None and class_name != checkpoint_class_name: logger.info( diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py index 94b5d237f..511066434 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py @@ -31,9 +31,9 @@ class SoundTokenizerLoader(PlainStateDictComponentLoader): component_model_path, server_args, component_name ) class_name = config.pop("_class_name", None) or self.component_architecture - assert ( - class_name is not None - ), "Sound tokenizer class name must be available from component config." + assert class_name is not None, ( + "Sound tokenizer class name must be available from component config." + ) server_args.model_paths[component_name] = component_model_path diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py index 590c2c1f0..c0c44e5ae 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py @@ -882,8 +882,9 @@ class TextEncoderLoader(OnlineQuantizationComponentLoader): model_device = local_torch_device encoder_tp_group = get_folding_tp_group(model_config) - with use_tensor_parallel_group(encoder_tp_group), set_default_torch_dtype( - PRECISION_TO_TYPE[dtype] + with ( + use_tensor_parallel_group(encoder_tp_group), + set_default_torch_dtype(PRECISION_TO_TYPE[dtype]), ): with model_device, skip_init_modules(): architectures = getattr(model_config, "architectures", []) diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py index a07a33e8a..5171b5402 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py @@ -467,9 +467,9 @@ class VAELoader(WeightOverrideComponentLoader): ) class_name = config.pop("_class_name", None) - assert ( - class_name is not None - ), "Model config does not contain a _class_name attribute. Only diffusers format is supported." + assert class_name is not None, ( + "Model config does not contain a _class_name attribute. Only diffusers format is supported." + ) component_type = self.structural_component_type(component_name) if component_type in ("vae", "video_vae"): @@ -577,9 +577,9 @@ class VAELoader(WeightOverrideComponentLoader): vae_precision, ) - assert ( - len(safetensors_list) >= 1 - ), f"Found no safetensors files in {component_weights_path}" + assert len(safetensors_list) >= 1, ( + f"Found no safetensors files in {component_weights_path}" + ) if direct_gpu_weight_loading: _assign_direct_gpu_vae_state( vae, diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vl_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vl_encoder_loader.py index c684e0e2b..3b4159794 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vl_encoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vl_encoder_loader.py @@ -25,7 +25,6 @@ class VisionLanguageEncoderLoader(ComponentLoader): component_name: str = "vision_language_encoder", ) -> Any: if self.structural_component_type(component_name) == "vision_language_encoder": - if server_args.srt_encoder_url is not None: health_url = server_args.srt_encoder_url.rstrip("/") + "/health" try: diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py index b21ccaaae..e616ea4f8 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py @@ -29,9 +29,9 @@ class VocoderLoader(PlainStateDictComponentLoader): component_model_path, server_args, component_name ) class_name = config.pop("_class_name", None) or self.component_architecture - assert ( - class_name is not None - ), "Vocoder class name must be available from component config or pipeline config." + assert class_name is not None, ( + "Vocoder class name must be available from component config or pipeline config." + ) server_args.model_paths[component_name] = component_model_path diff --git a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py index e0c14e77e..5676dd198 100644 --- a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py @@ -770,8 +770,7 @@ def resolve_transformer_quant_load_spec( ) if server_args.nunchaku_config is not None: raise ValueError( - "Per-layer checkpoint quantization and Nunchaku are mutually " - "exclusive" + "Per-layer checkpoint quantization and Nunchaku are mutually exclusive" ) quant_config = checkpoint_quant_config elif getattr(model_cls, "handles_checkpoint_quantization", False): diff --git a/python/sglang/multimodal_gen/runtime/loader/weight_utils.py b/python/sglang/multimodal_gen/runtime/loader/weight_utils.py index a409abb6f..44f6309a5 100644 --- a/python/sglang/multimodal_gen/runtime/loader/weight_utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/weight_utils.py @@ -77,7 +77,6 @@ temp_dir = tempfile.gettempdir() class DisabledTqdm(tqdm): - def __init__(self, *args, **kwargs): kwargs["disable"] = True super().__init__(*args, **kwargs) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index dab6e2855..2139b876f 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -434,8 +434,8 @@ class GPUWorker(GPUWorkerPostTrainingMixin): ), log_reqs=[req], return_req=False, - save_output_paths=lambda output_batch, req=req: self._save_output_paths( - req, output_batch + save_output_paths=lambda output_batch, req=req: ( + self._save_output_paths(req, output_batch) ), error_context=f"grouped request {req.request_id}", execution_start_time=group_start_time, diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency.py index cc5ffafb6..0835d331a 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency.py @@ -82,8 +82,7 @@ def normalize_component_residency( for raw_selector, raw_mode in entries: if not isinstance(raw_selector, str) or not isinstance(raw_mode, str): raise ComponentResidencyError( - "Invalid component residency assignment: " - f"{raw_selector!r}={raw_mode!r}" + f"Invalid component residency assignment: {raw_selector!r}={raw_mode!r}" ) selector = raw_selector.strip().replace("-", "_").lower() mode = raw_mode.strip().replace("_", "-").lower() @@ -149,8 +148,7 @@ def resolve_diffusers_pipeline_offload( return None if LAYERWISE_OFFLOAD in assignments.values(): raise ComponentResidencyError( - "--component-residency layerwise-offload requires the native SGLang " - "backend" + "--component-residency layerwise-offload requires the native SGLang backend" ) pipeline_mode = assignments.get(LAYERWISE_OFFLOAD_ALL_COMPONENTS) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py index 48c6619cc..220ad3b96 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py @@ -729,9 +729,9 @@ class LayerwiseOffloadManager: # below swaps that same object's storage for a (1,) # placeholder, leaving the placeholder in the store. # Keep an independent tensor over the mapped storage. - self._mapped_cpu_weights[layer_idx][ - name - ] = local_weight.detach().view_as(local_weight) + self._mapped_cpu_weights[layer_idx][name] = ( + local_weight.detach().view_as(local_weight) + ) self._weight_metadata[layer_idx][name] = { "dtype": local_weight.dtype, "shape": tuple(local_weight.shape), diff --git a/python/sglang/multimodal_gen/runtime/models/bridges/mova_dual_tower.py b/python/sglang/multimodal_gen/runtime/models/bridges/mova_dual_tower.py index c94191b29..240f25a88 100644 --- a/python/sglang/multimodal_gen/runtime/models/bridges/mova_dual_tower.py +++ b/python/sglang/multimodal_gen/runtime/models/bridges/mova_dual_tower.py @@ -118,10 +118,12 @@ class PerFrameAttentionPooling(nn.Module): """ B, L, D = x.shape T, H, W = grid_size - assert ( - D == self.dim - ), f"Input dimension D={D} does not match module dim={self.dim}" - assert L == T * H * W, f"Flattened length L={L} does not match T*H*W={T*H*W}" + assert D == self.dim, ( + f"Input dimension D={D} does not match module dim={self.dim}" + ) + assert L == T * H * W, ( + f"Flattened length L={L} does not match T*H*W={T * H * W}" + ) S = H * W x_bt_s_d = x.view(B, T, S, D).contiguous().view(B * T, S, D) # [B*T, S, D] diff --git a/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py index 6f3fc4494..30710d6d8 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py @@ -413,9 +413,10 @@ class CausalWanTransformerBlock(nn.Module): norm_hidden_states, hidden_states = self.self_attn_residual_norm( hidden_states, attn_output, gate_msa, null_shift, null_scale ) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype - ), hidden_states.to(orig_dtype) + norm_hidden_states, hidden_states = ( + norm_hidden_states.to(orig_dtype), + hidden_states.to(orig_dtype), + ) # 2. Cross-attention attn_output = self.attn2( @@ -427,9 +428,10 @@ class CausalWanTransformerBlock(nn.Module): norm_hidden_states, hidden_states = self.cross_attn_residual_norm( hidden_states, attn_output, 1, c_shift_msa, c_scale_msa ) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype - ), hidden_states.to(orig_dtype) + norm_hidden_states, hidden_states = ( + norm_hidden_states.to(orig_dtype), + hidden_states.to(orig_dtype), + ) # 3. Feed-forward ff_output = self.ffn(norm_hidden_states) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py b/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py index 3b6ebd693..be94f0f7e 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py @@ -228,9 +228,9 @@ class ErnieImageSelfAttention(nn.Module): tp_size = get_tp_world_size() self.num_local_heads = num_heads // tp_size - assert ( - num_heads % tp_size == 0 - ), f"num_heads ({num_heads}) must be divisible by tp_size ({tp_size})" + assert num_heads % tp_size == 0, ( + f"num_heads ({num_heads}) must be divisible by tp_size ({tp_size})" + ) self.to_q = ColumnParallelLinear( hidden_size, diff --git a/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py b/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py index 97fd79a73..56dbaa506 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py @@ -531,9 +531,9 @@ class GlmImageAttention(torch.nn.Module): self.out_dim = out_dim if out_dim is not None else query_dim tp_size = get_tp_world_size() - assert ( - self.heads % tp_size == 0 - ), f"heads ({self.heads}) must be divisible by tp_size ({tp_size})" + assert self.heads % tp_size == 0, ( + f"heads ({self.heads}) must be divisible by tp_size ({tp_size})" + ) self.num_local_heads = self.heads // tp_size self.num_local_kv_heads = self.num_local_heads @@ -673,9 +673,9 @@ class GlmImageAttention(torch.nn.Module): # 4. Attention if attention_mask is not None: text_attn_mask = attention_mask - assert ( - text_attn_mask.dim() == 2 - ), "the shape of text_attn_mask should be (batch_size, text_seq_length)" + assert text_attn_mask.dim() == 2, ( + "the shape of text_attn_mask should be (batch_size, text_seq_length)" + ) hidden_states = self.attn( query, key, value, num_replicated_prefix=text_seq_length ) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py index acdb6c7f9..9e064ccbf 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py @@ -960,9 +960,9 @@ class HunyuanVideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixi teacache_params = forward_batch.teacache_params assert teacache_params is not None, "teacache_params is not initialized" - assert isinstance( - teacache_params, TeaCacheParams - ), "teacache_params is not a TeaCacheParams" + assert isinstance(teacache_params, TeaCacheParams), ( + "teacache_params is not a TeaCacheParams" + ) num_inference_steps = forward_batch.num_inference_steps teache_thresh = teacache_params.teacache_thresh @@ -1006,9 +1006,7 @@ class HunyuanVideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixi img_mod2_shift, img_mod2_scale, img_mod2_gate, - ) = ( - self.double_blocks[0].img_mod(vec_).chunk(6, dim=-1) - ) + ) = self.double_blocks[0].img_mod(vec_).chunk(6, dim=-1) normed_inp = self.double_blocks[0].img_attn_norm.norm(inp) modulated_inp = modulate(normed_inp, shift=img_mod1_shift, scale=img_mod1_scale) if self.cnt == 0 or self.cnt == num_inference_steps - 1: @@ -1023,9 +1021,9 @@ class HunyuanVideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixi 9.61237896e-02, ] rescale_func = np.poly1d(coefficients) - assert ( - self.previous_modulated_input is not None - ), "previous_modulated_input is not initialized" + assert self.previous_modulated_input is not None, ( + "previous_modulated_input is not initialized" + ) self.accumulated_rel_l1_distance += rescale_func( ( (modulated_inp - self.previous_modulated_input).abs().mean() diff --git a/python/sglang/multimodal_gen/runtime/models/dits/krea2.py b/python/sglang/multimodal_gen/runtime/models/dits/krea2.py index eeff31c4e..2c0dc91a1 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/krea2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/krea2.py @@ -264,9 +264,9 @@ class Attention(nn.Module): # Parameter names match the released checkpoint (to_q/to_k/to_v/to_gate, # norm_q/norm_k, to_out.0) so the checkpoint loads with an identity mapping. tp = get_tp_world_size() - assert ( - self.heads % tp == 0 and self.kvheads % tp == 0 - ), f"heads={self.heads}, kvheads={self.kvheads} must be divisible by tp={tp}" + assert self.heads % tp == 0 and self.kvheads % tp == 0, ( + f"heads={self.heads}, kvheads={self.kvheads} must be divisible by tp={tp}" + ) self.local_heads = self.heads // tp self.local_kvheads = self.kvheads // tp diff --git a/python/sglang/multimodal_gen/runtime/models/dits/lingbot_video_moe.py b/python/sglang/multimodal_gen/runtime/models/dits/lingbot_video_moe.py index b01a48ab7..012cd4781 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/lingbot_video_moe.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/lingbot_video_moe.py @@ -427,9 +427,9 @@ class LingBotVideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixi hidden_size = config.hidden_size num_attention_heads = config.num_attention_heads head_dim = hidden_size // num_attention_heads - assert head_dim == sum( - config.axes_dims - ), f"head_dim {head_dim} != sum(axes_dims) {sum(config.axes_dims)}" + assert head_dim == sum(config.axes_dims), ( + f"head_dim {head_dim} != sum(axes_dims) {sum(config.axes_dims)}" + ) mlp_only_layers = tuple(config.mlp_only_layers) self.hidden_size = hidden_size diff --git a/python/sglang/multimodal_gen/runtime/models/dits/longcat_image.py b/python/sglang/multimodal_gen/runtime/models/dits/longcat_image.py index eb85876cd..986b793c7 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/longcat_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/longcat_image.py @@ -268,9 +268,9 @@ class _LongCatJointAttention(nn.Module): super().__init__() tp_size = get_tp_world_size() self.num_local_heads = num_attention_heads // tp_size - assert ( - num_attention_heads % tp_size == 0 - ), f"num_attention_heads ({num_attention_heads}) must be divisible by tp_size ({tp_size})" + assert num_attention_heads % tp_size == 0, ( + f"num_attention_heads ({num_attention_heads}) must be divisible by tp_size ({tp_size})" + ) self.head_dim = attention_head_dim inner_dim = num_attention_heads * attention_head_dim @@ -429,9 +429,9 @@ class _LongCatSingleAttention(nn.Module): super().__init__() tp_size = get_tp_world_size() self.num_local_heads = num_attention_heads // tp_size - assert ( - num_attention_heads % tp_size == 0 - ), f"num_attention_heads ({num_attention_heads}) must be divisible by tp_size ({tp_size})" + assert num_attention_heads % tp_size == 0, ( + f"num_attention_heads ({num_attention_heads}) must be divisible by tp_size ({tp_size})" + ) self.head_dim = attention_head_dim inner_dim = num_attention_heads * attention_head_dim diff --git a/python/sglang/multimodal_gen/runtime/models/dits/longlive2.py b/python/sglang/multimodal_gen/runtime/models/dits/longlive2.py index 19b56c138..c47acabcd 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/longlive2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/longlive2.py @@ -128,9 +128,10 @@ class LongLive2CausalWanTransformerBlock(CausalWanTransformerBlock): norm_hidden_states, hidden_states = self.self_attn_residual_norm( hidden_states, attn_output, gate_msa, null_shift, null_scale ) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype - ), hidden_states.to(orig_dtype) + norm_hidden_states, hidden_states = ( + norm_hidden_states.to(orig_dtype), + hidden_states.to(orig_dtype), + ) attn_output = self._cross_attn_with_cache( norm_hidden_states, @@ -140,9 +141,10 @@ class LongLive2CausalWanTransformerBlock(CausalWanTransformerBlock): norm_hidden_states, hidden_states = self.cross_attn_residual_norm( hidden_states, attn_output, 1, c_shift_msa, c_scale_msa ) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype - ), hidden_states.to(orig_dtype) + norm_hidden_states, hidden_states = ( + norm_hidden_states.to(orig_dtype), + hidden_states.to(orig_dtype), + ) ff_output = self.ffn(norm_hidden_states) hidden_states = self.mlp_residual(ff_output, c_gate_msa, hidden_states) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py index 50f572b3d..a72bcd81e 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py @@ -129,9 +129,12 @@ def _diffusers_h3_checkpoint( continue merge_dim = 1 if target_name.endswith((".qweight", ".qzeros", ".scales")) else 0 - yield target_name, torch.cat( - [pending[target_name][index] for index in range(merge_count)], - dim=merge_dim, + yield ( + target_name, + torch.cat( + [pending[target_name][index] for index in range(merge_count)], + dim=merge_dim, + ), ) del pending[target_name] @@ -1945,8 +1948,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): ): if value % tp_size: raise ValueError( - f"MiniMax H3 {name}={value} must be divisible by " - f"TP size {tp_size}." + f"MiniMax H3 {name}={value} must be divisible by TP size {tp_size}." ) @staticmethod @@ -2005,8 +2007,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): adaln_cache_path is not None or adaln_weight_files is not None ): raise ValueError( - "MiniMax H3 pruned curve checkpoints cannot use a separate " - "AdaLN cache" + "MiniMax H3 pruned curve checkpoints cannot use a separate AdaLN cache" ) self._adaln_precomputed = ( adaln_cache_path is not None or adaln_weight_files is not None diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index 944b67161..c4167020f 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py @@ -762,9 +762,9 @@ class QwenImageCrossAttention(nn.Module): self.inner_kv_dim = self.inner_dim tp_size = get_tp_world_size() - assert ( - self.num_heads % tp_size == 0 - ), f"num_heads ({self.num_heads}) must be divisible by tp_size ({tp_size})" + assert self.num_heads % tp_size == 0, ( + f"num_heads ({self.num_heads}) must be divisible by tp_size ({tp_size})" + ) self.local_num_heads = self.num_heads // tp_size self._unquantized_added_qkv_is_packed = False diff --git a/python/sglang/multimodal_gen/runtime/models/dits/sana.py b/python/sglang/multimodal_gen/runtime/models/dits/sana.py index 2af9923b2..f1e21cff8 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/sana.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/sana.py @@ -494,7 +494,6 @@ class SanaTransformerBlock(nn.Module): class SanaTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin): - _fsdp_shard_conditions = [ lambda n, m: isinstance(m, SanaTransformerBlock), ] diff --git a/python/sglang/multimodal_gen/runtime/models/dits/sana_wm_components.py b/python/sglang/multimodal_gen/runtime/models/dits/sana_wm_components.py index 1d9b1cc30..cad7a50c4 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/sana_wm_components.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/sana_wm_components.py @@ -1871,9 +1871,9 @@ class BidirectionalGDNUCPESinglePathLiteLA(nn.Module): ) -> None: super().__init__() out_dim = heads * head_dim - assert ( - out_dim == in_dim - ), f"in_dim ({in_dim}) must equal heads*head_dim ({out_dim})" + assert out_dim == in_dim, ( + f"in_dim ({in_dim}) must equal heads*head_dim ({out_dim})" + ) self.in_dim = in_dim self.out_dim = out_dim self.heads = heads @@ -2219,9 +2219,7 @@ class BidirectionalGDNUCPESinglePathLiteLA(nn.Module): if beta.ndim == 3 and beta.shape != (B, heads, T): return f"requires beta shape {(B, heads, T)}, got {tuple(beta.shape)}" if beta.ndim == 4 and beta.shape != (B, heads, T, S): - return ( - f"requires beta shape {(B, heads, T, S)}, " f"got {tuple(beta.shape)}" - ) + return f"requires beta shape {(B, heads, T, S)}, got {tuple(beta.shape)}" if decay.shape != (B, heads, T): return f"requires decay shape {(B, heads, T)}, got {tuple(decay.shape)}" if head_dim > 128: @@ -2252,8 +2250,7 @@ class BidirectionalGDNUCPESinglePathLiteLA(nn.Module): if precheck_reason is not None: if self.gdn_backend == "triton": raise RuntimeError( - "SANA-WM Triton camera GDN backend unavailable: " - f"{precheck_reason}" + f"SANA-WM Triton camera GDN backend unavailable: {precheck_reason}" ) return None diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index b80b22702..898add578 100755 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py @@ -679,9 +679,10 @@ class WanTransformerBlock(nn.Module): query = q_sbhd.view(query_shape) key = k_sbhd.view(key_shape) else: - query, key = _apply_rotary_emb( - query, cos, sin, is_neox_style=False - ), _apply_rotary_emb(key, cos, sin, is_neox_style=False) + query, key = ( + _apply_rotary_emb(query, cos, sin, is_neox_style=False), + _apply_rotary_emb(key, cos, sin, is_neox_style=False), + ) attn_output = self.attn1(query, key, value) attn_output = attn_output.flatten(2) attn_output, _ = self.to_out(attn_output) @@ -693,9 +694,10 @@ class WanTransformerBlock(nn.Module): norm_hidden_states, hidden_states = self.self_attn_residual_norm( hidden_states, attn_output, gate_msa, null_shift, null_scale ) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype - ), hidden_states.to(orig_dtype) + norm_hidden_states, hidden_states = ( + norm_hidden_states.to(orig_dtype), + hidden_states.to(orig_dtype), + ) # 2. Cross-attention attn_output = self.attn2( @@ -704,9 +706,10 @@ class WanTransformerBlock(nn.Module): norm_hidden_states, hidden_states = self.cross_attn_residual_norm( hidden_states, attn_output, 1, c_shift_msa, c_scale_msa ) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype - ), hidden_states.to(orig_dtype) + norm_hidden_states, hidden_states = ( + norm_hidden_states.to(orig_dtype), + hidden_states.to(orig_dtype), + ) # 3. Feed-forward ff_output = self.ffn(norm_hidden_states) @@ -930,9 +933,10 @@ class WanTransformerBlock_VSA(nn.Module): query = q_sbhd.view(query_shape) key = k_sbhd.view(key_shape) else: - query, key = _apply_rotary_emb( - query, cos, sin, is_neox_style=False - ), _apply_rotary_emb(key, cos, sin, is_neox_style=False) + query, key = ( + _apply_rotary_emb(query, cos, sin, is_neox_style=False), + _apply_rotary_emb(key, cos, sin, is_neox_style=False), + ) attn_output = self.attn1(query, key, value, gate_compress=gate_compress) attn_output = attn_output.flatten(2) @@ -943,9 +947,10 @@ class WanTransformerBlock_VSA(nn.Module): norm_hidden_states, hidden_states = self.self_attn_residual_norm( hidden_states, attn_output, gate_msa, null_shift, null_scale ) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype - ), hidden_states.to(orig_dtype) + norm_hidden_states, hidden_states = ( + norm_hidden_states.to(orig_dtype), + hidden_states.to(orig_dtype), + ) # 2. Cross-attention attn_output = self.attn2( @@ -954,9 +959,10 @@ class WanTransformerBlock_VSA(nn.Module): norm_hidden_states, hidden_states = self.cross_attn_residual_norm( hidden_states, attn_output, 1, c_shift_msa, c_scale_msa ) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype - ), hidden_states.to(orig_dtype) + norm_hidden_states, hidden_states = ( + norm_hidden_states.to(orig_dtype), + hidden_states.to(orig_dtype), + ) # 3. Feed-forward ff_output = self.ffn(norm_hidden_states) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py index a3fd410d9..3519b2583 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py @@ -255,12 +255,12 @@ class ZImageAttention(nn.Module): self.enable_zimage_qk_fusion = quant_config is None tp_size = get_tp_world_size() - assert ( - num_heads % tp_size == 0 - ), f"num_heads {num_heads} must be divisible by tp world size {tp_size}" - assert ( - num_kv_heads % tp_size == 0 - ), f"num_kv_heads {num_kv_heads} must be divisible by tp world size {tp_size}" + assert num_heads % tp_size == 0, ( + f"num_heads {num_heads} must be divisible by tp world size {tp_size}" + ) + assert num_kv_heads % tp_size == 0, ( + f"num_kv_heads {num_kv_heads} must be divisible by tp world size {tp_size}" + ) self.local_num_heads = num_heads // tp_size self.local_num_kv_heads = num_kv_heads // tp_size @@ -704,9 +704,9 @@ class RopeEmbedder: self.theta = theta self.axes_dims = axes_dims self.axes_lens = axes_lens - assert len(axes_dims) == len( - axes_lens - ), "axes_dims and axes_lens must have the same length" + assert len(axes_dims) == len(axes_lens), ( + "axes_dims and axes_lens must have the same length" + ) self.cos_cached = None self.sin_cached = None diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/clip.py b/python/sglang/multimodal_gen/runtime/models/encoders/clip.py index 75df48a7f..ce42af1d8 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/clip.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/clip.py @@ -36,7 +36,6 @@ def _srt_clip_param_name(name: str) -> str: class CLIPTextTransformer(nn.Module): - def __init__( self, config: CLIPTextConfig, @@ -145,7 +144,6 @@ class CLIPTextTransformer(nn.Module): class CLIPTextModel(TextEncoder): - def __init__( self, config: CLIPTextConfig, @@ -254,7 +252,6 @@ class CLIPTextModelWithProjection(CLIPTextModel): class CLIPVisionTransformer(nn.Module): - def __init__( self, config: CLIPVisionConfig, diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/hunyuan3d.py b/python/sglang/multimodal_gen/runtime/models/encoders/hunyuan3d.py index dc355f46e..f5548aa54 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/hunyuan3d.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/hunyuan3d.py @@ -209,7 +209,7 @@ def build_image_encoder(config): elif config["type"] == "DinoImageEncoderMV": return DinoImageEncoderMV(**config["kwargs"]) else: - raise ValueError(f'Unknown image encoder type: {config["type"]}') + raise ValueError(f"Unknown image encoder type: {config['type']}") class DualImageEncoder(nn.Module, LayerwiseOffloadableModuleMixin): diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/llama.py b/python/sglang/multimodal_gen/runtime/models/encoders/llama.py index 27ad1a518..5debd9735 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/llama.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/llama.py @@ -58,7 +58,6 @@ from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder class LlamaMLP(nn.Module): - def __init__( self, hidden_size: int, @@ -86,8 +85,7 @@ class LlamaMLP(nn.Module): ) if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() @@ -99,7 +97,6 @@ class LlamaMLP(nn.Module): class LlamaAttention(nn.Module): - def __init__( self, config: LlamaConfig, @@ -218,7 +215,6 @@ class LlamaAttention(nn.Module): class LlamaDecoderLayer(nn.Module): - def __init__( self, config: LlamaConfig, @@ -296,7 +292,6 @@ class LlamaDecoderLayer(nn.Module): class LlamaModel(TextEncoder): - def __init__( self, config: LlamaConfig, diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py b/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py index ab8b74f0e..c995291ca 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py @@ -1329,7 +1329,9 @@ class Qwen2_5_VLForConditionalGeneration(TextEncoder): eos_token_ids = ( [] if eos_token_id is None - else [eos_token_id] if isinstance(eos_token_id, int) else list(eos_token_id) + else [eos_token_id] + if isinstance(eos_token_id, int) + else list(eos_token_id) ) if pad_token_id is None: raise ValueError("pad_token_id must be set for generation") diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/qwen_vl_rope.py b/python/sglang/multimodal_gen/runtime/models/encoders/qwen_vl_rope.py index 85077a925..4e96370fd 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/qwen_vl_rope.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/qwen_vl_rope.py @@ -38,8 +38,7 @@ def apply_qwen_vl_text_rope( """Apply three-axis MRoPE to batched attention tensors.""" if query.ndim != 4 or key.ndim != 4: raise ValueError( - "Qwen-VL query and key must have shape " - "[batch, heads, sequence, head_dim]" + "Qwen-VL query and key must have shape [batch, heads, sequence, head_dim]" ) if position_ids.ndim != 3 or position_ids.shape[0] != 3: raise ValueError( diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/t5.py b/python/sglang/multimodal_gen/runtime/models/encoders/t5.py index 3834fb69e..7eb86621e 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/t5.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/t5.py @@ -72,7 +72,6 @@ class AttentionMetadata: class T5DenseActDense(nn.Module): - def __init__( self, config: T5Config, quant_config: QuantizationConfig | None = None ): @@ -98,7 +97,6 @@ class T5DenseActDense(nn.Module): class T5DenseGatedActDense(nn.Module): - def __init__( self, config: T5Config, quant_config: QuantizationConfig | None = None ): @@ -138,7 +136,6 @@ class T5DenseGatedActDense(nn.Module): class T5LayerFF(nn.Module): - def __init__( self, config: T5Config, quant_config: QuantizationConfig | None = None ): @@ -161,7 +158,6 @@ class T5LayerFF(nn.Module): # T5 has attn_bias and does not use softmax scaling class T5MultiHeadAttention(nn.Module): - def __init__(self) -> None: super().__init__() @@ -178,7 +174,6 @@ class T5MultiHeadAttention(nn.Module): class T5Attention(nn.Module): - def __init__( self, config: T5Config, @@ -378,7 +373,6 @@ class T5Attention(nn.Module): class T5LayerSelfAttention(nn.Module): - def __init__( self, config, @@ -416,7 +410,6 @@ class T5LayerSelfAttention(nn.Module): class T5LayerCrossAttention(nn.Module): - def __init__( self, config, quant_config: QuantizationConfig | None = None, prefix: str = "" ): @@ -445,7 +438,6 @@ class T5LayerCrossAttention(nn.Module): class T5Block(nn.Module): - def __init__( self, config: T5Config, @@ -505,7 +497,6 @@ class T5Block(nn.Module): class T5Stack(nn.Module): - def __init__( self, config: T5Config, diff --git a/python/sglang/multimodal_gen/runtime/models/registry.py b/python/sglang/multimodal_gen/runtime/models/registry.py index cb44d9fb6..0330acc75 100644 --- a/python/sglang/multimodal_gen/runtime/models/registry.py +++ b/python/sglang/multimodal_gen/runtime/models/registry.py @@ -177,7 +177,6 @@ class _ModelInfo: class _BaseRegisteredModel(ABC): - @abstractmethod def inspect_model_cls(self) -> _ModelInfo: raise NotImplementedError @@ -231,7 +230,7 @@ def _run_in_subprocess(fn: Callable[[], _T]) -> _T: except Exception as e: # wrap raised exception to provide more information raise RuntimeError( - f"Error raised in subprocess:\n" f"{returned.stderr.decode()}" + f"Error raised in subprocess:\n{returned.stderr.decode()}" ) from e with open(output_filepath, "rb") as f: diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_unipc_multistep.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_unipc_multistep.py index 5d9ea035d..c78d3ddf7 100644 --- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_unipc_multistep.py +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_unipc_multistep.py @@ -194,9 +194,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): assert num_inference_steps is not None sigmas = np.linspace( self.sigma_max, self.sigma_min, num_inference_steps + 1 - ).copy()[ - :-1 - ] # pyright: ignore + ).copy()[:-1] # pyright: ignore if self.config.use_dynamic_shifting: assert mu is not None @@ -217,9 +215,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): ) timesteps = sigmas * self.config.num_train_timesteps - sigmas = np.concatenate([sigmas, [sigma_last]]).astype( - np.float32 - ) # pyright: ignore + sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) # pyright: ignore self.sigmas = torch.from_numpy(sigmas).to(device=device) self.timesteps = torch.from_numpy(timesteps).to( @@ -476,18 +472,14 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): if self.predict_x0: x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 if D1s is not None: - pred_res = torch.einsum( - "k,bkc...->bc...", rhos_p, D1s - ) # pyright: ignore + pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) # pyright: ignore else: pred_res = 0 x_t = x_t_ - alpha_t * B_h * pred_res else: x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 if D1s is not None: - pred_res = torch.einsum( - "k,bkc...->bc...", rhos_p, D1s - ) # pyright: ignore + pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) # pyright: ignore else: pred_res = 0 x_t = x_t_ - sigma_t * B_h * pred_res diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_minimax_h3_euler_ancestral.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_minimax_h3_euler_ancestral.py index 6c067b0a7..81629fe14 100644 --- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_minimax_h3_euler_ancestral.py +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_minimax_h3_euler_ancestral.py @@ -138,8 +138,7 @@ class MiniMaxH3EulerAncestralEta0SchedulerAdapter: def __init__(self, **config: Any) -> None: if config: raise ValueError( - f"{type(self).__name__} does not accept config fields: " - f"{sorted(config)}" + f"{type(self).__name__} does not accept config fields: {sorted(config)}" ) def set_shift(self, _flow_shift: float) -> None: diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/common.py b/python/sglang/multimodal_gen/runtime/models/vaes/common.py index cbe9b6af1..28b032153 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/common.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/common.py @@ -749,7 +749,6 @@ class ParallelTiledVAE(ABC, nn.Module, LayerwiseOffloadableModuleMixin): # adapted from https://github.com/huggingface/diffusers/blob/e7ffeae0a191f710881d1fbde00cd6ff025e81f2/src/diffusers/models/autoencoders/vae.py#L691 class DiagonalGaussianDistribution: - def __init__(self, parameters: torch.Tensor, deterministic: bool = False): self.parameters = parameters self.mean, self.logvar = torch.chunk(parameters, 2, dim=1) diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/hunyuan3d_vae.py b/python/sglang/multimodal_gen/runtime/models/vaes/hunyuan3d_vae.py index 1de3203c4..8ad0a992a 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/hunyuan3d_vae.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/hunyuan3d_vae.py @@ -475,7 +475,6 @@ class Transformer(nn.Module): class CrossAttentionDecoder(nn.Module): - def __init__( self, *, diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/hunyuanvae.py b/python/sglang/multimodal_gen/runtime/models/vaes/hunyuanvae.py index 5067a890e..93d424f71 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/hunyuanvae.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/hunyuanvae.py @@ -100,7 +100,6 @@ def _apply_group_norm_silu( class HunyuanVAEAttention(nn.Module): - def __init__( self, in_channels, heads, dim_head, eps, norm_num_groups, bias ) -> None: @@ -167,7 +166,6 @@ class HunyuanVAEAttention(nn.Module): class HunyuanVideoCausalConv3d(nn.Module): - def __init__( self, in_channels: int, @@ -237,7 +235,6 @@ class HunyuanVideoCausalConv3d(nn.Module): class HunyuanVideoUpsampleCausal3D(nn.Module): - def __init__( self, in_channels: int, @@ -286,7 +283,6 @@ class HunyuanVideoUpsampleCausal3D(nn.Module): class HunyuanVideoDownsampleCausal3D(nn.Module): - def __init__( self, channels: int, @@ -309,7 +305,6 @@ class HunyuanVideoDownsampleCausal3D(nn.Module): class HunyuanVideoResnetBlockCausal3D(nn.Module): - def __init__( self, in_channels: int, @@ -361,7 +356,6 @@ class HunyuanVideoResnetBlockCausal3D(nn.Module): class HunyuanVideoMidBlock3D(nn.Module): - def __init__( self, in_channels: int, @@ -473,7 +467,6 @@ class HunyuanVideoMidBlock3D(nn.Module): class HunyuanVideoDownBlock3D(nn.Module): - def __init__( self, in_channels: int, @@ -537,7 +530,6 @@ class HunyuanVideoDownBlock3D(nn.Module): class HunyuanVideoUpBlock3D(nn.Module): - def __init__( self, in_channels: int, diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/processor.py b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/processor.py index 3fa00d9a7..9de0a673a 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/processor.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/processor.py @@ -54,7 +54,6 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name class VAEProcessor: - def __init__( self, *, diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/vae_vit.py b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/vae_vit.py index 39d37376e..9bcbade80 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/vae_vit.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/vae_vit.py @@ -241,8 +241,7 @@ class ViT3DDecoder(ViTBase): if dtype not in (torch.float16, torch.bfloat16): raise ValueError( - "MiniMax H3 decoder autocast weights require fp16 or bf16, " - f"got {dtype}" + f"MiniMax H3 decoder autocast weights require fp16 or bf16, got {dtype}" ) if self._autocast_linear_dtype == dtype: return 0 diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py b/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py index f631bdccd..c18bdf181 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py @@ -777,7 +777,6 @@ class WanMidBlock(nn.Module): class WanResidualDownBlock(nn.Module): - def __init__( self, in_dim, diff --git a/python/sglang/multimodal_gen/runtime/pipelines/comfyui_flux_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/comfyui_flux_pipeline.py index cd3c9f39f..db7305d6a 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/comfyui_flux_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/comfyui_flux_pipeline.py @@ -365,10 +365,22 @@ class ComfyUIFluxPipeline(LoRAPipeline, ComposedPipelineBase): k_tensor = qkv_tensor[hidden_size : 2 * hidden_size, :] v_tensor = qkv_tensor[2 * hidden_size : 3 * hidden_size, :] - yield f"single_transformer_blocks.{block_idx}.attn.to_q.{param_type}", q_tensor - yield f"single_transformer_blocks.{block_idx}.attn.to_k.{param_type}", k_tensor - yield f"single_transformer_blocks.{block_idx}.attn.to_v.{param_type}", v_tensor - yield f"single_transformer_blocks.{block_idx}.proj_mlp.{param_type}", mlp_tensor + yield ( + f"single_transformer_blocks.{block_idx}.attn.to_q.{param_type}", + q_tensor, + ) + yield ( + f"single_transformer_blocks.{block_idx}.attn.to_k.{param_type}", + k_tensor, + ) + yield ( + f"single_transformer_blocks.{block_idx}.attn.to_v.{param_type}", + v_tensor, + ) + yield ( + f"single_transformer_blocks.{block_idx}.proj_mlp.{param_type}", + mlp_tensor, + ) elif name == "final_layer.adaLN_modulation.1.weight": # ComfyUI: output order is [shift, scale] # AdaLayerNormContinuous: expects [scale, shift] diff --git a/python/sglang/multimodal_gen/runtime/pipelines/ernie_image.py b/python/sglang/multimodal_gen/runtime/pipelines/ernie_image.py index 207653035..63f9bd533 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/ernie_image.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/ernie_image.py @@ -27,7 +27,6 @@ logger = init_logger(__name__) class ErnieImagePipeline(LoRAPipeline, ComposedPipelineBase): - pipeline_name = "ErnieImagePipeline" _required_config_modules = [ @@ -153,9 +152,9 @@ class ErnieImagePipeline(LoRAPipeline, ComposedPipelineBase): hasattr(pipeline_config, "text_encoder_extra_args") and pipeline_config.text_encoder_extra_args ): - pipeline_config.text_encoder_extra_args[0][ - "max_length" - ] = text_model_max_length + pipeline_config.text_encoder_extra_args[0]["max_length"] = ( + text_model_max_length + ) logger.info( "Set text encoder model_max_length=%d from tokenizer/tokenizer_config.json", text_model_max_length, diff --git a/python/sglang/multimodal_gen/runtime/pipelines/hunyuan_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/hunyuan_pipeline.py index ef7fd15ce..75ea74255 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/hunyuan_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/hunyuan_pipeline.py @@ -19,7 +19,6 @@ from sglang.multimodal_gen.runtime.server_args import ServerArgs class HunyuanVideoPipeline(ComposedPipelineBase): - pipeline_name = "HunyuanVideoPipeline" _required_config_modules = [ diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py index 72266cfff..8d601c781 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py @@ -480,9 +480,9 @@ class ComposedPipelineBase(ABC): self._validate_direct_gpu_component_selection(model_index, server_args) # some sanity checks - assert ( - len(model_index) > 1 - ), "model_index.json must contain at least one pipeline module" + assert len(model_index) > 1, ( + "model_index.json must contain at least one pipeline module" + ) # In disagg mode, read HF config for skipped components (e.g., VAE) # so that update_model_arch + post_init can derive pipeline_config. diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/lora/format_adapter.py b/python/sglang/multimodal_gen/runtime/pipelines_core/lora/format_adapter.py index bc610cce1..557756d05 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/lora/format_adapter.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/lora/format_adapter.py @@ -337,8 +337,7 @@ def _convert_kohya_flux_via_diffusers( no_converter_warning="[LoRAFormatAdapter] No Kohya FLUX converter found.", success_info="[LoRAFormatAdapter] Converted Kohya FLUX LoRA using {name}", all_failed_warning=( - "[LoRAFormatAdapter] Kohya FLUX conversion failed; " - "last error: {last_err}" + "[LoRAFormatAdapter] Kohya FLUX conversion failed; last error: {last_err}" ), ) @@ -485,8 +484,7 @@ def _convert_ai_toolkit_flux_lora( sample = _sample_keys(final_out.keys(), 20) log.info( - "[LoRAFormatAdapter] after AI_TOOLKIT_FLUX conversion, " - "sample keys (<=20): %s", + "[LoRAFormatAdapter] after AI_TOOLKIT_FLUX conversion, sample keys (<=20): %s", ", ".join(sample), ) return final_out diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/lora/peft_adapter.py b/python/sglang/multimodal_gen/runtime/pipelines_core/lora/peft_adapter.py index a6aed9cfa..68ab86c5d 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/lora/peft_adapter.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/lora/peft_adapter.py @@ -264,7 +264,9 @@ def scale_fused_sections( alpha = float( alpha_parts[index].item() if index in alpha_parts - else default_alpha if default_alpha is not None else rank + else default_alpha + if default_alpha is not None + else rank ) scale = alpha / rank weight = b_parts[index] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py index 107c28e43..70d43d582 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py @@ -215,8 +215,7 @@ class PipelineStage(StageDedupMixin, ABC): return replace(use, target_dtype=target_dtype) return use raise ValueError( - f"{self.__class__.__name__} did not declare component use: " - f"{component_name}" + f"{self.__class__.__name__} did not declare component use: {component_name}" ) @contextmanager diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py index ca65b357d..ef91dfbe9 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py @@ -344,9 +344,9 @@ class DecodingStage(PipelineStage): # decode trajectory latents if needed if batch.return_trajectory_decoded: - assert ( - batch.trajectory_latents is not None - ), "batch should have trajectory latents" + assert batch.trajectory_latents is not None, ( + "batch should have trajectory latents" + ) # 1. Batch trajectory decoding to improve GPU utilization # batch.trajectory_latents is [batch_size, timesteps, channels, frames, height, width] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 03cc84fa0..69e61bc75 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -1521,9 +1521,9 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin): # 1. Prepare latent inputs in the model's compute dtype. latent_model_input = ctx.latents.to(ctx.target_dtype) if batch.image_latent is not None: - assert ( - not server_args.pipeline_config.task_type == ModelTaskType.TI2V - ), "image latents should not be provided for TI2V task" + assert not server_args.pipeline_config.task_type == ModelTaskType.TI2V, ( + "image latents should not be provided for TI2V task" + ) latent_model_input = torch.cat( [latent_model_input, batch.image_latent], dim=1 ).to(ctx.target_dtype) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py index 808c55559..5cefe0a7f 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py @@ -134,9 +134,9 @@ class DmdDenoisingStage(DenoisingStage): ], dim=2, ).to(target_dtype) - assert not torch.isnan( - latent_model_input - ).any(), "latent_model_input contains nan" + assert not torch.isnan(latent_model_input).any(), ( + "latent_model_input contains nan" + ) # Prepare inputs for transformer t_expand = t.repeat(latent_model_input.shape[0]) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py index 56dd05431..4b7e29e68 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py @@ -359,7 +359,9 @@ class InputValidationStage(PipelineStage): neg_prompt_state = ( "not set" if batch.negative_prompt is None - else "empty" if batch.negative_prompt == "" else "set" + else "empty" + if batch.negative_prompt == "" + else "set" ) raise ValueError( f"Server was launched with --enable-cfg-parallel but this " @@ -446,8 +448,10 @@ class InputValidationStage(PipelineStage): result.add_check( "prompt_or_embeds", None, - lambda _: V.string_or_list_strings(batch.prompt) - or V.list_not_empty(batch.prompt_embeds), + lambda _: ( + V.string_or_list_strings(batch.prompt) + or V.list_not_empty(batch.prompt_embeds) + ), ) if server_args.pipeline_config.task_type != ModelTaskType.I2M: diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py index 6b7f35cb6..ce07e10f8 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py @@ -342,9 +342,7 @@ class LatentPreparationStage(PipelineStage): server_args.pipeline_config.vae_config.use_temporal_scaling_frames ) if use_temporal_scaling_frames: - temporal_scale_factor = ( - server_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio - ) + temporal_scale_factor = server_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio latent_num_frames = (video_length - 1) // temporal_scale_factor + 1 return int(latent_num_frames) @@ -354,8 +352,10 @@ class LatentPreparationStage(PipelineStage): result.add_check( "prompt_or_embeds", None, - lambda _: V.string_or_list_strings(batch.prompt) - or V.list_not_empty(batch.prompt_embeds), + lambda _: ( + V.string_or_list_strings(batch.prompt) + or V.list_not_empty(batch.prompt_embeds) + ), ) result.add_check("prompt_embeds", batch.prompt_embeds, V.list_of_tensors) result.add_check( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py index 220916090..02e382703 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py @@ -152,8 +152,7 @@ def _resize_center_crop_uint8_cthw( """Resize and center-crop ``uint8 [3, T, H, W]`` transfer frames.""" if frames.ndim != 4 or frames.shape[0] != 3: raise ValueError( - "Transfer frames must have shape [3, T, H, W], got " - f"{tuple(frames.shape)}" + f"Transfer frames must have shape [3, T, H, W], got {tuple(frames.shape)}" ) orig_h, orig_w = int(frames.shape[2]), int(frames.shape[3]) scale = max(width / orig_w, height / orig_h) @@ -178,8 +177,7 @@ def _pad_transfer_frames(video: torch.Tensor, target_frames: int) -> torch.Tenso """Pad ``[1, 3, T, H, W]`` with reflected temporal content.""" if video.ndim != 5 or video.shape[0] != 1 or video.shape[1] != 3: raise ValueError( - "Transfer video must have shape [1, 3, T, H, W], got " - f"{tuple(video.shape)}" + f"Transfer video must have shape [1, 3, T, H, W], got {tuple(video.shape)}" ) if target_frames <= 0: raise ValueError("Transfer target frame count must be positive") @@ -242,8 +240,7 @@ class Cosmos3ImagePreprocessStage(PipelineStage): stride = frames_per_chunk - conditional_frames if stride <= 0: raise ValueError( - "num_conditional_frames must be smaller than " - "num_video_frames_per_chunk" + "num_conditional_frames must be smaller than num_video_frames_per_chunk" ) remaining = total_frames - frames_per_chunk return 1 + math.ceil(remaining / stride), stride diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ernie_image_pe.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ernie_image_pe.py index 7eab9e803..a713d569d 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ernie_image_pe.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ernie_image_pe.py @@ -19,7 +19,6 @@ logger = init_logger(__name__) class PromptEnhancementStage(PipelineStage): - def __init__(self, pe_model, pe_tokenizer): super().__init__() self.pe_model = pe_model diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py index e6c60bd96..162767f5d 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py @@ -2719,8 +2719,8 @@ class LTX2DenoisingStage(DenoisingStage): def _get_negative_prompt_embeds_validator(self, batch: Req): """Allow either tensor or list negative prompt embeddings for LTX-2 CFG.""" - return ( - lambda x: (not batch.do_classifier_free_guidance) + return lambda x: ( + (not batch.do_classifier_free_guidance) or V.is_tensor(x) or V.list_not_empty(x) ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/latent_preparation_av.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/latent_preparation_av.py index c2066c078..1fafae6d9 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/latent_preparation_av.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/latent_preparation_av.py @@ -40,9 +40,11 @@ class LTX2AVLatentPreparationStage(LatentPreparationStage): result.add_check( "prompt_or_embeds", None, - lambda _: V.string_or_list_strings(batch.prompt) - or V.list_not_empty(batch.prompt_embeds) - or V.is_tensor(batch.prompt_embeds), + lambda _: ( + V.string_or_list_strings(batch.prompt) + or V.list_not_empty(batch.prompt_embeds) + or V.is_tensor(batch.prompt_embeds) + ), ) if isinstance(batch.prompt_embeds, list): diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/text_connector.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/text_connector.py index e4d3d4e5d..8babe1856 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/text_connector.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/text_connector.py @@ -55,8 +55,7 @@ class LTX2TextConnectorStage(PipelineStage): if prompt_embeds is None or prompt_attention_mask is None: raise ValueError( - "LTX2TextConnectorStage requires prompt embeddings and " - "attention mask." + "LTX2TextConnectorStage requires prompt embeddings and attention mask." ) if batch.do_classifier_free_guidance: diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/condition_noise.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/condition_noise.py index 3a5e038a1..cf7f5aad7 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/condition_noise.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/condition_noise.py @@ -95,8 +95,7 @@ def minimax_h3_imgvid_cond_noise_aug_rows( full_t = target_latent_t + imgvid_cond_num_frames if full_t < latent_t: raise ValueError( - f"condition latent_t {latent_t} exceeds the noise draw " - f"length {full_t}" + f"condition latent_t {latent_t} exceeds the noise draw length {full_t}" ) generator = torch.Generator(device="cpu").manual_seed(int(seed)) noise = torch.randn( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/release_metadata.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/release_metadata.py index bd8d0e7d0..8df5e5a72 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/release_metadata.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/release_metadata.py @@ -63,7 +63,7 @@ class MiniMaxH3ReleaseMetadata: partition = raw.get("partition") if partition not in {"fl2va", "ref2va"}: raise ValueError( - "model_index.json._minimax_h3.partition must be one of " "fl2va, ref2va" + "model_index.json._minimax_h3.partition must be one of fl2va, ref2va" ) tasks = _string_list(raw.get("tasks"), "model_index.json._minimax_h3.tasks") aliases = raw.get("task_aliases", {}) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/request_validation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/request_validation.py index d18f0d7f2..707be26ff 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/request_validation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/request_validation.py @@ -196,7 +196,7 @@ def _validate_conditions( MINIMAX_H3_CONDITION_ROLE_REFERENCE, ): raise ValueError( - f"{cpath}.role must be keyframe or reference, " f"got {role!r}" + f"{cpath}.role must be keyframe or reference, got {role!r}" ) cond_type = _require_str(cond.get("type"), f"{cpath}.type") try: diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/latent_preparation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/latent_preparation.py index 86adbe91b..a3e9f67f1 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/latent_preparation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/latent_preparation.py @@ -145,8 +145,10 @@ class MiniMaxH3LatentPreparationStage(PipelineStage): result.add_check( "prompt_or_embeds", None, - lambda _: V.string_or_list_strings(batch.prompt) - or V.list_not_empty(batch.prompt_embeds), + lambda _: ( + V.string_or_list_strings(batch.prompt) + or V.list_not_empty(batch.prompt_embeds) + ), ) result.add_check("prompt_embeds", batch.prompt_embeds, V.list_of_tensors) result.add_check( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/timestep_preparation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/timestep_preparation.py index 512fb5aa9..bc4712420 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/timestep_preparation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/timestep_preparation.py @@ -120,8 +120,7 @@ class MiniMaxH3TimestepPreparationStage(PipelineStage): ) -> float: value = request_value source = ( - "request " - f"{'flow_shift' if modality == 'video' else 'audio_flow_shift'}" + f"request {'flow_shift' if modality == 'video' else 'audio_flow_shift'}" ) if value is None and model_scales is not None: value = model_scales.get(modality) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/video_adapter.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/video_adapter.py index 32dfd26d1..dc4da9c46 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/video_adapter.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/video_adapter.py @@ -70,8 +70,7 @@ class MiniMaxH3VideoModelAdapter: def validate_task_gate(self, task: Any, *, provided: bool) -> None: if not provided or task is None: raise ValueError( - "task is required for MiniMax H3; supported tasks: " - "fl2va, ref2va, t2va" + "task is required for MiniMax H3; supported tasks: fl2va, ref2va, t2va" ) if not isinstance(task, str): raise ValueError("task must be a non-empty string for MiniMax H3") diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/denoising.py index 32c3285e0..44aee43d7 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/denoising.py @@ -88,7 +88,7 @@ def _activation_time(P: float, delta: float) -> float: denom = P * (1.0 + P - delta) if denom <= 0 or delta >= 1.0 + P: raise ValueError( - f"delta={delta} >= 1+P={1+P:.4f}; criterion trivially satisfied." + f"delta={delta} >= 1+P={1 + P:.4f}; criterion trivially satisfied." ) return 1.0 / (1.0 + math.sqrt(delta / denom)) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py index 8662f2cfb..733968065 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py @@ -480,9 +480,11 @@ class TextEncodingStage(ConditionEncodingStage): result.add_check( "negative_prompt", batch.negative_prompt, - lambda x: not batch.do_classifier_free_guidance - or V.string_not_none(x) - or isinstance(x, str), + lambda x: ( + not batch.do_classifier_free_guidance + or V.string_not_none(x) + or isinstance(x, str) + ), ) result.add_check( "do_classifier_free_guidance", @@ -868,8 +870,10 @@ class TextEncodingStage(ConditionEncodingStage): result.add_check( "negative_prompt_embeds", batch.negative_prompt_embeds, - lambda x: not batch.do_classifier_free_guidance - or V.list_of_tensors_with_min_dims(x, 2), + lambda x: ( + not batch.do_classifier_free_guidance + or V.list_of_tensors_with_min_dims(x, 2) + ), ) if batch.debug: logger.debug(f"{batch.prompt_embeds=}") diff --git a/python/sglang/multimodal_gen/runtime/platforms/__init__.py b/python/sglang/multimodal_gen/runtime/platforms/__init__.py index 91d7c6a9c..e47eeac80 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/__init__.py +++ b/python/sglang/multimodal_gen/runtime/platforms/__init__.py @@ -247,7 +247,7 @@ def resolve_current_platform_cls_qualname() -> str: if platform_cls_qualname is not None: return platform_cls_qualname - raise RuntimeError("No platform plugin found. Please check your " "installation.") + raise RuntimeError("No platform plugin found. Please check your installation.") _current_platform: Platform | None = None diff --git a/python/sglang/multimodal_gen/runtime/platforms/interface.py b/python/sglang/multimodal_gen/runtime/platforms/interface.py index c6efec9ca..04eb9baf2 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/interface.py +++ b/python/sglang/multimodal_gen/runtime/platforms/interface.py @@ -392,8 +392,7 @@ class Platform: """ if cls.supported_quantization and quant not in cls.supported_quantization: raise ValueError( - f"{quant} quantization is currently not supported in " - f"{cls.device_name}." + f"{quant} quantization is currently not supported in {cls.device_name}." ) @classmethod diff --git a/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py index 8b7545198..674fb1a87 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py +++ b/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py @@ -37,7 +37,6 @@ def _kwargs_to_cpu(d: Any) -> Any: class RolloutDenoisingMixin: - def _maybe_prepare_rollout(self, batch: Req): """Prepare denoising loop for rollout.""" scheduler = batch.scheduler diff --git a/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_mixin.py index 3777ad4c4..aa52a25e6 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_mixin.py +++ b/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_mixin.py @@ -88,9 +88,9 @@ class SchedulerRLMixin(SchedulerRLDebugMixin): assert B == 1, "Generator must be a list if batch size is not 1" generator = [generator] else: - assert ( - len(generator) == B - ), "Generator list must have the same length as batch size" + assert len(generator) == B, ( + "Generator list must have the same length as batch size" + ) buffer = self._get_or_create_rollout_noise_buffer( rollout_session_data, rollout_session_data.latents_shape, device, dtype @@ -136,9 +136,9 @@ class SchedulerRLMixin(SchedulerRLDebugMixin): debug_mode = bool(getattr(batch, "rollout_debug_mode", False)) if not log_prob_no_const and sde_type != "ode": - assert ( - noise_level > 0 - ), "True log-probability computation requires a non-zero noise level." + assert noise_level > 0, ( + "True log-probability computation requires a non-zero noise level." + ) dt = next_sigma - current_sigma @@ -230,9 +230,9 @@ class SchedulerRLMixin(SchedulerRLDebugMixin): # Only enforce the "no full log-prob with ODE" constraint when the # user explicitly chose ODE globally. if sde_type == "ode": - assert ( - log_prob_no_const - ), "p_ode is always 0, true log_prob is meaningless, set rollout_log_prob_no_const to True to enable log_prob computation" + assert log_prob_no_const, ( + "p_ode is always 0, true log_prob is meaningless, set rollout_log_prob_no_const to True to enable log_prob computation" + ) else: raise ValueError(f"Unsupported sde_type: {sde_type}") diff --git a/python/sglang/multimodal_gen/runtime/server_args/server_args.py b/python/sglang/multimodal_gen/runtime/server_args/server_args.py index bd0bc7f73..41eef744e 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -2545,7 +2545,7 @@ class ServerArgs(DisaggServerArgsMixin): type=int, default=None, choices=[0, 1], - help="Quantize the attention sink too (1, default) " "or keep it bf16 (0).", + help="Quantize the attention sink too (1, default) or keep it bf16 (0).", ) parser.add_argument( "--kv-cache-quant-sink-keep", diff --git a/python/sglang/multimodal_gen/runtime/utils/distributed.py b/python/sglang/multimodal_gen/runtime/utils/distributed.py index 80afa0721..70f9a4e47 100644 --- a/python/sglang/multimodal_gen/runtime/utils/distributed.py +++ b/python/sglang/multimodal_gen/runtime/utils/distributed.py @@ -131,9 +131,9 @@ def generate_masked_orthogonal_rank_groups( idx = [(index // d) % s for s, d in zip(shape, stride)] # stride is a prefix_product result. And the value of stride[-1] # is not used. - assert ( - sum([x * y for x, y in zip(idx, stride[:-1])]) == index - ), "idx {} with shape {} mismatch the return idx {}".format(index, shape, idx) + assert sum([x * y for x, y in zip(idx, stride[:-1])]) == index, ( + "idx {} with shape {} mismatch the return idx {}".format(index, shape, idx) + ) return idx masked_shape = [s for s, m in zip(parallel_size, mask) if m] diff --git a/python/sglang/multimodal_gen/runtime/utils/logging_utils.py b/python/sglang/multimodal_gen/runtime/utils/logging_utils.py index 10ccc11cc..18b80cbd8 100644 --- a/python/sglang/multimodal_gen/runtime/utils/logging_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/logging_utils.py @@ -454,13 +454,13 @@ def enable_trace_function_call(log_file_path: str, root_dir: str | None = None): def set_uvicorn_logging_configs(server_args=None): from uvicorn.config import LOGGING_CONFIG - LOGGING_CONFIG["formatters"]["default"][ - "fmt" - ] = "[%(asctime)s] %(levelprefix)s %(message)s" + LOGGING_CONFIG["formatters"]["default"]["fmt"] = ( + "[%(asctime)s] %(levelprefix)s %(message)s" + ) LOGGING_CONFIG["formatters"]["default"]["datefmt"] = "%Y-%m-%d %H:%M:%S" - LOGGING_CONFIG["formatters"]["access"][ - "fmt" - ] = '[%(asctime)s] %(levelprefix)s %(client_addr)s - "%(request_line)s" %(status_code)s' + LOGGING_CONFIG["formatters"]["access"]["fmt"] = ( + '[%(asctime)s] %(levelprefix)s %(client_addr)s - "%(request_line)s" %(status_code)s' + ) LOGGING_CONFIG["formatters"]["access"]["datefmt"] = "%Y-%m-%d %H:%M:%S" # Install access log path filter into LOGGING_CONFIG so it survives diff --git a/python/sglang/multimodal_gen/runtime/weights/source.py b/python/sglang/multimodal_gen/runtime/weights/source.py index d5412158c..3b246c615 100644 --- a/python/sglang/multimodal_gen/runtime/weights/source.py +++ b/python/sglang/multimodal_gen/runtime/weights/source.py @@ -263,8 +263,7 @@ def _select_named_file(candidates: tuple[str, ...], weight_name: str) -> str: if not basename_matches: raise FileNotFoundError(f"Requested weight {weight_name!r} was not found") raise ValueError( - f"Weight name {weight_name!r} matches multiple files: " - f"{list(basename_matches)}" + f"Weight name {weight_name!r} matches multiple files: {list(basename_matches)}" ) diff --git a/python/sglang/multimodal_gen/test/run_suite.py b/python/sglang/multimodal_gen/test/run_suite.py index aa8f33a02..c5aec5f82 100644 --- a/python/sglang/multimodal_gen/test/run_suite.py +++ b/python/sglang/multimodal_gen/test/run_suite.py @@ -213,9 +213,7 @@ def _format_standalone_estimate_snippet( suite: str, standalone_file: str, measured_full_test_time_s: float ) -> str: return ( - f'"{suite}": {{\n' - f' "{standalone_file}": {measured_full_test_time_s:.1f},\n' - f"}}" + f'"{suite}": {{\n "{standalone_file}": {measured_full_test_time_s:.1f},\n}}' ) @@ -228,13 +226,13 @@ def _print_missing_standalone_estimate_message( suite, standalone_file, measured_full_test_time_s ) logger.error( - f'\n{"=" * 60}\n' + f"\n{'=' * 60}\n" f'Add standalone estimate for suite "{suite}" and file "{standalone_file}":\n\n' f"File: python/sglang/multimodal_gen/test/run_suite.py\n\n" f"Current partition used fallback estimate: " f"{DEFAULT_STANDALONE_EST_TIME_SECONDS:.1f}s\n\n" f"{snippet}\n" - f'{"=" * 60}\n' + f"{'=' * 60}\n" ) @@ -579,8 +577,7 @@ def _run_partition_assignment( else "" ) print( - f" - standalone: {standalone_file} " - f"({est_time:.1f}s{fallback_suffix})" + f" - standalone: {standalone_file} ({est_time:.1f}s{fallback_suffix})" ) for standalone_file in assignment.standalone_files: diff --git a/python/sglang/multimodal_gen/test/server/common/slack.py b/python/sglang/multimodal_gen/test/server/common/slack.py index 040df1ec9..4b59eed92 100644 --- a/python/sglang/multimodal_gen/test/server/common/slack.py +++ b/python/sglang/multimodal_gen/test/server/common/slack.py @@ -145,15 +145,13 @@ def upload_file_to_slack( title = ( "Original Image" if len(final_origin_paths) == 1 - else f"Original Image {i+1}" + else f"Original Image {i + 1}" ) uploads.append({"file": path, "title": title}) uploads.append({"file": file_path, "title": "Generated Image"}) - message = ( - f"*Case ID:* `{case_id}`\n" f"*Model:* `{model}`\n" f"*Prompt:* {prompt}" - ) + message = f"*Case ID:* `{case_id}`\n*Model:* `{model}`\n*Prompt:* {prompt}" client = WebClient(token=token, timeout=60) channel_id = "C0A02NDF7UY" diff --git a/python/sglang/multimodal_gen/test/server/test_server_common.py b/python/sglang/multimodal_gen/test/server/test_server_common.py index d1f524daf..c409ac675 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_common.py +++ b/python/sglang/multimodal_gen/test/server/test_server_common.py @@ -244,14 +244,14 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext: if needs_estimated_time and not is_baseline_generation_mode: _MISSING_ESTIMATED_TIME_CASES.add(case.id) logger.error( - f'\n{"=" * 60}\n' + f"\n{'=' * 60}\n" f'Add "estimated_full_test_time_s" to scenario "{case.id}":\n\n' f"File: {get_perf_baseline_update_path()}\n\n" f' "{case.id}": {{\n' f" ...\n" f' "estimated_full_test_time_s": {_measured_full_time:.1f}\n' f" }}\n" - f'{"=" * 60}\n' + f"{'=' * 60}\n" ) _print_case_log_separator(case.id, "END diffusion testcase") @@ -319,8 +319,7 @@ class DiffusionServerBase: tail = ctx.log_tail() message = ( - f"{case_id}: server process exited during generation " - f"(code {returncode})." + f"{case_id}: server process exited during generation (code {returncode})." ) if tail: message += f"\n\nServer log tail:\n{tail}" @@ -1184,7 +1183,9 @@ Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION} ), "loaded_adapters should be a non-empty list" assert any( a.get("nickname") == "default" for a in lora_info["loaded_adapters"] - ), f"nickname 'default' not found in loaded_adapters: {lora_info['loaded_adapters']}" + ), ( + f"nickname 'default' not found in loaded_adapters: {lora_info['loaded_adapters']}" + ) logger.info("[LoRA E2E] list_loras returned expected LoRA adapters") logger.info("[LoRA E2E] All LoRA API E2E tests passed for %s", case.id) @@ -1226,9 +1227,9 @@ Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION} json={"lora_nickname": "lora2", "lora_path": second_lora_path}, timeout=_CONTROL_API_TIMEOUT_SECS, ) - assert ( - resp.status_code == 200 - ), f"set_lora to second adapter failed: {resp.text}" + assert resp.status_code == 200, ( + f"set_lora to second adapter failed: {resp.text}" + ) logger.info( "[LoRA Switch E2E] Verifying generation with second LoRA for %s", case.id @@ -1320,9 +1321,9 @@ Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION} }, timeout=_CONTROL_API_TIMEOUT_SECS, ) - assert ( - resp.status_code == 200 - ), f"set_lora with multiple adapters failed: {resp.text}" + assert resp.status_code == 200, ( + f"set_lora with multiple adapters failed: {resp.text}" + ) rid, _ = self._run_generation_with_server_watchdog( ctx, case.id, generate_fn, client ) @@ -1339,9 +1340,9 @@ Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION} }, timeout=_CONTROL_API_TIMEOUT_SECS, ) - assert ( - resp.status_code == 200 - ), f"set_lora with different strengths failed: {resp.text}" + assert resp.status_code == 200, ( + f"set_lora with different strengths failed: {resp.text}" + ) rid, _ = self._run_generation_with_server_watchdog( ctx, case.id, generate_fn, client ) @@ -1363,9 +1364,9 @@ Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION} }, timeout=_CONTROL_API_TIMEOUT_SECS, ) - assert ( - resp.status_code == 200 - ), f"set_lora with cached adapters failed: {resp.text}" + assert resp.status_code == 200, ( + f"set_lora with cached adapters failed: {resp.text}" + ) rid, _ = self._run_generation_with_server_watchdog( ctx, case.id, generate_fn, client ) @@ -1377,9 +1378,9 @@ Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION} json={"lora_nickname": "default"}, timeout=_CONTROL_API_TIMEOUT_SECS, ) - assert ( - resp.status_code == 200 - ), f"set_lora back to single adapter failed: {resp.text}" + assert resp.status_code == 200, ( + f"set_lora back to single adapter failed: {resp.text}" + ) rid, content = self._run_generation_with_server_watchdog( ctx, case.id, generate_fn, client ) @@ -1403,28 +1404,28 @@ Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION} assert resp.status_code == 200, f"/v1/models failed: {resp.text}" data = resp.json() - assert ( - data["object"] == "list" - ), f"Expected object='list', got {data.get('object')}" + assert data["object"] == "list", ( + f"Expected object='list', got {data.get('object')}" + ) assert len(data["data"]) >= 1, "Expected at least one model in response" model = data["data"][0] assert "id" in model, "Model missing 'id' field" - assert ( - model["object"] == "model" - ), f"Expected object='model', got {model.get('object')}" - assert ( - model["id"] == case.server_args.model_path - ), f"Model ID mismatch: expected {case.server_args.model_path}, got {model['id']}" + assert model["object"] == "model", ( + f"Expected object='model', got {model.get('object')}" + ) + assert model["id"] == case.server_args.model_path, ( + f"Model ID mismatch: expected {case.server_args.model_path}, got {model['id']}" + ) # Verify extended diffusion-specific fields assert "num_gpus" in model, "Model missing 'num_gpus' field" assert "task_type" in model, "Model missing 'task_type' field" assert "dit_precision" in model, "Model missing 'dit_precision' field" assert "vae_precision" in model, "Model missing 'vae_precision' field" - assert ( - model["num_gpus"] == case.server_args.num_gpus - ), f"num_gpus mismatch: expected {case.server_args.num_gpus}, got {model['num_gpus']}" + assert model["num_gpus"] == case.server_args.num_gpus, ( + f"num_gpus mismatch: expected {case.server_args.num_gpus}, got {model['num_gpus']}" + ) expected_task_type = get_model_task_type_for_server_args(case.server_args).name assert model["task_type"] == expected_task_type, ( f"task_type mismatch: expected {expected_task_type}, " @@ -1466,9 +1467,9 @@ Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION} assert resp.status_code == 404, f"Expected 404, got {resp.status_code}" error_data = resp.json() assert "error" in error_data, "404 response missing 'error' field" - assert ( - error_data["error"]["code"] == "model_not_found" - ), f"Incorrect error code: {error_data['error'].get('code')}" + assert error_data["error"]["code"] == "model_not_found", ( + f"Incorrect error code: {error_data['error'].get('code')}" + ) logger.info("[Models API] GET /v1/models/non_existent returns 404 as expected") logger.info("[Models API] All /v1/models tests passed for %s", case.id) @@ -1500,13 +1501,13 @@ Pinned revision used by this check: {SGL_TEST_FILES_CI_DATA_REVISION} json=payload, timeout=_CONTROL_API_TIMEOUT_SECS, ) - assert ( - resp.status_code == 400 - ), f"Expected 400 for T2V input_reference, got {resp.status_code}: {resp.text}" + assert resp.status_code == 400, ( + f"Expected 400 for T2V input_reference, got {resp.status_code}: {resp.text}" + ) detail = resp.json().get("detail", "") - assert ( - "input_reference is not supported" in detail - ), f"Unexpected error detail for T2V input_reference: {detail}" + assert "input_reference is not supported" in detail, ( + f"Unexpected error detail for T2V input_reference: {detail}" + ) def test_diffusion_generation( self, diff --git a/python/sglang/multimodal_gen/test/single_test_file/component_accuracy/engine.py b/python/sglang/multimodal_gen/test/single_test_file/component_accuracy/engine.py index 4dfb4ade1..cf2a79055 100644 --- a/python/sglang/multimodal_gen/test/single_test_file/component_accuracy/engine.py +++ b/python/sglang/multimodal_gen/test/single_test_file/component_accuracy/engine.py @@ -431,9 +431,9 @@ class AccuracyEngine: ).item() rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 logger.info("[%s] Rank %s CosSim=%.6f", name, rank, cos_sim) - assert ( - cos_sim > threshold - ), f"Accuracy failure in {name}: CosSim {cos_sim:.4f} < {threshold}" + assert cos_sim > threshold, ( + f"Accuracy failure in {name}: CosSim {cos_sim:.4f} < {threshold}" + ) @staticmethod def transfer_weights( diff --git a/python/sglang/multimodal_gen/test/single_test_file/test_ar_models.py b/python/sglang/multimodal_gen/test/single_test_file/test_ar_models.py index 4a4b0f27d..768aafc40 100644 --- a/python/sglang/multimodal_gen/test/single_test_file/test_ar_models.py +++ b/python/sglang/multimodal_gen/test/single_test_file/test_ar_models.py @@ -94,8 +94,7 @@ class ARCluster(DisaggCluster): ) except Exception as e: raise RuntimeError( - f"AR model failed to start for {self.name}. Log tail:\n" - f"{_tail_log(log)}" + f"AR model failed to start for {self.name}. Log tail:\n{_tail_log(log)}" ) from e def _launch_server_head(self) -> None: @@ -144,7 +143,6 @@ class ARCluster(DisaggCluster): class _ARTestBase(_DisaggTestBase): - @classmethod def setUpClass(cls) -> None: super(CustomTestCase, cls).setUpClass() diff --git a/python/sglang/multimodal_gen/test/single_test_file/test_update_weights_from_disk.py b/python/sglang/multimodal_gen/test/single_test_file/test_update_weights_from_disk.py index 1f60b0772..11d30af1d 100644 --- a/python/sglang/multimodal_gen/test/single_test_file/test_update_weights_from_disk.py +++ b/python/sglang/multimodal_gen/test/single_test_file/test_update_weights_from_disk.py @@ -217,9 +217,9 @@ def _compute_checksum_from_disk(model_path: str, module_name: str) -> str: """ local_path = maybe_download_model(model_path) weights_dir = os.path.join(local_path, module_name) - assert os.path.exists( - weights_dir - ), f"No weights dir for {module_name} in {local_path}" + assert os.path.exists(weights_dir), ( + f"No weights dir for {module_name} in {local_path}" + ) safetensors_files = _list_safetensors_files(weights_dir) assert safetensors_files, f"No safetensors files in {weights_dir}" @@ -323,9 +323,9 @@ class _UpdateWeightsApiMixin: json=payload, timeout=timeout, ) - assert ( - response.status_code == 200 - ), f"get_weights_checksum failed: {response.status_code} {response.text}" + assert response.status_code == 200, ( + f"get_weights_checksum failed: {response.status_code} {response.text}" + ) return response.json() def _assert_server_matches_model( @@ -542,9 +542,9 @@ class TestUpdateWeightsFromDisk(_UpdateWeightsApiMixin): and perturbed_checksums.get(name) != "not_found" and base_checksums.get(name) != "not_found" ) - assert ( - text_encoder_modules - ), "Expected at least one text encoder module checksum" + assert text_encoder_modules, ( + "Expected at least one text encoder module checksum" + ) # perturbed → corrupted (should fail and rollback) rollback_targets = [_TRANSFORMER_MODULE, _VAE_MODULE] @@ -553,18 +553,18 @@ class TestUpdateWeightsFromDisk(_UpdateWeightsApiMixin): corrupted_vae_model_dir, target_modules=rollback_targets, ) - assert ( - status_code == 400 - ), f"Expected 400 on corrupted weights, got {status_code}" + assert status_code == 400, ( + f"Expected 400 on corrupted weights, got {status_code}" + ) assert not result.get("success", True) message = result.get("message", "") assert "rolled back" in message.lower() # The updater reports the first failing module in the error message. # With ordered target_modules=[transformer, vae], this makes the # failure point explicit: transformer is processed first, then vae fails. - assert ( - "Failed to update module 'vae'" in message - ), f"Expected vae to be the explicit failure point, got: {message}" + assert "Failed to update module 'vae'" in message, ( + f"Expected vae to be the explicit failure point, got: {message}" + ) rolled_back_checksums = self._get_weights_checksum(base_url) # 1) transformer: server == perturbed != base @@ -583,12 +583,12 @@ class TestUpdateWeightsFromDisk(_UpdateWeightsApiMixin): # 3) text encoder(s): server == base == perturbed for name in text_encoder_modules: - assert rolled_back_checksums.get(name) == perturbed_checksums.get( - name - ), f"Text encoder module '{name}' should stay equal to perturbed" - assert rolled_back_checksums.get(name) == base_checksums.get( - name - ), f"Text encoder module '{name}' should stay equal to base" + assert rolled_back_checksums.get(name) == perturbed_checksums.get(name), ( + f"Text encoder module '{name}' should stay equal to perturbed" + ) + assert rolled_back_checksums.get(name) == base_checksums.get(name), ( + f"Text encoder module '{name}' should stay equal to base" + ) class TestUpdateWeightsFromDiskWithOffload(_UpdateWeightsApiMixin): diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index b43a29fbb..727e51516 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/python/sglang/multimodal_gen/test/test_utils.py @@ -575,9 +575,9 @@ def validate_image_file( assert os.path.exists(file_path), f"Image file does not exist: {file_path}" # 2. Extension check - assert file_path.endswith( - f".{expected_ext}" - ), f"Expected .{expected_ext} extension, got: {file_path}" + assert file_path.endswith(f".{expected_ext}"), ( + f"Expected .{expected_ext} extension, got: {file_path}" + ) # 3. File size > 0 file_size = os.path.getsize(file_path) @@ -585,9 +585,9 @@ def validate_image_file( # 4. Filename validation actual_filename = os.path.basename(file_path) - assert ( - actual_filename == expected_filename - ), f"Filename mismatch: expected '{expected_filename}', got '{actual_filename}'" + assert actual_filename == expected_filename, ( + f"Filename mismatch: expected '{expected_filename}', got '{actual_filename}'" + ) # 5. Image format validation (magic bytes check based on expected format) with open(file_path, "rb") as f: @@ -603,12 +603,12 @@ def validate_image_file( if expected_width is not None and expected_height is not None: with Image.open(file_path) as img: width, height = img.size - assert ( - width == expected_width - ), f"Width mismatch: expected {expected_width}, got {width}" - assert ( - height == expected_height - ), f"Height mismatch: expected {expected_height}, got {height}" + assert width == expected_width, ( + f"Width mismatch: expected {expected_width}, got {width}" + ) + assert height == expected_height, ( + f"Height mismatch: expected {expected_height}, got {height}" + ) def _get_video_dimensions_from_metadata( @@ -707,9 +707,9 @@ def validate_video_file( # 4. Filename validation actual_filename = os.path.basename(file_path) - assert ( - actual_filename == expected_filename - ), f"Filename mismatch: expected '{expected_filename}', got '{actual_filename}'" + assert actual_filename == expected_filename, ( + f"Filename mismatch: expected '{expected_filename}', got '{actual_filename}'" + ) # 5. Video format validation (reuse is_mp4) with open(file_path, "rb") as f: @@ -719,12 +719,12 @@ def validate_video_file( # 6. Video dimension validation (using OpenCV) if expected_width is not None and expected_height is not None: actual_width, actual_height = get_video_dimensions(file_path) - assert ( - actual_width == expected_width - ), f"Video width mismatch: expected {expected_width}, got {actual_width}" - assert ( - actual_height == expected_height - ), f"Video height mismatch: expected {expected_height}, got {actual_height}" + assert actual_width == expected_width, ( + f"Video width mismatch: expected {expected_width}, got {actual_width}" + ) + assert actual_height == expected_height, ( + f"Video height mismatch: expected {expected_height}, got {actual_height}" + ) @dataclass(frozen=True) @@ -776,9 +776,9 @@ def probe_audio_stream(file_path: str) -> AudioStreamInfo: ) assert sample_rate > 0, f"Audio stream has invalid sample rate: {sample_rate}" assert channels > 0, f"Audio stream has invalid channel count: {channels}" - assert ( - math.isfinite(duration) and duration > 0 - ), f"Audio stream has invalid duration: {duration}" + assert math.isfinite(duration) and duration > 0, ( + f"Audio stream has invalid duration: {duration}" + ) return AudioStreamInfo(sample_rate, channels, duration) diff --git a/python/sglang/multimodal_gen/test/unit/test_attention_backend_override.py b/python/sglang/multimodal_gen/test/unit/test_attention_backend_override.py index 83daa8316..c44ad2d41 100644 --- a/python/sglang/multimodal_gen/test/unit/test_attention_backend_override.py +++ b/python/sglang/multimodal_gen/test/unit/test_attention_backend_override.py @@ -26,7 +26,7 @@ def _fake_backend_cls(enum, *, ring_capable=True): return SimpleNamespace( get_enum=lambda: enum, supports_ring_rotation=lambda: ring_capable, - get_impl_cls=lambda: (lambda **kwargs: f"{enum.name.lower()}_impl"), + get_impl_cls=lambda: lambda **kwargs: f"{enum.name.lower()}_impl", ) diff --git a/python/sglang/multimodal_gen/test/unit/test_component_quantization_admission.py b/python/sglang/multimodal_gen/test/unit/test_component_quantization_admission.py index 9f1077188..df61480f6 100644 --- a/python/sglang/multimodal_gen/test/unit/test_component_quantization_admission.py +++ b/python/sglang/multimodal_gen/test/unit/test_component_quantization_admission.py @@ -101,8 +101,9 @@ class TestComponentQuantizationAdmission(unittest.TestCase): component_precisions={}, component_quantizations={}, component_weights_paths={}, - should_direct_gpu_weight_load_component=lambda component: component - == "vocoder", + should_direct_gpu_weight_load_component=lambda component: ( + component == "vocoder" + ), ) with self.assertRaisesRegex( @@ -117,8 +118,9 @@ class TestComponentQuantizationAdmission(unittest.TestCase): component_precisions={}, component_quantizations={}, component_weights_paths={}, - should_direct_gpu_weight_load_component=lambda component: component - == "audio_vae", + should_direct_gpu_weight_load_component=lambda component: ( + component == "audio_vae" + ), ) with self.assertRaisesRegex( diff --git a/python/sglang/multimodal_gen/test/unit/test_consistency_metrics.py b/python/sglang/multimodal_gen/test/unit/test_consistency_metrics.py index e00572b24..1383b38db 100644 --- a/python/sglang/multimodal_gen/test/unit/test_consistency_metrics.py +++ b/python/sglang/multimodal_gen/test/unit/test_consistency_metrics.py @@ -263,8 +263,10 @@ def test_remote_platform_video_gt_prefers_platform_sglang_before_default_officia monkeypatch.setattr( test_utils, "_remote_file_exists", - lambda url: url.startswith(sglang_platform_prefix) - or url.startswith(official_default_prefix), + lambda url: ( + url.startswith(sglang_platform_prefix) + or url.startswith(official_default_prefix) + ), ) files = test_utils._find_remote_consistency_gt_files( diff --git a/python/sglang/multimodal_gen/test/unit/test_cosmos3_rollout.py b/python/sglang/multimodal_gen/test/unit/test_cosmos3_rollout.py index 70ca35733..ffd1788b7 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cosmos3_rollout.py +++ b/python/sglang/multimodal_gen/test/unit/test_cosmos3_rollout.py @@ -48,7 +48,6 @@ def _prepare(serving, batch, explicit_shift): class TestPrepareRolloutRequestScheduler(unittest.TestCase): - def test_inherits_serving_grid_without_explicit_shift(self): serving = _serving_scheduler() batch = _rollout_batch() @@ -126,7 +125,6 @@ class _FusedParamModule(torch.nn.Module): class TestWeightsUpdaterFusedParams(unittest.TestCase): - def test_merge_index_reaches_weight_loader_as_shard_id(self): module = _FusedParamModule() calls = [] diff --git a/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_tp_graph_capture.py b/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_tp_graph_capture.py index 1ce940dea..ebc591f02 100644 --- a/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_tp_graph_capture.py +++ b/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_tp_graph_capture.py @@ -58,15 +58,19 @@ class TestBCGTPGraphCapture(CustomTestCase): custom_allreduce_cls = MagicMock(return_value=expected) group = SimpleNamespace(cpu_group=object(), device=torch.device("cuda:0")) - with patch( - "sglang.multimodal_gen.runtime.distributed.group_coordinator.current_platform.is_cuda", - return_value=True, - ), patch( - "sglang.srt.distributed.device_communicators.custom_all_reduce.dispatch_custom_allreduce", - return_value=custom_allreduce_cls, - ) as dispatch, patch( - "sglang.srt.distributed.device_communicators.custom_all_reduce_v2.CustomAllReduceV2", - custom_allreduce_cls, + with ( + patch( + "sglang.multimodal_gen.runtime.distributed.group_coordinator.current_platform.is_cuda", + return_value=True, + ), + patch( + "sglang.srt.distributed.device_communicators.custom_all_reduce.dispatch_custom_allreduce", + return_value=custom_allreduce_cls, + ) as dispatch, + patch( + "sglang.srt.distributed.device_communicators.custom_all_reduce_v2.CustomAllReduceV2", + custom_allreduce_cls, + ), ): GroupCoordinator._init_srt_custom_allreduce(group) @@ -83,12 +87,15 @@ class TestBCGTPGraphCapture(CustomTestCase): custom_allreduce_cls = MagicMock(return_value=expected) group = SimpleNamespace(cpu_group=object(), device=object()) - with patch( - "sglang.multimodal_gen.runtime.distributed.group_coordinator.current_platform.is_cuda", - return_value=False, - ), patch( - "sglang.srt.distributed.device_communicators.custom_all_reduce.CustomAllreduce", - custom_allreduce_cls, + with ( + patch( + "sglang.multimodal_gen.runtime.distributed.group_coordinator.current_platform.is_cuda", + return_value=False, + ), + patch( + "sglang.srt.distributed.device_communicators.custom_all_reduce.CustomAllreduce", + custom_allreduce_cls, + ), ): GroupCoordinator._init_srt_custom_allreduce(group) @@ -120,10 +127,14 @@ class TestBCGTPGraphCapture(CustomTestCase): """Drive the CUDA branch of graph_capture() with a fake custom AR.""" group = SimpleNamespace(srt_custom_allreduce=custom_ar) ctx = GraphCaptureContext(MagicMock()) - with patch( - "sglang.multimodal_gen.runtime.distributed.group_coordinator.current_platform.is_cuda_alike", - return_value=True, - ), patch("torch.cuda.stream"), patch("torch.cuda.current_stream"): + with ( + patch( + "sglang.multimodal_gen.runtime.distributed.group_coordinator.current_platform.is_cuda_alike", + return_value=True, + ), + patch("torch.cuda.stream"), + patch("torch.cuda.current_stream"), + ): with GroupCoordinator.graph_capture(group, ctx) as yielded: events.append("body") return yielded, ctx @@ -157,12 +168,15 @@ class TestBCGTPGraphCapture(CustomTestCase): tp_group.graph_capture = _recording_context(events, "tp") runner = SimpleNamespace(_capture_stream=capture_stream) - with patch( - "sglang.multimodal_gen.runtime.distributed.parallel_state.model_parallel_is_initialized", - return_value=initialized, - ), patch( - "sglang.multimodal_gen.runtime.distributed.parallel_state.get_tp_group", - return_value=tp_group, + with ( + patch( + "sglang.multimodal_gen.runtime.distributed.parallel_state.model_parallel_is_initialized", + return_value=initialized, + ), + patch( + "sglang.multimodal_gen.runtime.distributed.parallel_state.get_tp_group", + return_value=tp_group, + ), ): with BaseBreakableCudaGraphRunner._tp_graph_capture(runner): events.append("body") @@ -190,12 +204,15 @@ class TestBCGTPGraphCapture(CustomTestCase): tp_group.graph_capture = _graph_capture runner = SimpleNamespace(_capture_stream=capture_stream) - with patch( - "sglang.multimodal_gen.runtime.distributed.parallel_state.model_parallel_is_initialized", - return_value=True, - ), patch( - "sglang.multimodal_gen.runtime.distributed.parallel_state.get_tp_group", - return_value=tp_group, + with ( + patch( + "sglang.multimodal_gen.runtime.distributed.parallel_state.model_parallel_is_initialized", + return_value=True, + ), + patch( + "sglang.multimodal_gen.runtime.distributed.parallel_state.get_tp_group", + return_value=tp_group, + ), ): with BaseBreakableCudaGraphRunner._tp_graph_capture(runner): pass @@ -232,20 +249,23 @@ class TestBCGTPGraphCapture(CustomTestCase): graph._segments = [] kwargs = {"hidden_states": torch.zeros(1)} - with patch.object( - BaseBreakableCudaGraphRunner, - "_tp_graph_capture", - _recording_context(events, "tp"), - ), patch.object( - runner_mod, "BreakableCUDAGraph", return_value=graph - ), patch.object( - runner_mod, - "enable_breakable_cuda_graph", - _recording_context(events, "bcg_enable"), - ), patch.object( - runner_mod, - "BreakableCUDAGraphCapture", - _recording_context(events, "bcg_capture"), + with ( + patch.object( + BaseBreakableCudaGraphRunner, + "_tp_graph_capture", + _recording_context(events, "tp"), + ), + patch.object(runner_mod, "BreakableCUDAGraph", return_value=graph), + patch.object( + runner_mod, + "enable_breakable_cuda_graph", + _recording_context(events, "bcg_enable"), + ), + patch.object( + runner_mod, + "BreakableCUDAGraphCapture", + _recording_context(events, "bcg_capture"), + ), ): runner._capture(kwargs, key=runner_mod._signature_kwargs(kwargs)) diff --git a/python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py b/python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py index 5caeb91e3..4888bbf44 100644 --- a/python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py +++ b/python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py @@ -296,6 +296,6 @@ def test_the_forced_host_size_behaves_like_a_machine_of_that_size(monkeypatch): # and a larger pretend machine reports more room, same process monkeypatch.setenv("SGLANG_DIFFUSION_TEST_FORCE_HOST_AVAILABLE_GIB", "64") larger = host_memory_budget.host_memory_available_bytes() - assert ( - abs((larger - available) - 32 * 1024**3) < 512 * 1024**2 - ), "the same process on a machine twice the size has one machine more room" + assert abs((larger - available) - 32 * 1024**3) < 512 * 1024**2, ( + "the same process on a machine twice the size has one machine more room" + ) diff --git a/python/sglang/multimodal_gen/test/unit/test_image_encoder_loader.py b/python/sglang/multimodal_gen/test/unit/test_image_encoder_loader.py index e0c4ef961..4f0861a81 100644 --- a/python/sglang/multimodal_gen/test/unit/test_image_encoder_loader.py +++ b/python/sglang/multimodal_gen/test/unit/test_image_encoder_loader.py @@ -110,8 +110,9 @@ class TestImageEncoderQuantizationAdmission(unittest.TestCase): def test_native_only_quantized_architecture_does_not_fall_back(self): self.server_args.pipeline_config.native_only_components = ("image_encoder",) config = self._component_config("UnknownVisionModel", quantized=True) - with self._config_patch(config), self.assertRaises( - NativeComponentLoaderRequired + with ( + self._config_patch(config), + self.assertRaises(NativeComponentLoaderRequired), ): self._load() self.load_native.assert_not_called() @@ -121,8 +122,9 @@ class TestImageEncoderQuantizationAdmission(unittest.TestCase): "architectures": ["UnknownVisionModel"], "quantization_config": {"quant_method": "not-a-format"}, } - with self._config_patch(config), self.assertRaises( - ComponentCheckpointUnsupportedError + with ( + self._config_patch(config), + self.assertRaises(ComponentCheckpointUnsupportedError), ): self._load() self.load_native.assert_not_called() @@ -159,18 +161,22 @@ class TestImageEncoderNativeLoading(unittest.TestCase): ) loader = ImageEncoderLoader() - with mock.patch( - "sglang.multimodal_gen.runtime.loader.component_loaders." - "component_loader.get_hf_config", - return_value=component_config, - ), mock.patch.object( - loader, - "resolve_native_transformers_model_class", - return_value=model_class, - ), mock.patch.object( - loader, - "target_device", - return_value=torch.device("cuda:0"), + with ( + mock.patch( + "sglang.multimodal_gen.runtime.loader.component_loaders." + "component_loader.get_hf_config", + return_value=component_config, + ), + mock.patch.object( + loader, + "resolve_native_transformers_model_class", + return_value=model_class, + ), + mock.patch.object( + loader, + "target_device", + return_value=torch.device("cuda:0"), + ), ): component = loader.load_native( "/model/image_encoder", @@ -211,16 +217,20 @@ class TestImageEncoderNativeLoading(unittest.TestCase): ) loader = ImageEncoderLoader() - with mock.patch( - "sglang.multimodal_gen.runtime.loader.component_loaders." - "component_loader.get_hf_config", - return_value=component_config, - ), mock.patch.object( - loader, - "resolve_native_transformers_model_class", - return_value=model_class, - ), self.assertRaisesRegex( - ComponentCheckpointUnsupportedError, "requires resident placement" + with ( + mock.patch( + "sglang.multimodal_gen.runtime.loader.component_loaders." + "component_loader.get_hf_config", + return_value=component_config, + ), + mock.patch.object( + loader, + "resolve_native_transformers_model_class", + return_value=model_class, + ), + self.assertRaisesRegex( + ComponentCheckpointUnsupportedError, "requires resident placement" + ), ): loader.load_native( "/model/image_encoder", @@ -266,34 +276,42 @@ class TestImageEncoderNativeLoading(unittest.TestCase): ) loader = ImageEncoderLoader() - with mock.patch.object( - loader, - "load_customized", - side_effect=NativeComponentLoaderRequired("use Transformers"), - ), mock.patch.object( - loader, - "resolve_native_transformers_model_class", - return_value=model_class, - ), mock.patch.object( - loader, - "target_device", - return_value=torch.device("cuda:0"), - ), mock.patch( - "sglang.multimodal_gen.runtime.loader.component_loaders." - "component_loader.get_hf_config", - return_value=component_config, - ), mock.patch( - "sglang.multimodal_gen.runtime.loader.component_loaders." - "component_loader.current_platform.get_available_gpu_memory", - return_value=10.0, - ), mock.patch( - "sglang.multimodal_gen.runtime.loader.component_loaders." - "component_loader.get_memory_usage_of_component", - return_value=0.0, - ), mock.patch( - "sglang.multimodal_gen.runtime.loader.component_loaders." - "component_loader.format_component_residency", - return_value="resident", + with ( + mock.patch.object( + loader, + "load_customized", + side_effect=NativeComponentLoaderRequired("use Transformers"), + ), + mock.patch.object( + loader, + "resolve_native_transformers_model_class", + return_value=model_class, + ), + mock.patch.object( + loader, + "target_device", + return_value=torch.device("cuda:0"), + ), + mock.patch( + "sglang.multimodal_gen.runtime.loader.component_loaders." + "component_loader.get_hf_config", + return_value=component_config, + ), + mock.patch( + "sglang.multimodal_gen.runtime.loader.component_loaders." + "component_loader.current_platform.get_available_gpu_memory", + return_value=10.0, + ), + mock.patch( + "sglang.multimodal_gen.runtime.loader.component_loaders." + "component_loader.get_memory_usage_of_component", + return_value=0.0, + ), + mock.patch( + "sglang.multimodal_gen.runtime.loader.component_loaders." + "component_loader.format_component_residency", + return_value="resident", + ), ): component, _ = loader.load( "/model/image_encoder", diff --git a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py index 83fce212b..b8320f96b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py +++ b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py @@ -1375,15 +1375,15 @@ def test_mapped_layers_ship_through_the_courier(tmp_path, monkeypatch): manager.prefetch_layer(0, non_blocking=True) assert 0 in manager._courier_inflight, "an async prefetch hands the layer over" - assert ( - 0 not in manager._gpu_layers - ), "the layer is not ready until its tensors are bound on this thread" + assert 0 not in manager._gpu_layers, ( + "the layer is not ready until its tensors are bound on this thread" + ) manager.prefetch_layer(0, non_blocking=False) assert 0 in manager._gpu_layers and not manager._courier_inflight - assert torch.equal( - model.blocks[0].weight.detach().cpu(), expected - ), "the bytes that went through the courier's slot must be the checkpoint's" + assert torch.equal(model.blocks[0].weight.detach().cpu(), expected), ( + "the bytes that went through the courier's slot must be the checkpoint's" + ) def test_the_courier_kill_switch_forces_the_synchronous_path(tmp_path, monkeypatch): @@ -1397,9 +1397,9 @@ def test_the_courier_kill_switch_forces_the_synchronous_path(tmp_path, monkeypat manager.release_all() manager.prefetch_layer(0, non_blocking=True) assert not manager._courier_inflight and manager._mapped_courier is None - assert ( - 0 in manager._gpu_layers - ), "with the courier disabled the direct synchronous path serves the layer" + assert 0 in manager._gpu_layers, ( + "with the courier disabled the direct synchronous path serves the layer" + ) def test_release_all_drains_the_courier(tmp_path, monkeypatch): diff --git a/python/sglang/multimodal_gen/test/unit/test_lora_format_adapter.py b/python/sglang/multimodal_gen/test/unit/test_lora_format_adapter.py index 7dec429dd..57a2c87fe 100644 --- a/python/sglang/multimodal_gen/test/unit/test_lora_format_adapter.py +++ b/python/sglang/multimodal_gen/test/unit/test_lora_format_adapter.py @@ -339,9 +339,9 @@ def main() -> None: class TestLoRAFormatAdapter: def test_lora_format_adapter_all_formats(self): results = _run_all_tests() - assert all( - r["pass"] for r in results - ), "At least one LoRA format adapter case failed" + assert all(r["pass"] for r in results), ( + "At least one LoRA format adapter case failed" + ) if __name__ == "__main__": diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_vae_parallel_modes.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_vae_parallel_modes.py index 3f54206de..473bf04df 100644 --- a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_vae_parallel_modes.py +++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_vae_parallel_modes.py @@ -62,9 +62,7 @@ def test_unvalidated_decode_modes_are_rejected(mode): def test_vit_attention_uses_local_usp_backend_dispatch(): - module = ( - "sglang.multimodal_gen.runtime.models.vaes." "minimax_h3_video_vae.attention" - ) + module = "sglang.multimodal_gen.runtime.models.vaes.minimax_h3_video_vae.attention" with ( mock.patch(f"{module}.current_platform.is_cuda", return_value=True), mock.patch(f"{module}.USPAttention", autospec=True) as usp_attention, @@ -98,9 +96,7 @@ def test_audio_vae_attention_defaults_to_local_sdpa_and_allows_fa(): self.input_dtype = query.dtype return query - module = ( - "sglang.multimodal_gen.runtime.models.vaes." "minimax_h3_audio_vae.audio_vae" - ) + module = "sglang.multimodal_gen.runtime.models.vaes.minimax_h3_audio_vae.audio_vae" recording_fa = RecordingFA() with ( mock.patch(f"{module}.current_platform.is_cuda", return_value=True), diff --git a/python/sglang/multimodal_gen/test/unit/test_ministral3_generation.py b/python/sglang/multimodal_gen/test/unit/test_ministral3_generation.py index 9922743a8..cf94ed302 100644 --- a/python/sglang/multimodal_gen/test/unit/test_ministral3_generation.py +++ b/python/sglang/multimodal_gen/test/unit/test_ministral3_generation.py @@ -105,8 +105,9 @@ def test_native_ministral3_matches_hf_prefill_and_generation(): torch.testing.assert_close(native_output.logits, reference_output.logits) assert len(native_output.past_key_values.layers) == config.num_hidden_layers - with torch.no_grad(), set_forward_context( - current_timestep=0, attn_metadata=None + with ( + torch.no_grad(), + set_forward_context(current_timestep=0, attn_metadata=None), ): native_ids = native.generate(input_ids, max_new_tokens=2, do_sample=False) with torch.no_grad(): diff --git a/python/sglang/multimodal_gen/test/unit/test_modelopt_fp4_backend.py b/python/sglang/multimodal_gen/test/unit/test_modelopt_fp4_backend.py index db3a9109c..cecdb0a18 100644 --- a/python/sglang/multimodal_gen/test/unit/test_modelopt_fp4_backend.py +++ b/python/sglang/multimodal_gen/test/unit/test_modelopt_fp4_backend.py @@ -12,8 +12,9 @@ ENV_PATH = ( def _backend(env_value, *, sm120): CudaPlatform.get_modelopt_flashinfer_fp4_backend.cache_clear() try: - with patch(ENV_PATH, env_value), patch.object( - CudaPlatform, "is_sm120", classmethod(lambda cls: sm120) + with ( + patch(ENV_PATH, env_value), + patch.object(CudaPlatform, "is_sm120", classmethod(lambda cls: sm120)), ): return CudaPlatform.get_modelopt_flashinfer_fp4_backend() finally: diff --git a/python/sglang/multimodal_gen/test/unit/test_nvtx_pytorch_hooks.py b/python/sglang/multimodal_gen/test/unit/test_nvtx_pytorch_hooks.py index fdd3fe1dd..864b62509 100644 --- a/python/sglang/multimodal_gen/test/unit/test_nvtx_pytorch_hooks.py +++ b/python/sglang/multimodal_gen/test/unit/test_nvtx_pytorch_hooks.py @@ -236,8 +236,8 @@ class TestComponentResidencyNvtxHooks(unittest.TestCase): manager.strategy_for = lambda _component_name, _module: ResidentStrategy() self.assertTrue(manager._should_keep_single_dit("transformer", module)) - manager.strategy_for = ( - lambda _component_name, _module: ComponentOffloadStrategy() + manager.strategy_for = lambda _component_name, _module: ( + ComponentOffloadStrategy() ) self.assertFalse(manager._should_keep_single_dit("transformer", module)) diff --git a/python/sglang/multimodal_gen/test/unit/test_pi05_runtime_helpers.py b/python/sglang/multimodal_gen/test/unit/test_pi05_runtime_helpers.py index d4ff2de69..733885e5e 100644 --- a/python/sglang/multimodal_gen/test/unit/test_pi05_runtime_helpers.py +++ b/python/sglang/multimodal_gen/test/unit/test_pi05_runtime_helpers.py @@ -688,10 +688,10 @@ def test_sample_actions_only_hoists_denoise_layout_for_eager(): model.denoise_step = lambda _ctx, x_t, _t, **_kwargs: torch.zeros_like(x_t) layout_calls = [] model.core_model = SimpleNamespace( - prepare_denoise_layout=lambda *args, **kwargs: layout_calls.append( - (args, kwargs) + prepare_denoise_layout=lambda *args, **kwargs: ( + layout_calls.append((args, kwargs)) + or (None, torch.zeros(1, 2, dtype=torch.long)) ) - or (None, torch.zeros(1, 2, dtype=torch.long)) ) observation = SimpleNamespace(batch_size=1) prefix_context = _prefix_context(1.0, "prompt") diff --git a/python/sglang/multimodal_gen/test/unit/test_regional_torch_compile.py b/python/sglang/multimodal_gen/test/unit/test_regional_torch_compile.py index 878b74b8f..80eeff9de 100644 --- a/python/sglang/multimodal_gen/test/unit/test_regional_torch_compile.py +++ b/python/sglang/multimodal_gen/test/unit/test_regional_torch_compile.py @@ -28,8 +28,9 @@ class _CompilableModule(nn.Module): class _RegionalModel(_CompilableModule): _compile_conditions = [ - lambda name, _module: name.startswith("transformer_blocks.") - and name.count(".") == 1 + lambda name, _module: ( + name.startswith("transformer_blocks.") and name.count(".") == 1 + ) ] def __init__(self): diff --git a/python/sglang/multimodal_gen/test/unit/test_rollout_api.py b/python/sglang/multimodal_gen/test/unit/test_rollout_api.py index 8f31e6531..64b0298c0 100644 --- a/python/sglang/multimodal_gen/test/unit/test_rollout_api.py +++ b/python/sglang/multimodal_gen/test/unit/test_rollout_api.py @@ -21,7 +21,6 @@ from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( class TestTensorToBytesRoundtrip(unittest.TestCase): - def _roundtrip(self, t: torch.Tensor): encoded = tensor_to_bytes(t) self.assertIsInstance(encoded, bytes) diff --git a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py index e7f8f0b30..21cdb2136 100644 --- a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py +++ b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py @@ -427,9 +427,9 @@ class TestSchedulerFlowGRPOStepAlignmentUnit(unittest.TestCase): def _mock_variance_noise(_batch, *_args, **_kwargs): variance_noise_call_count["n"] += 1 - scheduler._get_rollout_session_data(_batch).noise_buffer = ( - variance_noise_ref - ) + scheduler._get_rollout_session_data( + _batch + ).noise_buffer = variance_noise_ref return variance_noise_ref scheduler._rollout_variance_noise = ( # type: ignore[method-assign] diff --git a/python/sglang/multimodal_gen/test/unit/test_server_args.py b/python/sglang/multimodal_gen/test/unit/test_server_args.py index ac41dceed..470516511 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py @@ -870,7 +870,6 @@ class TestDiffusionModelDetection(unittest.TestCase): class TestMiniMaxH3Routing(unittest.TestCase): - def test_semantic_variants_map_to_checkpoint_partitions(self): self.assertEqual( MiniMaxH3Pipeline.model_subfolder_for_variant("fl2va"), "FL2VA" diff --git a/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py b/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py index e349daaa7..c6fb023ca 100644 --- a/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py +++ b/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py @@ -166,18 +166,22 @@ class TestTextEncoderClassResolution(unittest.TestCase): } loader = TextEncoderLoader() - with mock.patch.object( - TextEncoderLoader, - "resolve_native_transformers_model_class", - return_value=transformers_model_class, - ), mock.patch.object( - loader, - "target_device", - return_value=torch.device("cuda:0"), - ), mock.patch( - "sglang.multimodal_gen.runtime.loader.component_loaders." - "component_loader.get_hf_config", - return_value=component_config, + with ( + mock.patch.object( + TextEncoderLoader, + "resolve_native_transformers_model_class", + return_value=transformers_model_class, + ), + mock.patch.object( + loader, + "target_device", + return_value=torch.device("cuda:0"), + ), + mock.patch( + "sglang.multimodal_gen.runtime.loader.component_loaders." + "component_loader.get_hf_config", + return_value=component_config, + ), ): encoder = loader.load_native( "/model/text_encoder", @@ -719,9 +723,12 @@ class TestTextEncoderQuantization(unittest.TestCase): "CLIPTextModel", "ThirdPartyTextEncoder", ): - with self.subTest(architecture=architecture), self.assertRaisesRegex( - NativeComponentLoaderRequired, - "delegates serialized quant_method='bitsandbytes' checkpoint", + with ( + self.subTest(architecture=architecture), + self.assertRaisesRegex( + NativeComponentLoaderRequired, + "delegates serialized quant_method='bitsandbytes' checkpoint", + ), ): _resolve_and_configure_encoder_quantization( SimpleNamespace(architectures=[architecture], quant_config=None), diff --git a/python/sglang/multimodal_gen/test/unit/test_usp_ipc_a2a_guard.py b/python/sglang/multimodal_gen/test/unit/test_usp_ipc_a2a_guard.py index e2553c650..840b9818b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_usp_ipc_a2a_guard.py +++ b/python/sglang/multimodal_gen/test/unit/test_usp_ipc_a2a_guard.py @@ -37,9 +37,10 @@ class TestIpcInputA2AQkvGuard(unittest.TestCase): # group lookup to prove the guard is not what rejected it. q = torch.zeros(1, 128, 8, 64) group = mock.MagicMock(return_value=None) - with mock.patch.object( - usp, "get_ulysses_parallel_world_size", lambda: 2 - ), mock.patch.object(usp, "_ipc_ready_group", group): + with ( + mock.patch.object(usp, "get_ulysses_parallel_world_size", lambda: 2), + mock.patch.object(usp, "_ipc_ready_group", group), + ): self.assertIsNone(usp._ipc_input_a2a_qkv(q, q.clone(), q.clone())) # Reached the group lookup, so the shape guard did not reject it. self.assertEqual(group.call_count, 1) diff --git a/python/sglang/multimodal_gen/test/unit/test_usp_ring_replicated.py b/python/sglang/multimodal_gen/test/unit/test_usp_ring_replicated.py index d3a985cb1..2a557d35c 100644 --- a/python/sglang/multimodal_gen/test/unit/test_usp_ring_replicated.py +++ b/python/sglang/multimodal_gen/test/unit/test_usp_ring_replicated.py @@ -77,8 +77,8 @@ class RingReplicatedBase(unittest.TestCase): ), patch( f"{_LAYER}.ring_attn", - side_effect=lambda q, k, v, impl, return_softmax_lse: _ring_pair_via_impl( - impl, q, k, v + side_effect=lambda q, k, v, impl, return_softmax_lse: ( + _ring_pair_via_impl(impl, q, k, v) ), ), ) @@ -142,14 +142,19 @@ class TestRingReplicatedSuffix(RingReplicatedBase): o.copy_(t) ps = self._patches(ring_ws=1) - with ps[0], ps[4], ps[5], ps[6], patch( - f"{_LAYER}.get_ring_parallel_world_size", return_value=1 - ), patch(f"{_LAYER}.get_sequence_parallel_world_size", return_value=2), patch( - f"{_LAYER}.get_ulysses_parallel_world_size", return_value=2 - ), patch( - f"{_LAYER}.get_sp_group", return_value=SimpleNamespace(ulysses_group=None) - ), patch( - "torch.distributed.all_gather", side_effect=_fake_gather + with ( + ps[0], + ps[4], + ps[5], + ps[6], + patch(f"{_LAYER}.get_ring_parallel_world_size", return_value=1), + patch(f"{_LAYER}.get_sequence_parallel_world_size", return_value=2), + patch(f"{_LAYER}.get_ulysses_parallel_world_size", return_value=2), + patch( + f"{_LAYER}.get_sp_group", + return_value=SimpleNamespace(ulysses_group=None), + ), + patch("torch.distributed.all_gather", side_effect=_fake_gather), ): # Identity-mocked collectives don't reproduce head-shard shapes, # so the final concat may fail — the kernel K order is recorded diff --git a/python/sglang/multimodal_gen/utils.py b/python/sglang/multimodal_gen/utils.py index 6729138ee..be12a1551 100644 --- a/python/sglang/multimodal_gen/utils.py +++ b/python/sglang/multimodal_gen/utils.py @@ -135,7 +135,6 @@ def current_stream() -> torch.cuda.Stream | None: class StoreBoolean(argparse.Action): - def __init__(self, option_strings, dest, default=False, required=False, help=None): super().__init__( option_strings=option_strings, @@ -157,7 +156,7 @@ class StoreBoolean(argparse.Action): setattr(namespace, self.dest, False) else: raise ValueError( - f"Invalid boolean value: {values}. " "Expected 'true' or 'false'." + f"Invalid boolean value: {values}. Expected 'true' or 'false'." ) else: setattr(namespace, self.dest, bool(values)) @@ -295,8 +294,7 @@ class FlexibleArgumentParser(argparse.ArgumentParser): if args[0] == "serve": if index == 1: raise ValueError( - "No model_tag specified! Please check your command-line" - " arguments." + "No model_tag specified! Please check your command-line arguments." ) command = args_before_config[0] model_tag = args_before_config[1] @@ -484,7 +482,7 @@ def update_environment_variables(envs: dict[str, str]): for k, v in envs.items(): if k in os.environ and os.environ[k] != v: logger.warning( - "Overwriting environment variable %s " "from '%s' to '%s'", + "Overwriting environment variable %s from '%s' to '%s'", k, os.environ[k], v, @@ -509,7 +507,7 @@ def run_method( func = getattr(obj, method) except AttributeError: raise NotImplementedError( - f"Method {method!r} is not" " implemented." + f"Method {method!r} is not implemented." ) from None else: func = partial(method, obj) # type: ignore @@ -549,7 +547,6 @@ def get_exception_traceback() -> str: class TypeBasedDispatcher: - def __init__(self, mapping: list[tuple[type, Callable]]): self._mapping = mapping @@ -626,9 +623,9 @@ def dict_to_3d_list( """ # Case 1: no data, but fixed shape requested if mask_strategy is None: - assert ( - t_max is not None and l_max is not None and h_max is not None - ), "If mask_strategy is None, you must provide t_max, l_max, and h_max" + assert t_max is not None and l_max is not None and h_max is not None, ( + "If mask_strategy is None, you must provide t_max, l_max, and h_max" + ) return [ [[None for _ in range(h_max)] for _ in range(l_max)] for _ in range(t_max) ] diff --git a/python/sglang/srt/arg_groups/attention_hook.py b/python/sglang/srt/arg_groups/attention_hook.py index 282563416..d9bfab8c2 100644 --- a/python/sglang/srt/arg_groups/attention_hook.py +++ b/python/sglang/srt/arg_groups/attention_hook.py @@ -91,9 +91,9 @@ def handle_attention_backend_compatibility(server_args: Any): cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED ), ) - assert ( - cfg.speculative_algorithm is None - ), "Speculative decoding is currently not supported with Flex Attention backend" + assert cfg.speculative_algorithm is None, ( + "Speculative decoding is currently not supported with Flex Attention backend" + ) # Whisper's encoder token padding conflicts with prefix caching. # Only disable for Whisper; other encoder-decoder models (e.g., mllama) use radix cache. diff --git a/python/sglang/srt/arg_groups/cuda_graph_hook.py b/python/sglang/srt/arg_groups/cuda_graph_hook.py index 788f08292..4b3a5dfba 100644 --- a/python/sglang/srt/arg_groups/cuda_graph_hook.py +++ b/python/sglang/srt/arg_groups/cuda_graph_hook.py @@ -148,7 +148,7 @@ def apply_cuda_graph_compatibility(server_args: Any): and attention_backends_of(resolved_view(server_args))[0] != "trtllm_mla" ): logger.info( - "Using tc_piecewise CUDA graph for validated multimodal " "decoder prefill." + "Using tc_piecewise CUDA graph for validated multimodal decoder prefill." ) declare_resolution( server_args, @@ -183,16 +183,20 @@ def disable_tc_piecewise_cudagraph_if_incompatible(server_args: Any): ("pipeline parallelism (pp_size > 1)", lambda: cfg.pp_size > 1), ( "non-CUDA hardware (HIP/NPU/CPU/MPS/XPU)", - lambda: get_platform().is_hip - or get_platform().is_npu - or is_cpu() - or is_mps() - or get_platform().is_xpu, + lambda: ( + get_platform().is_hip + or get_platform().is_npu + or is_cpu() + or is_mps() + or get_platform().is_xpu + ), ), ( "OOT platform without piecewise support", - lambda: current_platform.is_out_of_tree() - and not current_platform.support_piecewise_cuda_graph(), + lambda: ( + current_platform.is_out_of_tree() + and not current_platform.support_piecewise_cuda_graph() + ), ), ( "MoE A2A backend", @@ -203,16 +207,20 @@ def disable_tc_piecewise_cudagraph_if_incompatible(server_args: Any): ("LoRA", lambda: bool(cfg.lora_paths) or cfg.enable_lora), ( "multimodal model", - lambda: model_config_of(server_args).is_multimodal - and not model_config_of( - server_args - ).is_multimodal_piecewise_cuda_graph_supported, + lambda: ( + model_config_of(server_args).is_multimodal + and not model_config_of( + server_args + ).is_multimodal_piecewise_cuda_graph_supported + ), ), ( "GGUF quantization", - lambda: cfg.load_format == "gguf" - or resolved_view(server_args).quantization == "gguf" - or check_gguf_file(cfg.model_path), + lambda: ( + cfg.load_format == "gguf" + or resolved_view(server_args).quantization == "gguf" + or check_gguf_file(cfg.model_path) + ), ), ("DLLM (diffusion LLM)", lambda: cfg.dllm_algorithm is not None), ( @@ -227,8 +235,9 @@ def disable_tc_piecewise_cudagraph_if_incompatible(server_args: Any): ("symmetric memory", lambda: cfg.enable_symm_mem), ( "expert distribution recorder", - lambda: cfg.enable_eplb - or cfg.expert_distribution_recorder_mode is not None, + lambda: ( + cfg.enable_eplb or cfg.expert_distribution_recorder_mode is not None + ), ), ( "context parallel (attn_cp_size > 1)", @@ -279,8 +288,10 @@ def disable_breakable_cudagraph_if_incompatible(server_args: Any): # CP all_gather replay size mismatch under BCG. ( "context parallel (attn_cp_size > 1)", - lambda: resolved_view(server_args).attn_cp_size > 1 - and not supports_prefill_cp_bcg(server_args), + lambda: ( + resolved_view(server_args).attn_cp_size > 1 + and not supports_prefill_cp_bcg(server_args) + ), ), # Capture builds a dummy extend forward with attn_dcp_metadata=None. ( @@ -294,16 +305,20 @@ def disable_breakable_cudagraph_if_incompatible(server_args: Any): ), ( "unvalidated a2a backend", - lambda: resolved_view(server_args).moe_a2a_backend - not in ("none", "deepep", "megamoe", "flashinfer"), + lambda: ( + resolved_view(server_args).moe_a2a_backend + not in ("none", "deepep", "megamoe", "flashinfer") + ), ), # Multimodal prefill replay faults under BCG; allowlisted archs opt back in. ( "multimodal model", - lambda: model_config_of(server_args).is_multimodal - and not model_config_of( - server_args - ).is_multimodal_breakable_cuda_graph_supported, + lambda: ( + model_config_of(server_args).is_multimodal + and not model_config_of( + server_args + ).is_multimodal_breakable_cuda_graph_supported + ), ), ] for name, predicate in rules: diff --git a/python/sglang/srt/arg_groups/deepseek_v4_hook.py b/python/sglang/srt/arg_groups/deepseek_v4_hook.py index 2c991845f..551625c6d 100644 --- a/python/sglang/srt/arg_groups/deepseek_v4_hook.py +++ b/python/sglang/srt/arg_groups/deepseek_v4_hook.py @@ -148,11 +148,13 @@ def apply_deepseek_v4_defaults(server_args: ServerArgs, model_arch: str) -> None assert cfg.speculative_algorithm in ( "EAGLE", "DSPARK", - ), f"Only EAGLE and DSPARK speculative algorithms are supported for {model_arch}" + ), ( + f"Only EAGLE and DSPARK speculative algorithms are supported for {model_arch}" + ) if cfg.speculative_algorithm == "EAGLE": - assert ( - cfg.speculative_eagle_topk == 1 - ), f"Only EAGLE speculative algorithm with topk == 1 is supported for {model_arch}" + assert cfg.speculative_eagle_topk == 1, ( + f"Only EAGLE speculative algorithm with topk == 1 is supported for {model_arch}" + ) def validate_deepseek_v4_cp(server_args: ServerArgs) -> None: @@ -163,7 +165,7 @@ def validate_deepseek_v4_cp(server_args: ServerArgs) -> None: if cfg.cp_strategy != "interleave": raise ValueError( - "DeepSeekV4 only supports interleave CP strategy, " f"got {cfg.cp_strategy}" + f"DeepSeekV4 only supports interleave CP strategy, got {cfg.cp_strategy}" ) declare_resolution( @@ -196,12 +198,12 @@ def validate_deepseek_v4_cp(server_args: ServerArgs) -> None: "validate_deepseek_v4_cp", attn_cp_size=cfg.tp_size // cfg.dp_size, ) - assert ( - cfg.dp_size == 1 - ), "For round-robin split mode, dp attention is not supported." - assert ( - cfg.tp_size <= 8 - ), "Context parallel only supports single machine (tp_size <= 8). Cross-machine CP has precision issues." + assert cfg.dp_size == 1, ( + "For round-robin split mode, dp attention is not supported." + ) + assert cfg.tp_size <= 8, ( + "Context parallel only supports single machine (tp_size <= 8). Cross-machine CP has precision issues." + ) supported_a2a_backends = ("none", "deepep", "megamoe", "mori") if cfg.moe_a2a_backend not in supported_a2a_backends: raise ValueError( diff --git a/python/sglang/srt/arg_groups/hisparse_hook.py b/python/sglang/srt/arg_groups/hisparse_hook.py index 3b38b1e40..cdb841dff 100644 --- a/python/sglang/srt/arg_groups/hisparse_hook.py +++ b/python/sglang/srt/arg_groups/hisparse_hook.py @@ -96,9 +96,9 @@ def validate_hisparse(server_args: ServerArgs) -> None: "models (e.g., DeepSeek V3.2, GLM-5) and DeepSeek V4 now. " ) - assert ( - cfg.disable_radix_cache - ), "Hierarchical sparse attention currently requires --disable-radix-cache." + assert cfg.disable_radix_cache, ( + "Hierarchical sparse attention currently requires --disable-radix-cache." + ) # DSv4 hisparse handles its own dtype/backend pairing elsewhere; the dtype- # aware checks below only apply to the DSA hisparse path. diff --git a/python/sglang/srt/arg_groups/lora_hook.py b/python/sglang/srt/arg_groups/lora_hook.py index c349db689..82df754e1 100644 --- a/python/sglang/srt/arg_groups/lora_hook.py +++ b/python/sglang/srt/arg_groups/lora_hook.py @@ -76,9 +76,9 @@ def check_lora_server_args(server_args: Any): pinned=False, ) elif isinstance(lora_path, dict): - assert ( - "lora_name" in lora_path and "lora_path" in lora_path - ), f"When providing LoRA paths as a list of dict, each dict should contain 'lora_name' and 'lora_path' keys. Got: {lora_path}" + assert "lora_name" in lora_path and "lora_path" in lora_path, ( + f"When providing LoRA paths as a list of dict, each dict should contain 'lora_name' and 'lora_path' keys. Got: {lora_path}" + ) lora_ref = LoRARef( lora_id=LoRARef.deterministic_id( lora_path["lora_name"], lora_path["lora_path"] @@ -129,14 +129,14 @@ def check_lora_server_args(server_args: Any): lora_target_modules=set(cfg.lora_target_modules), ) if "all" in cfg.lora_target_modules: - assert ( - len(cfg.lora_target_modules) == 1 - ), "If 'all' is specified in --lora-target-modules, it should be the only module specified." + assert len(cfg.lora_target_modules) == 1, ( + "If 'all' is specified in --lora-target-modules, it should be the only module specified." + ) # Ensure sufficient information is provided for LoRA initialization. - assert cfg.lora_paths or ( - cfg.max_lora_rank and cfg.lora_target_modules - ), "When no initial --lora-paths is provided, you need to specify both --max-lora-rank and --lora-target-modules for LoRA initialization." + assert cfg.lora_paths or (cfg.max_lora_rank and cfg.lora_target_modules), ( + "When no initial --lora-paths is provided, you need to specify both --max-lora-rank and --lora-target-modules for LoRA initialization." + ) # Validate max_loaded_loras if cfg.max_loaded_loras is not None: @@ -158,9 +158,9 @@ def check_lora_server_args(server_args: Any): if cfg.lora_use_virtual_experts: logger.info("Virtual expert computation enabled.") - assert ( - cfg.lora_drain_wait_threshold >= 0.0 - ), "--lora-drain-wait-threshold must be non-negative." + assert cfg.lora_drain_wait_threshold >= 0.0, ( + "--lora-drain-wait-threshold must be non-negative." + ) def check_lora_speculative_compatibility(server_args: Any): diff --git a/python/sglang/srt/arg_groups/mamba_hook.py b/python/sglang/srt/arg_groups/mamba_hook.py index 39e28cc9e..9107f432e 100644 --- a/python/sglang/srt/arg_groups/mamba_hook.py +++ b/python/sglang/srt/arg_groups/mamba_hook.py @@ -97,9 +97,9 @@ def handle_int8_mamba_checkpoint(server_args: Any): def validate_mamba_extra_buffer(view, model_arch: str, *, mamba_cache_chunk_size_of): - assert supports_mamba_cache_extra_buffer( - view, model_arch - ), f"extra_buffer is not supported for {model_arch}; use no_buffer." + assert supports_mamba_cache_extra_buffer(view, model_arch), ( + f"extra_buffer is not supported for {model_arch}; use no_buffer." + ) assert ( get_platform().is_cuda or get_platform().is_musa @@ -142,9 +142,9 @@ def validate_mamba_extra_buffer(view, model_arch: str, *, mamba_cache_chunk_size def validate_mamba_no_buffer(view, model_arch: str): assert view.page_size in (1, None), "no_buffer only supports page_size=1." - assert ( - view.disable_overlap_schedule - ), "no_buffer do not support overlap schedule. Try to set disable_overlap_schedule=True." - assert ( - view.attention_backend != "trtllm_mha" - ), "no_buffer do not support trtllm_mha attention backend." + assert view.disable_overlap_schedule, ( + "no_buffer do not support overlap schedule. Try to set disable_overlap_schedule=True." + ) + assert view.attention_backend != "trtllm_mha", ( + "no_buffer do not support trtllm_mha attention backend." + ) diff --git a/python/sglang/srt/arg_groups/memory_hook.py b/python/sglang/srt/arg_groups/memory_hook.py index 39c5b2ad8..0685fe805 100644 --- a/python/sglang/srt/arg_groups/memory_hook.py +++ b/python/sglang/srt/arg_groups/memory_hook.py @@ -177,9 +177,9 @@ def handle_gpu_memory_settings(server_args: Any, gpu_mem): ) decode_cuda_graph_config.bs = generate_cpu_graph_batch_sizes(server_args) - assert ( - cfg.torch_compile_max_bs > 0 - ), "cuda_graph_config[decode].bs should contain positive batch sizes" + assert cfg.torch_compile_max_bs > 0, ( + "cuda_graph_config[decode].bs should contain positive batch sizes" + ) decode_cuda_graph_config.max_bs = cfg.torch_compile_max_bs if prefill_cuda_graph_config.max_bs is None: diff --git a/python/sglang/srt/arg_groups/model_hook.py b/python/sglang/srt/arg_groups/model_hook.py index 56bfd3b87..a56fe3f31 100644 --- a/python/sglang/srt/arg_groups/model_hook.py +++ b/python/sglang/srt/arg_groups/model_hook.py @@ -218,9 +218,9 @@ def handle_model_specific_adjustments(server_args: Any): run_post_process_pass(server_args, _dsa_split_backend_resolution) if cfg.enable_prefill_cp: - assert ( - cfg.disaggregation_mode != "decode" - ), "CP is only supported for prefill when PD disaggregation, please remove --enable-prefill-cp." + assert cfg.disaggregation_mode != "decode", ( + "CP is only supported for prefill when PD disaggregation, please remove --enable-prefill-cp." + ) if ( cfg.enable_dsa_cache_layer_split and cfg.disaggregation_mode != "prefill" @@ -423,9 +423,9 @@ def handle_model_specific_adjustments(server_args: Any): # (arg_groups/overrides.py: _gpt_oss_overrides). if resolved_view(server_args).moe_runner_backend == "triton_kernel": - assert ( - resolved_view(server_args).ep_size == 1 - ), "Triton kernel MoE is only supported when ep_size == 1" + assert resolved_view(server_args).ep_size == 1, ( + "Triton kernel MoE is only supported when ep_size == 1" + ) elif model_arch in ("MiMoV2ForCausalLM", "MiMoV2FlashForCausalLM"): if model_arch == "MiMoV2ForCausalLM" and not cfg.encoder_only: @@ -481,7 +481,9 @@ def handle_model_specific_adjustments(server_args: Any): "ascend", "trtllm_mha", "intel_xpu", - }, f"fa3, aiter, triton, ascend, trtllm_mha or intel_xpu is required for Llama4 model but got {attention_backend}" + }, ( + f"fa3, aiter, triton, ascend, trtllm_mha or intel_xpu is required for Llama4 model but got {attention_backend}" + ) # The moe_runner_backend selection moved to the override registry # (arg_groups/overrides.py: _llama4_overrides). # Gemma2/Gemma3 (disable_hybrid_swa_memory) moved to the override registry @@ -523,9 +525,9 @@ def handle_model_specific_adjustments(server_args: Any): # https://docs.sglang.ai/advanced_features/attention_backend.html accepted_backends = ["fa3", "triton", "trtllm_mha"] attention_backend = resolved_view(server_args).attention_backend - assert ( - attention_backend in accepted_backends - ), f"One of the attention backends in {accepted_backends} is required for {model_arch}, but got {attention_backend}" + assert attention_backend in accepted_backends, ( + f"One of the attention backends in {accepted_backends} is required for {model_arch}, but got {attention_backend}" + ) elif model_arch in ["Olmo2ForCausalLM"]: # disable_hybrid_swa_memory + attention backend selection moved to # the override registry (arg_groups/overrides.py: _olmo2_overrides). @@ -534,9 +536,9 @@ def handle_model_specific_adjustments(server_args: Any): # is used for the Olmo2 architecture. Olmo2 does not use sliding window attention # but Olmo3 does. attention_backend = resolved_view(server_args).attention_backend - assert ( - attention_backend != "flashinfer" - ), "FlashInfer backend can significantly degrade the performance of Olmo3 models." + assert attention_backend != "flashinfer", ( + "FlashInfer backend can significantly degrade the performance of Olmo3 models." + ) logger.info(f"Using {attention_backend} as attention backend for {model_arch}.") elif model_arch in [ diff --git a/python/sglang/srt/arg_groups/model_overrides/deepseek_v2.py b/python/sglang/srt/arg_groups/model_overrides/deepseek_v2.py index e50d65be1..19d48b659 100644 --- a/python/sglang/srt/arg_groups/model_overrides/deepseek_v2.py +++ b/python/sglang/srt/arg_groups/model_overrides/deepseek_v2.py @@ -59,12 +59,12 @@ def _deepseek_family_overrides(server_args: Any, hf_config: Any) -> dict: "moe_a2a_backend=deepep, ep_size=tp_size, batch_size=1." ) else: - assert ( - cfg.dp_size == 1 - ), "interleave DSA CP does not support DP attention." - assert ( - cfg.tp_size <= 8 - ), "Context parallel only supports single machine (tp_size <= 8). Cross-machine CP has precision issues." + assert cfg.dp_size == 1, ( + "interleave DSA CP does not support DP attention." + ) + assert cfg.tp_size <= 8, ( + "Context parallel only supports single machine (tp_size <= 8). Cross-machine CP has precision issues." + ) # Note(kpham-sgl): Keep attn_tp_size == 1 under DSA CP. # DSACPLayerCommunicator does not all-reduce attention-TP # partial o_proj outputs before replicated dense FFNs. diff --git a/python/sglang/srt/arg_groups/model_overrides/deepseek_v4.py b/python/sglang/srt/arg_groups/model_overrides/deepseek_v4.py index 558260893..3aa0c4711 100644 --- a/python/sglang/srt/arg_groups/model_overrides/deepseek_v4.py +++ b/python/sglang/srt/arg_groups/model_overrides/deepseek_v4.py @@ -70,7 +70,5 @@ def _deepseek_v4_overrides(server_args: Any, hf_config: Any) -> dict: ) ): overrides["moe_runner_backend"] = "flashinfer_mxfp4" - logger.info( - "Use flashinfer_mxfp4 as MoE runner backend for " f"{model_arch}." - ) + logger.info(f"Use flashinfer_mxfp4 as MoE runner backend for {model_arch}.") return overrides diff --git a/python/sglang/srt/arg_groups/model_overrides/gpt_oss.py b/python/sglang/srt/arg_groups/model_overrides/gpt_oss.py index 447702749..0bee69abc 100644 --- a/python/sglang/srt/arg_groups/model_overrides/gpt_oss.py +++ b/python/sglang/srt/arg_groups/model_overrides/gpt_oss.py @@ -67,7 +67,6 @@ def _gpt_oss_overrides(server_args: Any, hf_config: Any) -> dict: # use bf16 for mxfp4 triton kernels overrides["dtype"] = "bfloat16" if cfg.moe_runner_backend == "auto": - if get_platform().is_sm100 and is_mxfp4_quant_format: overrides["moe_runner_backend"] = "flashinfer_mxfp4" logger.warning( diff --git a/python/sglang/srt/arg_groups/moe_hook.py b/python/sglang/srt/arg_groups/moe_hook.py index 7eeafb0c4..b1146dd32 100644 --- a/python/sglang/srt/arg_groups/moe_hook.py +++ b/python/sglang/srt/arg_groups/moe_hook.py @@ -47,7 +47,9 @@ def handle_moe_kernel_config(server_args: Any): "modelopt_fp8", "modelopt_mixed", None, - ], f"Invalid quantization '{view.quantization}'. \nFlashInfer Cutlass MOE supports only: 'modelopt_fp4', 'modelopt_fp8', 'modelopt_mixed', or bfloat16 (None)." + ], ( + f"Invalid quantization '{view.quantization}'. \nFlashInfer Cutlass MOE supports only: 'modelopt_fp4', 'modelopt_fp8', 'modelopt_mixed', or bfloat16 (None)." + ) assert view.ep_size in [ 1, cfg.tp_size, @@ -58,7 +60,9 @@ def handle_moe_kernel_config(server_args: Any): assert ( view.quantization in ["modelopt_fp4", "modelopt_mixed", "nvfp4_online"] or model_config_of(server_args).nvfp4_moe_meta is not None - ), f"Invalid quantization '{view.quantization}'. \nFlashInfer CuteDSL MOE currently supports only: 'modelopt_fp4', 'modelopt_mixed' (with NVFP4 MoE layers), 'nvfp4_online', or hybrid NVFP4 models." + ), ( + f"Invalid quantization '{view.quantization}'. \nFlashInfer CuteDSL MOE currently supports only: 'modelopt_fp4', 'modelopt_mixed' (with NVFP4 MoE layers), 'nvfp4_online', or hybrid NVFP4 models." + ) assert view.ep_size in [ 1, cfg.tp_size, @@ -90,7 +94,9 @@ def handle_moe_kernel_config(server_args: Any): "modelopt_mixed", "compressed-tensors", None, - ], f"Invalid quantization '{view.quantization}'. \nFlashInfer TRTLLM MOE supports only: 'modelopt_fp4', 'nvfp4_online', 'fp8', 'modelopt_fp8', 'modelopt_mixed', 'compressed-tensors', or bfloat16 (None)." + ], ( + f"Invalid quantization '{view.quantization}'. \nFlashInfer TRTLLM MOE supports only: 'modelopt_fp4', 'nvfp4_online', 'fp8', 'modelopt_fp8', 'modelopt_mixed', 'compressed-tensors', or bfloat16 (None)." + ) if view.moe_runner_backend == "flashinfer_trtllm_routed": assert view.quantization in [ @@ -100,7 +106,9 @@ def handle_moe_kernel_config(server_args: Any): "modelopt_mixed", "nvfp4_online", None, - ], f"Invalid quantization '{view.quantization}'. \nFlashInfer TRTLLM routed MOE supports only: 'fp8', 'mxfp8', 'modelopt_fp4', 'modelopt_mixed', 'nvfp4_online', or bfloat16 (None)." + ], ( + f"Invalid quantization '{view.quantization}'. \nFlashInfer TRTLLM routed MOE supports only: 'fp8', 'mxfp8', 'modelopt_fp4', 'modelopt_mixed', 'nvfp4_online', or bfloat16 (None)." + ) # The runner-driven shared-experts fusion disables moved to the # pipeline (arg_groups/overrides.py: _moe_runner_fusion_disable), @@ -113,9 +121,9 @@ def handle_moe_kernel_config(server_args: Any): "fp8", "mxfp8", ]: - assert ( - resolved_view(server_args).ep_size == 1 - ), "FP8/MXFP8 Cutlass MoE is only supported with ep_size == 1" + assert resolved_view(server_args).ep_size == 1, ( + "FP8/MXFP8 Cutlass MoE is only supported with ep_size == 1" + ) def handle_a2a_moe(server_args: Any): @@ -256,7 +264,9 @@ def handle_a2a_moe(server_args: Any): assert ( resolved_view(server_args).enable_dp_attention and cfg.dp_size == cfg.tp_size - ), "Flashinfer MoE A2A is only supported with dp_size == tp_size and --enable-dp-attention" + ), ( + "Flashinfer MoE A2A is only supported with dp_size == tp_size and --enable-dp-attention" + ) if cfg.deepep_mode != "auto": logger.warning("--deepep-mode is ignored for Flashinfer MoE A2A") use_cutedsl_w4a16 = ( diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 9debcb6d7..9ebae0e1c 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -449,8 +449,9 @@ import sglang.srt.arg_groups.model_overrides # noqa: F401 @register_model_override_predicate( - lambda arch: "Step3p5ForCausalLM" in arch - or "Step3p7ForConditionalGeneration" in arch + lambda arch: ( + "Step3p5ForCausalLM" in arch or "Step3p7ForConditionalGeneration" in arch + ) ) def _step3p_overrides(server_args: Any, hf_config: Any) -> dict: cfg = resolving_view(server_args) @@ -1261,9 +1262,9 @@ def _cutedsl_prefill_backend_fill(view: Any) -> dict: or view.prefill_attention_backend == "cutedsl_mla" ): return {} - assert ( - view.prefill_attention_backend != "cutedsl_mla" - ), "CuteDSL MLA only supports decoding for now" + assert view.prefill_attention_backend != "cutedsl_mla", ( + "CuteDSL MLA only supports decoding for now" + ) if not get_platform().is_sm100: raise ValueError( "CuteDSL MLA backend is only supported on Blackwell GPUs (SM100). Please use a different backend." @@ -1435,12 +1436,12 @@ def _dp_lm_head_validation(view: Any) -> dict: dp LM head and the TP LM-head all-to-all path. Reads the mid-resolution values through the view.""" if view.enable_dp_lm_head: - assert ( - view.enable_dp_attention - ), "Please enable dp attention when setting enable_dp_lm_head. " + assert view.enable_dp_attention, ( + "Please enable dp attention when setting enable_dp_lm_head. " + ) if view.enable_tp_lm_head_all_to_all: assert view.enable_dp_attention, ( - "Please enable dp attention when setting " "enable_tp_lm_head_all_to_all." + "Please enable dp attention when setting enable_tp_lm_head_all_to_all." ) assert not view.enable_dp_lm_head, ( "--enable-tp-lm-head-all-to-all uses a TP-sharded LM head and is " @@ -1500,7 +1501,7 @@ def _moe_runner_backend_quant_constraints(view: Any) -> dict: moe_runner_backend = mxfp8_default elif moe_runner_backend not in allowed: logger.warning( - "mxfp8 quantization supports only %s backends. " "Overriding %r.", + "mxfp8 quantization supports only %s backends. Overriding %r.", ", ".join(allowed), moe_runner_backend, ) @@ -1845,7 +1846,6 @@ def mamba_cache_chunk_size(server_args: Any) -> int: from sglang.srt.arg_groups.overrides import model_config_of if not hasattr(server_args, "_mamba_cache_chunk_size"): - try: from sglang.kernels.ops.attention.fla.chunk_delta_h import ( CHUNK_SIZE as FLA_CHUNK_SIZE, @@ -1857,9 +1857,9 @@ def mamba_cache_chunk_size(server_args: Any) -> int: hf_config = model_config_of(server_args).hf_config chunk_size = getattr(hf_config, "mamba_chunk_size", FLA_CHUNK_SIZE) page_size = resolved_view(server_args).page_size - assert ( - max(chunk_size, page_size) % min(chunk_size, page_size) == 0 - ), f"For SSM models, either chunk_size or page_size must be divisible by the other, got {chunk_size=}, {page_size=}" + assert max(chunk_size, page_size) % min(chunk_size, page_size) == 0, ( + f"For SSM models, either chunk_size or page_size must be divisible by the other, got {chunk_size=}, {page_size=}" + ) if not getattr(server_args, "_resolution_finished", False): return max(chunk_size, page_size) server_args._mamba_cache_chunk_size = max(chunk_size, page_size) diff --git a/python/sglang/srt/arg_groups/parallel_hook.py b/python/sglang/srt/arg_groups/parallel_hook.py index b066a56cb..7573a92ed 100644 --- a/python/sglang/srt/arg_groups/parallel_hook.py +++ b/python/sglang/srt/arg_groups/parallel_hook.py @@ -59,8 +59,7 @@ def handle_context_parallelism(server_args: Any): and not cfg.language_model_only ): raise ValueError( - "MiMo V2 CP-v2 only supports text inference; add " - "--language-only." + "MiMo V2 CP-v2 only supports text inference; add --language-only." ) if cfg.enable_prefill_cp and cfg.cp_strategy is None: @@ -81,40 +80,40 @@ def handle_context_parallelism(server_args: Any): view = resolved_view(server_args) if view.attn_cp_size > 1: # The tp_size is the world size, not the real tensor parallel size - assert ( - cfg.tp_size % view.attn_cp_size == 0 - ), "tp_size must be divisible by attn_cp_size" - assert ( - cfg.tp_size % (cfg.dp_size * view.attn_cp_size) == 0 - ), "tp_size must be divisible by dp_size * attn_cp_size" + assert cfg.tp_size % view.attn_cp_size == 0, ( + "tp_size must be divisible by attn_cp_size" + ) + assert cfg.tp_size % (cfg.dp_size * view.attn_cp_size) == 0, ( + "tp_size must be divisible by dp_size * attn_cp_size" + ) - assert ( - not cfg.enable_aiter_allreduce_fusion - ), "Aiter allreduce fusion is not supported with context parallelism" + assert not cfg.enable_aiter_allreduce_fusion, ( + "Aiter allreduce fusion is not supported with context parallelism" + ) if cfg.moe_dp_size > 1: # The tp_size is the world size, not the real tensor parallel size - assert ( - cfg.tp_size % cfg.moe_dp_size == 0 - ), "tp_size must be divisible by moe_dp_size" - assert ( - view.ep_size * cfg.moe_dp_size <= cfg.tp_size - ), "ep_size * moe_dp_size must be less than or equal to tp_size" + assert cfg.tp_size % cfg.moe_dp_size == 0, ( + "tp_size must be divisible by moe_dp_size" + ) + assert view.ep_size * cfg.moe_dp_size <= cfg.tp_size, ( + "ep_size * moe_dp_size must be less than or equal to tp_size" + ) assert cfg.pp_size == 1, "PP is not supported with context parallelism" if view.ep_size > 1: - assert ( - view.ep_size * cfg.moe_dp_size == cfg.tp_size - ), "ep_size * moe_dp_size must be equal to tp_size" + assert view.ep_size * cfg.moe_dp_size == cfg.tp_size, ( + "ep_size * moe_dp_size must be equal to tp_size" + ) - assert ( - not cfg.enable_aiter_allreduce_fusion - ), "Aiter allreduce fusion is not supported with context parallelism" + assert not cfg.enable_aiter_allreduce_fusion, ( + "Aiter allreduce fusion is not supported with context parallelism" + ) if view.attn_cp_size != cfg.moe_dp_size: - assert ( - cfg.moe_dp_size == 1 - ), "attn_cp_size != moe_dp_size is only supported when moe_dp_size == 1" + assert cfg.moe_dp_size == 1, ( + "attn_cp_size != moe_dp_size is only supported when moe_dp_size == 1" + ) from sglang.srt.layers.cp.base import init_cp_strategy @@ -244,26 +243,26 @@ def handle_dwdp(server_args: Any): if cfg.dwdp_size <= 1: return - assert ( - cfg.dwdp_size >= 2 - ), f"dwdp_size must be >= 2 when enabled, got {cfg.dwdp_size}" - assert ( - cfg.dwdp_size == cfg.tp_size - ), f"dwdp_size ({cfg.dwdp_size}) must equal tp_size ({cfg.tp_size})" + assert cfg.dwdp_size >= 2, ( + f"dwdp_size must be >= 2 when enabled, got {cfg.dwdp_size}" + ) + assert cfg.dwdp_size == cfg.tp_size, ( + f"dwdp_size ({cfg.dwdp_size}) must equal tp_size ({cfg.tp_size})" + ) assert cfg.disaggregation_mode in ( "null", "prefill", ), "DWDP requires --disaggregation-mode null or prefill" - assert ( - not cfg.enable_eplb - ), "EPLB dynamic migration conflicts with static DWDP partitioning" - assert ( - cfg.speculative_algorithm is None - ), "DWDP does not support speculative decoding (MTP/draft workers)" + assert not cfg.enable_eplb, ( + "EPLB dynamic migration conflicts with static DWDP partitioning" + ) + assert cfg.speculative_algorithm is None, ( + "DWDP does not support speculative decoding (MTP/draft workers)" + ) assert cfg.pp_size == 1, "DWDP requires pp_size == 1" - assert ( - not cfg.enable_two_batch_overlap - ), "DWDP's prefetch event protocol does not support two-batch overlap" + assert not cfg.enable_two_batch_overlap, ( + "DWDP's prefetch event protocol does not support two-batch overlap" + ) if cfg.disaggregation_mode == "null": logger.warning( @@ -359,7 +358,9 @@ def handle_elastic_ep(server_args: Any): assert cfg.eplb_algorithm in [ "elasticity_aware", "elasticity_aware_hierarchical", - ], "Elastic EP requires eplb_algorithm to be set to 'auto' or 'elasticity_aware(_hierarchical)'." + ], ( + "Elastic EP requires eplb_algorithm to be set to 'auto' or 'elasticity_aware(_hierarchical)'." + ) assert cfg.pp_size == 1, "PP size should be set to 1 under elastic EP" @@ -370,9 +371,9 @@ def handle_elastic_ep(server_args: Any): mooncake_ib_device=validate_ib_devices(cfg.mooncake_ib_device), ) if cfg.ep_join_mode is not None: - assert ( - cfg.elastic_ep_backend is not None - ), "--elastic-ep-join-mode requires --elastic-ep-backend to be set." + assert cfg.elastic_ep_backend is not None, ( + "--elastic-ep-join-mode requires --elastic-ep-backend to be set." + ) if cfg.ep_join_mode == "scale": assert cfg.node_rank == 1, ( "Elastic EP scale-up requires one joining TP group at " @@ -390,9 +391,9 @@ def handle_elastic_ep(server_args: Any): ) assert cfg.ep_join_rank_offset >= 0, "elastic EP join rank offset must be >= 0." if cfg.max_ep_size is not None: - assert ( - cfg.elastic_ep_backend is not None - ), "--max-ep-size requires --elastic-ep-backend to be set." + assert cfg.elastic_ep_backend is not None, ( + "--max-ep-size requires --elastic-ep-backend to be set." + ) assert cfg.max_ep_size > 0, "--max-ep-size must be a positive integer." scaling_active = ( @@ -407,16 +408,15 @@ def handle_elastic_ep(server_args: Any): ) if scaling_active: resolved = resolved_view(server_args) - assert ( - cfg.elastic_ep_scale_timeout > 0 - ), "--elastic-ep-scale-timeout must be greater than zero." + assert cfg.elastic_ep_scale_timeout > 0, ( + "--elastic-ep-scale-timeout must be greater than zero." + ) assert cfg.tokenizer_worker_num == 1, ( - "Elastic EP runtime scale-up currently requires " - "--tokenizer-worker-num 1." + "Elastic EP runtime scale-up currently requires --tokenizer-worker-num 1." + ) + assert not cfg.use_ray, ( + "Elastic EP runtime scale-up does not support --use-ray." ) - assert ( - not cfg.use_ray - ), "Elastic EP runtime scale-up does not support --use-ray." assert not cfg.enable_elastic_expert_backup, ( "Elastic EP runtime scale-up does not support " "--enable-elastic-expert-backup." diff --git a/python/sglang/srt/arg_groups/pd_disaggregation_hook.py b/python/sglang/srt/arg_groups/pd_disaggregation_hook.py index 7e05b9176..a615e59ea 100644 --- a/python/sglang/srt/arg_groups/pd_disaggregation_hook.py +++ b/python/sglang/srt/arg_groups/pd_disaggregation_hook.py @@ -129,9 +129,9 @@ def handle_pd_disaggregation(server_args: ServerArgs) -> None: ) elif cfg.disaggregation_mode == "prefill": - assert ( - cfg.disaggregation_transfer_backend != "fake" - ), "Prefill server does not support 'fake' as the transfer backend" + assert cfg.disaggregation_transfer_backend != "fake", ( + "Prefill server does not support 'fake' as the transfer backend" + ) if envs.SGLANG_RUST_SERVER.get(): _alias_bootstrap_port_to_api_port(server_args) diff --git a/python/sglang/srt/arg_groups/serving_hook.py b/python/sglang/srt/arg_groups/serving_hook.py index 39c3f63ca..404509bb3 100644 --- a/python/sglang/srt/arg_groups/serving_hook.py +++ b/python/sglang/srt/arg_groups/serving_hook.py @@ -75,7 +75,7 @@ def handle_ssl_validation(server_args: Any): if cfg.enable_http2: if not 0 < cfg.http2_max_concurrent_streams < 2**32: raise ValueError( - "--http2-max-concurrent-streams must be between 1 and " "4294967295." + "--http2-max-concurrent-streams must be between 1 and 4294967295." ) if not 1024 <= cfg.http2_initial_connection_window_size < 2**31: raise ValueError( @@ -343,8 +343,7 @@ def handle_deprecated_args(server_args: Any): ) if cfg.grpc_worker_threads is not None and cfg.grpc_worker_threads < 1: raise ValueError( - "SGLANG_GRPC_WORKER_THREADS " - f"({cfg.grpc_worker_threads}) must be >= 1" + f"SGLANG_GRPC_WORKER_THREADS ({cfg.grpc_worker_threads}) must be >= 1" ) # Native gRPC is incompatible with launch paths it doesn't wire into. @@ -482,8 +481,7 @@ def handle_other_validations(server_args: Any): ) elif resolved_view(server_args).uses_mamba_radix_cache: logger.warning( - "Optimistic prefill does not support models that use " - "mamba radix cache." + "Optimistic prefill does not support models that use mamba radix cache." ) declare_resolution( server_args, @@ -851,8 +849,7 @@ def handle_multimodal_feature_transport(server_args: Any): raise ValueError("--mm-feature-transport=cuda_vmm requires NVIDIA CUDA.") if cfg.pp_size != 1: raise ValueError( - "--mm-feature-transport=cuda_vmm does not support pipeline " - "parallelism." + "--mm-feature-transport=cuda_vmm does not support pipeline parallelism." ) if envs.SGLANG_RUST_SERVER.get(): raise ValueError( diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index b7aefdc4a..3726bf44e 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -559,7 +559,6 @@ def _resolve_dflash_draft_attention_backend(server_args: ServerArgs) -> None: draft_backend = cfg.speculative_draft_attention_backend if draft_backend is None: - draft_backend, _ = attention_backends_of(resolved_view(server_args)) if draft_backend is None: draft_backend = fallback_backend diff --git a/python/sglang/srt/arg_groups/validation_hook.py b/python/sglang/srt/arg_groups/validation_hook.py index e02cccded..eef15b96d 100644 --- a/python/sglang/srt/arg_groups/validation_hook.py +++ b/python/sglang/srt/arg_groups/validation_hook.py @@ -31,9 +31,9 @@ def check_server_args(server_args: Any): # Check parallel size constraints if cfg.ep_join_mode != "scale": - assert ( - cfg.tp_size * cfg.pp_size - ) % cfg.nnodes == 0, "tp_size must be divisible by number of nodes" + assert (cfg.tp_size * cfg.pp_size) % cfg.nnodes == 0, ( + "tp_size must be divisible by number of nodes" + ) assert cfg.pp_max_micro_batch_size is None or cfg.pp_max_micro_batch_size >= 1, ( "pp_max_micro_batch_size must be a positive integer or None (for auto-compute). " @@ -49,18 +49,18 @@ def check_server_args(server_args: Any): ) if cfg.pp_size > 1: - assert ( - cfg.disable_overlap_schedule and cfg.speculative_algorithm is None - ), "Pipeline parallelism is not compatible with overlap schedule, speculative decoding" + assert cfg.disable_overlap_schedule and cfg.speculative_algorithm is None, ( + "Pipeline parallelism is not compatible with overlap schedule, speculative decoding" + ) assert cfg.min_free_slots_delay is None, ( "--min-free-slots-delay is not supported with pipeline " "parallelism: allocatable slots per microbatch are bounded by " "pp-max-micro-batch-size, so the threshold may never be reached" ) - assert not ( - cfg.dp_size > 1 and cfg.nnodes != 1 and not cfg.enable_dp_attention - ), "multi-node data parallel is not supported unless dp attention!" + assert not (cfg.dp_size > 1 and cfg.nnodes != 1 and not cfg.enable_dp_attention), ( + "multi-node data parallel is not supported unless dp attention!" + ) assert cfg.base_gpu_id >= 0, "base_gpu_id must be non-negative" assert cfg.gpu_id_step >= 1, "gpu_id_step must be positive" @@ -102,24 +102,24 @@ def check_server_args(server_args: Any): # Skip validation if chunked prefill is disabled (i.e., size <= 0). # Skip validation if disaggregation mode is decode. if cfg.chunked_prefill_size > 0 and cfg.disaggregation_mode != "decode": - assert ( - cfg.chunked_prefill_size % cfg.page_size == 0 - ), "chunked_prefill_size must be divisible by page_size" + assert cfg.chunked_prefill_size % cfg.page_size == 0, ( + "chunked_prefill_size must be divisible by page_size" + ) # Check pdmux if cfg.enable_pdmux: - assert ( - cfg.pp_size == 1 - ), "PD-Multiplexing is only supported with pipeline parallelism disabled (pp_size=1)." - assert ( - cfg.chunked_prefill_size == -1 - ), "PD-Multiplexing is not compatible with chunked prefill." - assert ( - cfg.disaggregation_mode == "null" - ), "PD-Multiplexing is not compatible with disaggregation mode." - assert ( - cfg.disable_overlap_schedule - ), "PD-Multiplexing is not compatible with overlap schedule." + assert cfg.pp_size == 1, ( + "PD-Multiplexing is only supported with pipeline parallelism disabled (pp_size=1)." + ) + assert cfg.chunked_prefill_size == -1, ( + "PD-Multiplexing is not compatible with chunked prefill." + ) + assert cfg.disaggregation_mode == "null", ( + "PD-Multiplexing is not compatible with disaggregation mode." + ) + assert cfg.disable_overlap_schedule, ( + "PD-Multiplexing is not compatible with overlap schedule." + ) # NOTE: CUDA Green Context may encounter potential issues with CudaGraph on torch 2.7.x – 2.8.x, leading to performance degradation. import torch @@ -143,7 +143,9 @@ def check_server_args(server_args: Any): assert cfg.schedule_policy in [ "fcfs", "lof", - ], f"To use priority scheduling, schedule_policy must be 'fcfs' or 'lof'. '{cfg.schedule_policy}' is not supported." + ], ( + f"To use priority scheduling, schedule_policy must be 'fcfs' or 'lof'. '{cfg.schedule_policy}' is not supported." + ) if cfg.default_priority_value is None: logger.warning( "--default-priority-value is not set while --enable-priority-scheduling is enabled. " @@ -170,14 +172,14 @@ def check_server_args(server_args: Any): run_post_process_pass(server_args, _hisparse_validation) - assert ( - cfg.schedule_conservativeness >= 0 - ), "schedule_conservativeness must be non-negative" + assert cfg.schedule_conservativeness >= 0, ( + "schedule_conservativeness must be non-negative" + ) if cfg.model_impl == "mindspore": - assert ( - get_platform().is_npu - ), "MindSpore model impl is only supported on Ascend npu." + assert get_platform().is_npu, ( + "MindSpore model impl is only supported on Ascend npu." + ) # Check metrics labels if ( @@ -239,43 +241,45 @@ def validate_buckets_rule(arg_name: str, buckets_rule: List[str]): "tse", "default", "custom", - ], f"Unsupported {arg_name} rule type: '{rule}'. Must be one of: 'tse', 'default', 'custom'" + ], ( + f"Unsupported {arg_name} rule type: '{rule}'. Must be one of: 'tse', 'default', 'custom'" + ) if rule == "tse": - assert ( - len(buckets_rule) == 4 - ), f"{arg_name} TSE rule requires exactly 4 parameters: ['tse', middle, base, count], got {len(buckets_rule)}" + assert len(buckets_rule) == 4, ( + f"{arg_name} TSE rule requires exactly 4 parameters: ['tse', middle, base, count], got {len(buckets_rule)}" + ) try: middle = float(buckets_rule[1]) base = float(buckets_rule[2]) count = int(buckets_rule[3]) except (ValueError, IndexError): - assert ( - False - ), f"{arg_name} TSE rule parameters must be: ['tse', , , ]" + assert False, ( + f"{arg_name} TSE rule parameters must be: ['tse', , , ]" + ) assert base > 1, f"{arg_name} TSE base must be larger than 1, got: {base}" assert count > 0, f"{arg_name} TSE count must be positive, got: {count}" assert middle > 0, f"{arg_name} TSE middle must be positive, got: {middle}" elif rule == "default": - assert ( - len(buckets_rule) == 1 - ), f"{arg_name} default rule should only have one parameter: ['default'], got {len(buckets_rule)}" + assert len(buckets_rule) == 1, ( + f"{arg_name} default rule should only have one parameter: ['default'], got {len(buckets_rule)}" + ) elif rule == "custom": - assert ( - len(buckets_rule) >= 2 - ), f"{arg_name} custom rule requires at least one bucket value: ['custom', value1, ...]" + assert len(buckets_rule) >= 2, ( + f"{arg_name} custom rule requires at least one bucket value: ['custom', value1, ...]" + ) try: bucket_values = [float(x) for x in buckets_rule[1:]] except ValueError: assert False, f"{arg_name} custom rule bucket values must be numeric" - assert len(set(bucket_values)) == len( - bucket_values - ), f"{arg_name} custom rule bucket values should not contain duplicates" - assert all( - val >= 0 for val in bucket_values - ), f"{arg_name} custom rule bucket values should be non-negative" + assert len(set(bucket_values)) == len(bucket_values), ( + f"{arg_name} custom rule bucket values should not contain duplicates" + ) + assert all(val >= 0 for val in bucket_values), ( + f"{arg_name} custom rule bucket values should be non-negative" + ) def check_load_publish_args(server_args: Any): diff --git a/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py b/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py index 51a911eb1..5176ff3eb 100644 --- a/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py +++ b/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py @@ -176,9 +176,9 @@ def _matmul_persistent_triton( # Check constraints. assert a.shape[1] == b.shape[0], "Incompatible dimensions" assert a.dtype == b.dtype, "Incompatible dtypes" - assert ( - bias is None or bias.dim() == 1 - ), "Currently assuming bias is 1D, let Horace know if you run into this" + assert bias is None or bias.dim() == 1, ( + "Currently assuming bias is 1D, let Horace know if you run into this" + ) NUM_SMS = get_device_core_count() M, K = a.shape K, N = b.shape @@ -501,9 +501,9 @@ def mean_dim( """ # Validate inputs assert input.is_cuda or input.is_xpu, "Input must be a CUDA or XPU tensor" - assert ( - -input.ndim <= dim < input.ndim - ), f"Invalid dimension {dim} for tensor with {input.ndim} dimensions" + assert -input.ndim <= dim < input.ndim, ( + f"Invalid dimension {dim} for tensor with {input.ndim} dimensions" + ) # Handle negative dim if dim < 0: diff --git a/python/sglang/srt/batch_overlap/operations.py b/python/sglang/srt/batch_overlap/operations.py index df3decc7d..5e8829277 100644 --- a/python/sglang/srt/batch_overlap/operations.py +++ b/python/sglang/srt/batch_overlap/operations.py @@ -191,9 +191,9 @@ class _StateDict: if key == "_data": super().__setattr__(key, value) return - assert ( - key not in self._data - ), f"`{key}` already exist, are you sure you want to override it?" + assert key not in self._data, ( + f"`{key}` already exist, are you sure you want to override it?" + ) self._data[key] = value def __getattr__(self, item): diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index 8fb19f291..6e9a7ba10 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/python/sglang/srt/batch_overlap/two_batch_overlap.py @@ -630,9 +630,9 @@ class TboForwardBatchPreparer: sum_field="extend_num_tokens", ) - assert ( - child_a.extend_num_tokens == half_seq_lens_sum - ), f"{child_a.extend_num_tokens=}, {half_seq_lens_sum=}" + assert child_a.extend_num_tokens == half_seq_lens_sum, ( + f"{child_a.extend_num_tokens=}, {half_seq_lens_sum=}" + ) child_a.seq_lens_cpu = copy.deepcopy(child_a.seq_lens_cpu) child_a.seq_lens_cpu[-1] = ( @@ -674,9 +674,9 @@ class TboForwardBatchPreparer: out_num_token_non_padded: torch.Tensor, out_num_token_non_padded_cpu: Optional[int] = None, ): - assert ( - end_token_index >= start_token_index - ), f"{end_token_index=}, {start_token_index=}, batch={batch}" + assert end_token_index >= start_token_index, ( + f"{end_token_index=}, {start_token_index=}, batch={batch}" + ) num_tokens = batch.input_ids.shape[0] num_seqs = batch.batch_size @@ -688,9 +688,9 @@ class TboForwardBatchPreparer: "out_cache_loc", ]: old_value = getattr(batch, key) - assert ( - old_value.shape[0] == num_tokens - ), f"{key=} {old_value=} {num_tokens=} {batch=}" + assert old_value.shape[0] == num_tokens, ( + f"{key=} {old_value=} {num_tokens=} {batch=}" + ) output_dict[key] = old_value[start_token_index:end_token_index] attention_tp_size = get_parallel().attn_tp_size @@ -736,9 +736,9 @@ class TboForwardBatchPreparer: start_seq_index : min(end_seq_index, len(old_value)) ] continue - assert ( - len(old_value) == num_seqs - ), f"{key=} {old_value=} {num_seqs=} {batch=}" + assert len(old_value) == num_seqs, ( + f"{key=} {old_value=} {num_seqs=} {batch=}" + ) output_dict[key] = old_value[start_seq_index:end_seq_index] spec_info = getattr(batch, "spec_info") diff --git a/python/sglang/srt/beam_search/beam_group.py b/python/sglang/srt/beam_search/beam_group.py index 0552dd109..7ad2c0ea6 100644 --- a/python/sglang/srt/beam_search/beam_group.py +++ b/python/sglang/srt/beam_search/beam_group.py @@ -211,9 +211,9 @@ class BeamGroup: cums = sel.cum_logprobs.tolist() # A parent outside the committed frontier means a tick-gating bug let an # unsynchronized step through; fail rather than build a corrupt DAG. - assert not parents or max(parents) < len( - self.leaves - ), "beam commit consumed an unsynced or misordered step" + assert not parents or max(parents) < len(self.leaves), ( + "beam commit consumed an unsynced or misordered step" + ) for token, parent, cum in zip(tokens, parents, cums): leaf = BeamNode(token, self.leaves[parent]) self.completed.append(CompletedBeam(leaf, cum, new_len, matched_token=None)) diff --git a/python/sglang/srt/compilation/backend.py b/python/sglang/srt/compilation/backend.py index 08b871364..5f40fd296 100644 --- a/python/sglang/srt/compilation/backend.py +++ b/python/sglang/srt/compilation/backend.py @@ -121,15 +121,14 @@ class CompilerManager: ) if runtime_shape is None: logger.debug( - "Directly load the %s-th graph for dynamic shape from %s via " - "handle %s", + "Directly load the %s-th graph for dynamic shape from %s via handle %s", graph_index, self.compiler.name, handle, ) else: logger.debug( - "Directly load the %s-th graph for shape %s from %s via " "handle %s", + "Directly load the %s-th graph for shape %s from %s via handle %s", graph_index, str(runtime_shape), self.compiler.name, @@ -184,7 +183,7 @@ class CompilerManager: ) if runtime_shape is None: logger.debug( - "Store the %s-th graph for dynamic shape from %s via " "handle %s", + "Store the %s-th graph for dynamic shape from %s via handle %s", graph_index, self.compiler.name, handle, @@ -352,9 +351,9 @@ model_tag: str = "backbone" def set_model_tag(tag: str): """Context manager to set the model tag.""" global model_tag - assert ( - tag != model_tag - ), f"Model tag {tag} is the same as the current tag {model_tag}." + assert tag != model_tag, ( + f"Model tag {tag} is the same as the current tag {model_tag}." + ) old_tag = model_tag model_tag = tag try: @@ -364,7 +363,6 @@ def set_model_tag(tag: str): class SGLangBackend: - graph_pool: Any _called: bool = False # the graph we compiled diff --git a/python/sglang/srt/compilation/compiler_interface.py b/python/sglang/srt/compilation/compiler_interface.py index 87255aae4..ced289dfe 100644 --- a/python/sglang/srt/compilation/compiler_interface.py +++ b/python/sglang/srt/compilation/compiler_interface.py @@ -256,9 +256,7 @@ class InductorAdaptor(CompilerInterface): break return inductor_compiled_graph - hijacked_compile_fx_inner = ( - torch._inductor.compile_fx.compile_fx_inner - ) # noqa + hijacked_compile_fx_inner = torch._inductor.compile_fx.compile_fx_inner # noqa elif torch_release >= (2, 6): # function renamed in 2.6 original_load_name = None diff --git a/python/sglang/srt/compilation/cuda_piecewise_backend.py b/python/sglang/srt/compilation/cuda_piecewise_backend.py index f4192b4bc..75ce0e0dc 100644 --- a/python/sglang/srt/compilation/cuda_piecewise_backend.py +++ b/python/sglang/srt/compilation/cuda_piecewise_backend.py @@ -45,7 +45,6 @@ class ConcreteSizeEntry: class CUDAPiecewiseBackend: - def __init__( self, graph: fx.GraphModule, @@ -190,8 +189,9 @@ class CUDAPiecewiseBackend: stack.enter_context(patch("gc.collect", lambda: None)) stack.enter_context(patch("torch.cuda.empty_cache", lambda: None)) # mind-exploding: carefully manage the reference and memory. - with graph_pool_capture_scope(), torch.cuda.graph( - cudagraph, pool=self.graph_pool, stream=stream + with ( + graph_pool_capture_scope(), + torch.cuda.graph(cudagraph, pool=self.graph_pool, stream=stream), ): # `output` is managed by pytorch's cudagraph pool output = entry.runnable(*args) diff --git a/python/sglang/srt/compilation/fix_functionalization.py b/python/sglang/srt/compilation/fix_functionalization.py index afb1a8c12..5fdbea269 100644 --- a/python/sglang/srt/compilation/fix_functionalization.py +++ b/python/sglang/srt/compilation/fix_functionalization.py @@ -119,9 +119,9 @@ class FixFunctionalizationPass(SGLangInductorPass): :param args: If we cannot use kwargs, specify args directly. If an arg is a string, `node.kwargs[arg]` is used. """ # noqa: E501 - assert is_func( - node, auto_functionalized - ), f"node must be auto-functionalized, is {node} instead" + assert is_func(node, auto_functionalized), ( + f"node must be auto-functionalized, is {node} instead" + ) # Create a new call to the original function with graph.inserting_before(node): diff --git a/python/sglang/srt/compilation/inductor_pass.py b/python/sglang/srt/compilation/inductor_pass.py index 6b32cbc83..48c1f4b95 100644 --- a/python/sglang/srt/compilation/inductor_pass.py +++ b/python/sglang/srt/compilation/inductor_pass.py @@ -22,7 +22,6 @@ _pass_context = None class PassContext: - def __init__(self, runtime_shape: Optional[int]): self.runtime_shape = runtime_shape @@ -114,7 +113,6 @@ class CallableInductorPass(InductorPass): class SGLangInductorPass(InductorPass): - def __init__( self, ): @@ -133,7 +131,6 @@ class SGLangInductorPass(InductorPass): class PrinterInductorPass(SGLangInductorPass): - def __init__(self, name: str): super().__init__() self.name = name diff --git a/python/sglang/srt/configs/bailing_hybrid.py b/python/sglang/srt/configs/bailing_hybrid.py index 659b9af35..567c30963 100644 --- a/python/sglang/srt/configs/bailing_hybrid.py +++ b/python/sglang/srt/configs/bailing_hybrid.py @@ -37,7 +37,6 @@ class HybridLayerType(enum.Enum): class BailingHybridConfig(PretrainedConfig): - model_type = "bailing_hybrid" keys_to_ignore_at_inference = ["past_key_values"] @@ -175,9 +174,9 @@ class BailingHybridConfig(PretrainedConfig): layer_type_list.append(HybridLayerType.linear_attention.value) else: # Per-layer schedule: 1 marks a linear-attention layer. - assert ( - len(self.layer_group_size) == self.num_hidden_layers - ), "When layer_group_size is a list, its length must be equal to num_hidden_layers" + assert len(self.layer_group_size) == self.num_hidden_layers, ( + "When layer_group_size is a list, its length must be equal to num_hidden_layers" + ) for l in range(self.num_hidden_layers): if self.layer_group_size[l] == 1: layer_type_list.append(HybridLayerType.linear_attention.value) diff --git a/python/sglang/srt/configs/deepseek_ocr.py b/python/sglang/srt/configs/deepseek_ocr.py index 9a4041773..3cde7cc86 100644 --- a/python/sglang/srt/configs/deepseek_ocr.py +++ b/python/sglang/srt/configs/deepseek_ocr.py @@ -193,7 +193,6 @@ class ImageTransform(object): elif not x.is_floating_point(): x = x.to(torch.float32) if self.normalize: - import torchvision.transforms as T x = T.Normalize(self.mean, self.std)(x) @@ -620,9 +619,9 @@ class DeepseekOCRProcessor(ProcessorMixin): tokenized_str = tokenized_str + [self.eos_id] images_seq_mask = images_seq_mask + [False] - assert len(tokenized_str) == len( - images_seq_mask - ), f"tokenize_with_images func: tokenized_str's length {len(tokenized_str)} is not equal to imags_seq_mask's length {len(images_seq_mask)}" + assert len(tokenized_str) == len(images_seq_mask), ( + f"tokenize_with_images func: tokenized_str's length {len(tokenized_str)} is not equal to imags_seq_mask's length {len(images_seq_mask)}" + ) masked_tokenized_str = [] for token_index in tokenized_str: diff --git a/python/sglang/srt/configs/deepseekvl2.py b/python/sglang/srt/configs/deepseekvl2.py index e8f784258..6aa526b78 100644 --- a/python/sglang/srt/configs/deepseekvl2.py +++ b/python/sglang/srt/configs/deepseekvl2.py @@ -21,8 +21,9 @@ def select_best_resolution(image_size, candidate_resolutions): for width, height in candidate_resolutions: scale = min(width / original_width, height / original_height) - downscaled_width, downscaled_height = int(original_width * scale), int( - original_height * scale + downscaled_width, downscaled_height = ( + int(original_width * scale), + int(original_height * scale), ) effective_resolution = min( downscaled_width * downscaled_height, original_width * original_height @@ -205,9 +206,9 @@ class DeepseekVLV2Processor(ProcessorMixin): images_seq_mask += seq_mask images_spatial_crop += spatial_crop - assert len(tokenized_data) == len( - images_seq_mask - ), f"format_messages_v2: tokenized_str's length {len(tokenized_str)} is not equal to imags_seq_mask's length {len(images_seq_mask)}" + assert len(tokenized_data) == len(images_seq_mask), ( + f"format_messages_v2: tokenized_str's length {len(tokenized_str)} is not equal to imags_seq_mask's length {len(images_seq_mask)}" + ) return ( tokenized_data, @@ -274,9 +275,9 @@ class DeepseekVLV2Processor(ProcessorMixin): - num_image_tokens (List[int]): the number of image tokens """ - assert ( - prompt is None or conversations is None - ), "prompt and conversations cannot be used at the same time." + assert prompt is None or conversations is None, ( + "prompt and conversations cannot be used at the same time." + ) ( tokenized_str, @@ -458,9 +459,9 @@ class DeepseekVLV2Processor(ProcessorMixin): tokenized_str = tokenized_str + [self.eos_id] images_seq_mask = images_seq_mask + [False] - assert len(tokenized_str) == len( - images_seq_mask - ), f"tokenize_with_images func: tokenized_str's length {len(tokenized_str)} is not equal to imags_seq_mask's length {len(images_seq_mask)}" + assert len(tokenized_str) == len(images_seq_mask), ( + f"tokenize_with_images func: tokenized_str's length {len(tokenized_str)} is not equal to imags_seq_mask's length {len(images_seq_mask)}" + ) return tokenized_str, images_list, images_seq_mask, images_spatial_crop @@ -547,7 +548,6 @@ class DeepseekVL2MlpProjectorConfig(PretrainedConfig): class DeepseekV2Config(PretrainedConfig): - model_type = "deepseek_v2" keys_to_ignore_at_inference = ["past_key_values"] diff --git a/python/sglang/srt/configs/inkling.py b/python/sglang/srt/configs/inkling.py index 7c1859f1f..f4f8f115e 100644 --- a/python/sglang/srt/configs/inkling.py +++ b/python/sglang/srt/configs/inkling.py @@ -104,9 +104,9 @@ class InklingModelConfig(PretrainedConfig): mtp_swa_head_dim = swa_head_dim if mtp_local_layer_ids: local_id_set = set(mtp_local_layer_ids) - assert len(local_id_set) == len( - mtp_local_layer_ids - ), f"mtp_local_layer_ids must be unique: {mtp_local_layer_ids}" + assert len(local_id_set) == len(mtp_local_layer_ids), ( + f"mtp_local_layer_ids must be unique: {mtp_local_layer_ids}" + ) assert all(0 <= i < num_nextn_predict_layers for i in local_id_set), ( f"mtp_local_layer_ids must be in [0, {num_nextn_predict_layers}): " f"{mtp_local_layer_ids}" @@ -231,9 +231,9 @@ class InklingModelConfig(PretrainedConfig): if get_exec().comm.enable_scattered_sconv: # Scattered sconv: the attn/mlp output sconvs run on the [T, H/P] # hidden shard, so their conv-state caches shard with them. - assert ( - self.hidden_size % tp_size == 0 - ), f"hidden_size {self.hidden_size} not divisible by attn tp {tp_size}" + assert self.hidden_size % tp_size == 0, ( + f"hidden_size {self.hidden_size} not divisible by attn tp {tp_size}" + ) stream_dim = self.hidden_size // tp_size conv_len = self.sconv_kernel_size - 1 shape = InklingConvStateShape( diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 1682879d3..2a5a4f459 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -1593,7 +1593,6 @@ class ModelConfig: # of an NVFP4/mixed checkpoint) must not be overridden back to the # source format if self.quantization not in REQUANTIZATION_METHODS: - # Detect which checkpoint is it if not preserve_online_draft_quantization: for _, method in QUANTIZATION_METHODS.items(): @@ -2242,12 +2241,14 @@ def get_hybrid_layer_ids( elif "InklingForConditionalGeneration" in model_architectures: local_layer_ids = hf_text_config.local_layer_ids local_layer_id_set = set(local_layer_ids) - assert len(local_layer_id_set) == len( - local_layer_ids - ), f"Inkling local_layer_ids must be unique: {local_layer_ids}" + assert len(local_layer_id_set) == len(local_layer_ids), ( + f"Inkling local_layer_ids must be unique: {local_layer_ids}" + ) assert all( 0 <= layer_id < num_hidden_layers for layer_id in local_layer_id_set - ), f"Inkling local_layer_ids must be in [0, {num_hidden_layers}): {local_layer_ids}" + ), ( + f"Inkling local_layer_ids must be in [0, {num_hidden_layers}): {local_layer_ids}" + ) swa_attention_layer_ids = [ i for i in range(num_hidden_layers) if i in local_layer_id_set ] diff --git a/python/sglang/srt/configs/muse_glimmer.py b/python/sglang/srt/configs/muse_glimmer.py index 468d0f89b..549a95caa 100644 --- a/python/sglang/srt/configs/muse_glimmer.py +++ b/python/sglang/srt/configs/muse_glimmer.py @@ -26,7 +26,6 @@ _ARCH = "muse-glimmer" class MuseGlimmerAssistantConfig(PretrainedConfig): - model_type = "muse_glimmer_assistant" is_causal = False # The DFlash draft has no head; draft_worker_common borrows the target's. @@ -34,7 +33,6 @@ class MuseGlimmerAssistantConfig(PretrainedConfig): class MuseGlimmerVisionConfig(PretrainedConfig): - model_type = "muse_glimmer_vision" def __init__( diff --git a/python/sglang/srt/configs/nano_nemotron_vl.py b/python/sglang/srt/configs/nano_nemotron_vl.py index 888f94a1b..30c6b4390 100644 --- a/python/sglang/srt/configs/nano_nemotron_vl.py +++ b/python/sglang/srt/configs/nano_nemotron_vl.py @@ -24,9 +24,9 @@ from sglang.srt.multimodal.internvl_utils import IMAGENET_MEAN, IMAGENET_STD def float_triplet(seq: Any): a, b, c = tuple(seq) - assert ( - isinstance(a, float) and isinstance(b, float) and isinstance(c, float) - ), "expected three floats" + assert isinstance(a, float) and isinstance(b, float) and isinstance(c, float), ( + "expected three floats" + ) return a, b, c diff --git a/python/sglang/srt/configs/olmo3.py b/python/sglang/srt/configs/olmo3.py index 864034219..fcfce4f82 100644 --- a/python/sglang/srt/configs/olmo3.py +++ b/python/sglang/srt/configs/olmo3.py @@ -28,7 +28,6 @@ class Olmo3LayerType(enum.Enum): class Olmo3Config(PretrainedConfig): - model_type = "olmo3" keys_to_ignore_at_inference = ["past_key_values"] diff --git a/python/sglang/srt/configs/qwen3_omni.py b/python/sglang/srt/configs/qwen3_omni.py index 8baea8923..cb7f2cd2e 100644 --- a/python/sglang/srt/configs/qwen3_omni.py +++ b/python/sglang/srt/configs/qwen3_omni.py @@ -235,7 +235,6 @@ class Qwen3OmniMoeThinkerConfig(PretrainedConfig): class Qwen3OmniMoeTalkerCodePredictorConfig(PretrainedConfig): - model_type = "qwen3_omni_moe_talker_code_predictor" keys_to_ignore_at_inference = ["past_key_values"] @@ -325,7 +324,6 @@ class Qwen3OmniMoeTalkerCodePredictorConfig(PretrainedConfig): class Qwen3OmniMoeTalkerTextConfig(PretrainedConfig): - model_type = "qwen3_omni_moe_talker_text" keys_to_ignore_at_inference = ["past_key_values"] @@ -415,7 +413,6 @@ class Qwen3OmniMoeTalkerTextConfig(PretrainedConfig): class Qwen3OmniMoeTalkerConfig(PretrainedConfig): - sub_configs = { "code_predictor_config": Qwen3OmniMoeTalkerCodePredictorConfig, "text_config": Qwen3OmniMoeTalkerTextConfig, @@ -486,7 +483,6 @@ class Qwen3OmniMoeTalkerConfig(PretrainedConfig): class Qwen3OmniMoeCode2WavConfig(PretrainedConfig): - def __init__( self, codebook_size=2048, @@ -538,7 +534,6 @@ class Qwen3OmniMoeCode2WavConfig(PretrainedConfig): class Qwen3OmniMoeConfig(PretrainedConfig): - model_type = "qwen3_omni_moe" sub_configs = { "thinker_config": Qwen3OmniMoeThinkerConfig, diff --git a/python/sglang/srt/configs/update_config.py b/python/sglang/srt/configs/update_config.py index e0f0de2b0..fa2c803b2 100644 --- a/python/sglang/srt/configs/update_config.py +++ b/python/sglang/srt/configs/update_config.py @@ -38,9 +38,9 @@ def get_moe_padding_size(weight_block_size): 2, ], "Only len(weight_block_size) in [1, 2] is supported" if len(weight_block_size) == 2: - assert ( - weight_block_size[0] == weight_block_size[1] - ), "Only weight_block_size[0] == weight_block_size[1] is supported" + assert weight_block_size[0] == weight_block_size[1], ( + "Only weight_block_size[0] == weight_block_size[1] is supported" + ) return weight_block_size[0] return DEFAULT_MOE_PADDING_SIZE @@ -238,7 +238,6 @@ def adjust_config_with_unaligned_cpu_tp( model_config.num_attention_heads % tp_size != 0 or model_config.get_total_num_kv_heads() % tp_size != 0 ): - if hasattr(model_config.hf_config, "qk_nope_head_dim") and hasattr( model_config.hf_config, "qk_rope_head_dim" ): diff --git a/python/sglang/srt/configs/zaya.py b/python/sglang/srt/configs/zaya.py index 3ff25afe3..2c6dd8a30 100644 --- a/python/sglang/srt/configs/zaya.py +++ b/python/sglang/srt/configs/zaya.py @@ -167,9 +167,9 @@ class ZayaConfig(PretrainedConfig): self.head_dim = head_dim self.kv_channels = kv_channels if kv_channels is not None else head_dim assert self.head_dim is not None, "head_dim is required for ZayaConfig" - assert ( - self.num_query_groups == num_key_value_heads - ), "num_query_groups must equal num_key_value_heads for ZAYA1 checkpoints" + assert self.num_query_groups == num_key_value_heads, ( + "num_query_groups must equal num_key_value_heads for ZAYA1 checkpoints" + ) self.num_key_value_heads = num_key_value_heads self.activation_func = activation_func self.max_position_embeddings = max_position_embeddings @@ -266,7 +266,6 @@ class ZayaConfig(PretrainedConfig): # equals the global TP group (DP attention is unsupported), so the two # are always identical in practice. try: - tp_size = get_parallel().tp_size except (AssertionError, RuntimeError, ValueError): tp_size = 1 diff --git a/python/sglang/srt/connector/base_connector.py b/python/sglang/srt/connector/base_connector.py index c9a1c36e2..bbbcea481 100644 --- a/python/sglang/srt/connector/base_connector.py +++ b/python/sglang/srt/connector/base_connector.py @@ -73,7 +73,6 @@ class BaseConnector(ABC): class BaseKVConnector(BaseConnector): - @abstractmethod def get(self, key: str) -> Optional[torch.Tensor]: raise NotImplementedError() diff --git a/python/sglang/srt/connector/redis.py b/python/sglang/srt/connector/redis.py index cb1db3f7c..7824571a7 100644 --- a/python/sglang/srt/connector/redis.py +++ b/python/sglang/srt/connector/redis.py @@ -14,7 +14,6 @@ logger = logging.getLogger(__name__) class RedisConnector(BaseKVConnector): - def __init__(self, url: str): import redis diff --git a/python/sglang/srt/connector/remote_instance.py b/python/sglang/srt/connector/remote_instance.py index 318c362af..875063bd4 100644 --- a/python/sglang/srt/connector/remote_instance.py +++ b/python/sglang/srt/connector/remote_instance.py @@ -14,11 +14,10 @@ logger = logging.getLogger(__name__) class RemoteInstanceConnector(BaseConnector): - def __init__(self, url: str, device: torch.device = "cpu"): - assert ( - device.type == "cuda" or device.type == "npu" - ), "RemoteInstanceConnector only supports cuda device." + assert device.type == "cuda" or device.type == "npu", ( + "RemoteInstanceConnector only supports cuda device." + ) super().__init__(url) self.url = url self.device = device @@ -31,12 +30,12 @@ class RemoteInstanceConnector(BaseConnector): group_rank: int = 1, world_size: int = 2, ): - assert ( - self.device.type == "cuda" or self.device.type == "npu" - ), "RemoteInstanceConnector only supports cuda device." - assert ( - gpu_id != -1 and tp_rank != -1 - ), "gpu_id and tp_rank must be specified for RemoteInstanceConnector. " + assert self.device.type == "cuda" or self.device.type == "npu", ( + "RemoteInstanceConnector only supports cuda device." + ) + assert gpu_id != -1 and tp_rank != -1, ( + "gpu_id and tp_rank must be specified for RemoteInstanceConnector. " + ) self.device_id = torch.device(self.device.type, gpu_id) diff --git a/python/sglang/srt/connector/s3.py b/python/sglang/srt/connector/s3.py index 7bef8f5d5..e9955aff4 100644 --- a/python/sglang/srt/connector/s3.py +++ b/python/sglang/srt/connector/s3.py @@ -67,7 +67,6 @@ def list_files( class S3Connector(BaseFileConnector): - def __init__(self, url: str) -> None: import boto3 diff --git a/python/sglang/srt/connector/serde/safe_serde.py b/python/sglang/srt/connector/serde/safe_serde.py index 3e75f9bfc..2ae04a2cd 100644 --- a/python/sglang/srt/connector/serde/safe_serde.py +++ b/python/sglang/srt/connector/serde/safe_serde.py @@ -9,7 +9,6 @@ from sglang.srt.connector.serde.serde import Deserializer, Serializer class SafeSerializer(Serializer): - def __init__(self): super().__init__() @@ -18,7 +17,6 @@ class SafeSerializer(Serializer): class SafeDeserializer(Deserializer): - def __init__(self): # TODO: dtype options super().__init__(torch.float32) diff --git a/python/sglang/srt/connector/serde/serde.py b/python/sglang/srt/connector/serde/serde.py index 3d6f804d7..20cfcbf65 100644 --- a/python/sglang/srt/connector/serde/serde.py +++ b/python/sglang/srt/connector/serde/serde.py @@ -7,7 +7,6 @@ import torch class Serializer(ABC): - @abstractmethod def to_bytes(self, t: torch.Tensor) -> bytes: """ @@ -25,7 +24,6 @@ class Serializer(ABC): class Deserializer(metaclass=abc.ABCMeta): - def __init__(self, dtype): self.dtype = dtype diff --git a/python/sglang/srt/constrained/base_grammar_backend.py b/python/sglang/srt/constrained/base_grammar_backend.py index d6fdbf211..716f1a030 100644 --- a/python/sglang/srt/constrained/base_grammar_backend.py +++ b/python/sglang/srt/constrained/base_grammar_backend.py @@ -56,7 +56,6 @@ class GrammarRow(NamedTuple): class BaseGrammarObject: - def __init__(self): self._finished = False self.grammar_stats = None diff --git a/python/sglang/srt/constrained/llguidance_backend.py b/python/sglang/srt/constrained/llguidance_backend.py index 1f5bb72ad..658d96b39 100644 --- a/python/sglang/srt/constrained/llguidance_backend.py +++ b/python/sglang/srt/constrained/llguidance_backend.py @@ -113,7 +113,6 @@ def _create_llguidance_tokenizer( class GuidanceGrammar(BaseGrammarObject): - def __init__( self, llguidance_tokenizer: LLTokenizer, @@ -226,7 +225,6 @@ class GuidanceGrammar(BaseGrammarObject): class GuidanceBackend(BaseGrammarBackend): - def __init__( self, tokenizer, diff --git a/python/sglang/srt/constrained/xgrammar_backend.py b/python/sglang/srt/constrained/xgrammar_backend.py index 58b3f1540..1d79b8c33 100644 --- a/python/sglang/srt/constrained/xgrammar_backend.py +++ b/python/sglang/srt/constrained/xgrammar_backend.py @@ -71,7 +71,6 @@ def _allocate_token_bitmask(vocab_size: int, batch_size: int) -> torch.Tensor: class XGrammarGrammar(BaseGrammarObject): - def __init__( self, matcher: GrammarMatcher, diff --git a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/aux_plugins.py b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/aux_plugins.py index 8e497e9ec..7efa040aa 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/aux_plugins.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/aux_plugins.py @@ -135,15 +135,15 @@ class _SGLangPlugin(_AuxFrameworkPlugin): seq_lens = step_data["seq_lens"] rids_raw = step_data.get("rids") - assert isinstance( - input_ids, torch.Tensor - ), f"input_ids: expected Tensor, got {type(input_ids)}" - assert isinstance( - positions, torch.Tensor - ), f"positions: expected Tensor, got {type(positions)}" - assert isinstance( - seq_lens, torch.Tensor - ), f"seq_lens: expected Tensor, got {type(seq_lens)}" + assert isinstance(input_ids, torch.Tensor), ( + f"input_ids: expected Tensor, got {type(input_ids)}" + ) + assert isinstance(positions, torch.Tensor), ( + f"positions: expected Tensor, got {type(positions)}" + ) + assert isinstance(seq_lens, torch.Tensor), ( + f"seq_lens: expected Tensor, got {type(seq_lens)}" + ) seq_lens_list: list[int] = seq_lens.tolist() num_seqs: int = len(seq_lens_list) diff --git a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/executor.py b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/executor.py index 4cfcb4982..af7c92e3d 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/executor.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/executor.py @@ -108,7 +108,11 @@ def _build_bs_collapse_pattern( lhs: str = " ".join(names) # type: ignore[arg-type] - rhs_names: list[str] = list(names[:lo]) + [f"({BATCH_DIM_NAME} {SEQ_DIM_NAME})"] + list(names[hi + 1 :]) # type: ignore[misc] + rhs_names: list[str] = ( + list(names[:lo]) + + [f"({BATCH_DIM_NAME} {SEQ_DIM_NAME})"] + + list(names[hi + 1 :]) + ) # type: ignore[misc] rhs: str = " ".join(rhs_names) new_names: list[str | None] = ( diff --git a/python/sglang/srt/debug_utils/comparator/output_formatter.py b/python/sglang/srt/debug_utils/comparator/output_formatter.py index 37b1afe12..3b43b03c3 100644 --- a/python/sglang/srt/debug_utils/comparator/output_formatter.py +++ b/python/sglang/srt/debug_utils/comparator/output_formatter.py @@ -283,10 +283,7 @@ def _format_non_tensor_rich_body( target_val: str = escape(record.target_value) if record.values_equal: - return ( - f"═ {name}{suffix} = {baseline_val} " - f"({record.baseline_type}) [green]✓[/]" - ) + return f"═ {name}{suffix} = {baseline_val} ({record.baseline_type}) [green]✓[/]" return ( f"═ [bold red]{name}{suffix}[/]\n" f" baseline = {baseline_val} ({record.baseline_type})\n" diff --git a/python/sglang/srt/debug_utils/cuda_coredump.py b/python/sglang/srt/debug_utils/cuda_coredump.py index fa6daa44d..3ef9ea248 100644 --- a/python/sglang/srt/debug_utils/cuda_coredump.py +++ b/python/sglang/srt/debug_utils/cuda_coredump.py @@ -86,7 +86,7 @@ def report(): if not coredump_files: return - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"CUDA coredump(s) detected ({len(coredump_files)} file(s)):") for f in coredump_files: size_mb = os.path.getsize(f) / (1024 * 1024) @@ -98,7 +98,7 @@ def report(): repo = os.environ.get("GITHUB_REPOSITORY", "sgl-project/sglang") print(f"Download from CI: gh run download {run_id} --repo {repo}") - print(f"{'='*60}\n") + print(f"{'=' * 60}\n") # Auto-inject CUDA coredump env vars at import time. diff --git a/python/sglang/srt/debug_utils/dump_comparator.py b/python/sglang/srt/debug_utils/dump_comparator.py index e052ef7a1..8fe8d71b0 100644 --- a/python/sglang/srt/debug_utils/dump_comparator.py +++ b/python/sglang/srt/debug_utils/dump_comparator.py @@ -152,7 +152,7 @@ def check_tensor_pair( value_baseline = fn(x_baseline).item() value_target = fn(x_target).item() print( - f"[{name}] {value_baseline :.4f} vs {value_target:.4f} (diff: {value_target - value_baseline:.4f})" + f"[{name}] {value_baseline:.4f} vs {value_target:.4f} (diff: {value_target - value_baseline:.4f})" ) if x_baseline.shape != x_target.shape: diff --git a/python/sglang/srt/debug_utils/dump_loader.py b/python/sglang/srt/debug_utils/dump_loader.py index f35a455c2..d37d3f333 100644 --- a/python/sglang/srt/debug_utils/dump_loader.py +++ b/python/sglang/srt/debug_utils/dump_loader.py @@ -77,9 +77,9 @@ class DumpLoader: step = dumper._state.step conditions = dict(name=name, step=step, **kwargs) row = find_row(self._df, conditions=conditions) - assert ( - row is not None - ), f"DumpLoader cannot find row given query {name=} {kwargs=} {self._directory=}" + assert row is not None, ( + f"DumpLoader cannot find row given query {name=} {kwargs=} {self._directory=}" + ) path = self._directory / row["filename"] output = torch.load(path, weights_only=False) diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index 906b980c1..111217f10 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/python/sglang/srt/debug_utils/dumper.py @@ -179,9 +179,9 @@ class DumperConfig(_BaseConfig): f"grafter_role must be 'baseline' or 'target' when grafter_enable=True, " f"got {self.grafter_role!r}" ) - assert ( - self.grafter_master_address - ), "grafter_master_address must be set when grafter_enable=True" + assert self.grafter_master_address, ( + "grafter_master_address must be set when grafter_enable=True" + ) assert self.grafter_master_port > 0, ( f"grafter_master_port must be a positive port when grafter_enable=True, " f"got {self.grafter_master_port}" @@ -996,9 +996,9 @@ class _Grafter: return cfg = self._config - assert ( - dist.is_initialized() - ), "[Grafter] default torch.distributed must be initialized" + assert dist.is_initialized(), ( + "[Grafter] default torch.distributed must be initialized" + ) role = _GraftRole(cfg.grafter_role) local_world = dist.get_world_size() local_rank = dist.get_rank() @@ -1795,7 +1795,6 @@ class _SGLangPlugin(_FrameworkPlugin): return None try: - args = get_server_args() if args is None: return None diff --git a/python/sglang/srt/debug_utils/tensor_dump_forward_hook.py b/python/sglang/srt/debug_utils/tensor_dump_forward_hook.py index 66a266075..cd841b518 100644 --- a/python/sglang/srt/debug_utils/tensor_dump_forward_hook.py +++ b/python/sglang/srt/debug_utils/tensor_dump_forward_hook.py @@ -158,7 +158,7 @@ def register_forward_hook_for_model( model_top_level_module_matched, _ = tensor_dumper._add_hook_recursive( model, "", top_level_module_name, layers_module_name ) - assert ( - model_top_level_module_matched - ), f"model should have a module named {top_level_module_name}" + assert model_top_level_module_matched, ( + f"model should have a module named {top_level_module_name}" + ) return tensor_dumper diff --git a/python/sglang/srt/disaggregation/ascend/transfer_engine.py b/python/sglang/srt/disaggregation/ascend/transfer_engine.py index 23d18fcbf..e8ecb2105 100644 --- a/python/sglang/srt/disaggregation/ascend/transfer_engine.py +++ b/python/sglang/srt/disaggregation/ascend/transfer_engine.py @@ -22,7 +22,6 @@ logger = logging.getLogger(__name__) class AscendTransferEngine(MooncakeTransferEngine): - def __init__( self, hostname: str, diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index ef028be35..163f45fee 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -1010,9 +1010,9 @@ class CommonKVManager(BaseKVManager): """ start_layer = self.kv_args.prefill_start_layer end_layer = getattr(self.kv_args, "prefill_end_layer", None) - assert ( - end_layer is not None - ), "KVArgs.prefill_end_layer must be set when using compressed-MLA PD with PP" + assert end_layer is not None, ( + "KVArgs.prefill_end_layer must be set when using compressed-MLA PD with PP" + ) c4_full = sum(1 for r in mla_ratios if r == 4) c128_full = sum(1 for r in mla_ratios if r == 128) @@ -1066,8 +1066,7 @@ class CommonKVManager(BaseKVManager): list(dst_kv_ptrs[swa_s:swa_e]) + list( dst_kv_ptrs[ - compress_section_start - + c4_off_s : compress_section_start + compress_section_start + c4_off_s : compress_section_start + c4_off_e ] ) diff --git a/python/sglang/srt/disaggregation/common/staging_buffer.py b/python/sglang/srt/disaggregation/common/staging_buffer.py index baf0d42f4..7f3e97c1f 100644 --- a/python/sglang/srt/disaggregation/common/staging_buffer.py +++ b/python/sglang/srt/disaggregation/common/staging_buffer.py @@ -146,7 +146,7 @@ class StagingBuffer: self.data_ptr = self.buffer.data_ptr() logger.info( - f"StagingBuffer allocated: {size_bytes / (1024*1024):.1f} MB " + f"StagingBuffer allocated: {size_bytes / (1024 * 1024):.1f} MB " f"on {device}, method={alloc_method}, ptr=0x{self.data_ptr:x}" ) @@ -207,7 +207,7 @@ class StagingAllocator: logger.info( f"StagingAllocator (ring+overcommit): " - f"{total_size_bytes / (1024*1024):.1f} MB " + f"{total_size_bytes / (1024 * 1024):.1f} MB " f"on {device}, ptr=0x{self.base_ptr:x}" ) diff --git a/python/sglang/srt/disaggregation/common/staging_handler.py b/python/sglang/srt/disaggregation/common/staging_handler.py index d9b06eeff..a9ae31b85 100644 --- a/python/sglang/srt/disaggregation/common/staging_handler.py +++ b/python/sglang/srt/disaggregation/common/staging_handler.py @@ -301,7 +301,7 @@ class DecodeStagingHandler: receiver = self._room_to_receiver.get(room) if receiver is None: logger.warning( - "Staging chunk arrived for unregistered room=%s chunk=%d, " "skipping", + "Staging chunk arrived for unregistered room=%s chunk=%d, skipping", room, chunk_idx, ) @@ -901,9 +901,9 @@ class StagingManagerMixin: room = int(msg[1].decode("ascii")) session_id = msg[4].decode("ascii") handler = self._staging_handler - assert ( - handler is not None - ), "STAGING_REQ received before staging handler initialized" + assert handler is not None, ( + "STAGING_REQ received before staging handler initialized" + ) decode_req = handler._room_to_decode_req.get(room) if decode_req is None: logger.warning( diff --git a/python/sglang/srt/disaggregation/common/utils.py b/python/sglang/srt/disaggregation/common/utils.py index c42f6f65c..1571c3d19 100644 --- a/python/sglang/srt/disaggregation/common/utils.py +++ b/python/sglang/srt/disaggregation/common/utils.py @@ -40,7 +40,7 @@ def pack_list_of_buffers(buffers: List[bytes]) -> bytes: if not buffers: return b"" n = len(buffers) - header = struct.pack(f"<{n+1}I", n, *(len(b) for b in buffers)) + header = struct.pack(f"<{n + 1}I", n, *(len(b) for b in buffers)) return header + b"".join(buffers) @@ -64,7 +64,7 @@ def pack_int_lists(lists, fmt: str) -> bytes: def unpack_int_lists(buf: bytes, fmt: str) -> List[List[int]]: width = struct.calcsize(fmt) return [ - list(struct.unpack(f"<{len(b)//width}{fmt}", b)) + list(struct.unpack(f"<{len(b) // width}{fmt}", b)) for b in unpack_list_of_buffers(buf) ] diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 4cda16dcd..9a63da8ac 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -190,9 +190,9 @@ class DecodeReqToTokenPool: # Indices of reqs that already have a req_pool_idx and will reuse # their existing slot (e.g. chunked prefill continuing across chunks). reusing = [i for i, r in enumerate(reqs) if r.kv.holds_kv] - assert all( - reqs[i].kv.kv_allocated_len > 0 for i in reusing - ), "a reused row must carry allocated KV" + assert all(reqs[i].kv.kv_allocated_len > 0 for i in reusing), ( + "a reused row must carry allocated KV" + ) need_size = len(reqs) - len(reusing) if need_size > len(self.free_slots): @@ -1766,9 +1766,9 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): req_pool_indices = self.req_to_token_pool.alloc([req]) - assert ( - req_pool_indices is not None - ), "req_pool_indices is full! There is a bug in memory estimation." + assert req_pool_indices is not None, ( + "req_pool_indices is full! There is a bug in memory estimation." + ) fill_len = self._pre_alloc_fill_len(req) req.kv.kv_committed_len = fill_len @@ -2202,9 +2202,9 @@ class DecodeTransferQueue(DecodeHiCacheTransferMixin): ].tolist() ) if decode_req.req.return_sampling_mask: - assert ( - output_token_sampling_mask_idx is not None - ), "sampling mask buffer disabled on decode side" + assert output_token_sampling_mask_idx is not None, ( + "sampling mask buffer disabled on decode side" + ) sampling_mask_len = int(output_token_sampling_mask_len[0].item()) if sampling_mask_len < 0: decode_req.req.output_token_sampling_mask.append(None) diff --git a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py index 836c93332..07a786f76 100644 --- a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py +++ b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py @@ -19,7 +19,6 @@ if TYPE_CHECKING: class ScheduleBatchDisaggregationDecodeMixin: - def prepare_for_prebuilt(self: ScheduleBatch): """ Prepare a prebuilt extend by populate metadata @@ -49,18 +48,18 @@ class ScheduleBatchDisaggregationDecodeMixin: chunk = self.req_to_token_pool.req_to_token[req.kv.req_pool_idx][ pre_len : pre_len + req.extend_range.length ] - assert ( - offset + req.extend_range.length <= total_size - ), f"Exceeds total size: offset={offset}, req.extend_range.length={req.extend_range.length}, total_size={total_size}" + assert offset + req.extend_range.length <= total_size, ( + f"Exceeds total size: offset={offset}, req.extend_range.length={req.extend_range.length}, total_size={total_size}" + ) out_cache_loc[offset : offset + req.extend_range.length] = chunk offset += req.extend_range.length seq_len = len(req.origin_input_ids) + max(0, len(req.output_ids) - 1) seq_lens.append(seq_len) if len(req.output_ids) == 0: - assert ( - seq_len - pre_len == req.extend_range.length - ), f"seq_len={seq_len}, pre_len={pre_len}, req.extend_range.length={req.extend_range.length}" + assert seq_len - pre_len == req.extend_range.length, ( + f"seq_len={seq_len}, pre_len={pre_len}, req.extend_range.length={req.extend_range.length}" + ) if not req.retracted_stain: # Clamp to avoid double-counting: already_computed is seeded from diff --git a/python/sglang/srt/disaggregation/encoder/receiver.py b/python/sglang/srt/disaggregation/encoder/receiver.py index fd5a8564d..a1ef1d6ce 100644 --- a/python/sglang/srt/disaggregation/encoder/receiver.py +++ b/python/sglang/srt/disaggregation/encoder/receiver.py @@ -1465,8 +1465,7 @@ async def _extract_encoder_error(responses, endpoint, context, encode_requests=N if isinstance(resp, asyncio.TimeoutError): timeout_val = envs.SGLANG_ENCODER_HTTP_TIMEOUT.get() logger.error( - f"Encoder {endpoint} timeout ({timeout_val}s) for {ctx} " - f"(request {i})" + f"Encoder {endpoint} timeout ({timeout_val}s) for {ctx} (request {i})" ) return f"Encoder {endpoint} timeout ({timeout_val}s)" if isinstance(resp, Exception): @@ -1692,9 +1691,9 @@ def _view_pool_buffer_by_modality(raw_buffer, embedding_data, dtype): if info is None: mod_info[mod] = [start, end, shape[0], shape[1]] else: - assert ( - info[3] == shape[1] - ), f"hidden_dim mismatch in modality {mod}: {info[3]} vs {shape[1]}" + assert info[3] == shape[1], ( + f"hidden_dim mismatch in modality {mod}: {info[3]} vs {shape[1]}" + ) assert info[1] == start, f"non-contiguous parts in modality {mod}" info[1] = end info[2] += shape[0] diff --git a/python/sglang/srt/disaggregation/encoder/runtime.py b/python/sglang/srt/disaggregation/encoder/runtime.py index f30623084..f8ee1250a 100644 --- a/python/sglang/srt/disaggregation/encoder/runtime.py +++ b/python/sglang/srt/disaggregation/encoder/runtime.py @@ -892,8 +892,7 @@ class DPDispatcher: ) self._listener_failed = True self._fail_all_pending( - "encoder DP result listener stopped after repeated " - "recv errors", + "encoder DP result listener stopped after repeated recv errors", "ResultListenerStopped", ) return @@ -1301,8 +1300,7 @@ async def _dp_worker_handle_request( # Error envelope, not 200 + phantom count: the decoder must # fail fast instead of waiting for a ZMQ ack that never comes. raise MMError( - f"no staged embedding for /send req_id={req_id} " - f"(already released)" + f"no staged embedding for /send req_id={req_id} (already released)" ) # Releasing on the first /send breaks decoder TP > 1. No count means # a pre-refcount decoder: stay eager rather than pin until the sweep. diff --git a/python/sglang/srt/disaggregation/encoder/server.py b/python/sglang/srt/disaggregation/encoder/server.py index 9b63cf7e0..f76cef67a 100644 --- a/python/sglang/srt/disaggregation/encoder/server.py +++ b/python/sglang/srt/disaggregation/encoder/server.py @@ -950,9 +950,10 @@ class MMEncoder: modality_str = modality.name.lower() preprocess_start = time.perf_counter() try: - preprocess_result, items_per_req = ( - await self.preprocessor.process_batch_mm_items(requests, modality) - ) + ( + preprocess_result, + items_per_req, + ) = await self.preprocessor.process_batch_mm_items(requests, modality) except NotImplementedError as e: raise InternalError(f"Not implemented error: {str(e)}") except Exception as e: diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index afcc0ac47..2c03adead 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -162,8 +162,8 @@ class KVArgsRegisterInfo: endpoint=msg[1].decode("ascii"), dst_port=int(msg[2].decode("ascii")), mooncake_session_id=msg[3].decode("ascii"), - dst_kv_ptrs=list(struct.unpack(f"{len(msg[4])//8}Q", msg[4])), - dst_aux_ptrs=list(struct.unpack(f"{len(msg[5])//8}Q", msg[5])), + dst_kv_ptrs=list(struct.unpack(f"{len(msg[4]) // 8}Q", msg[4])), + dst_aux_ptrs=list(struct.unpack(f"{len(msg[5]) // 8}Q", msg[5])), dst_state_data_ptrs=unpack_int_lists(msg[6], "Q"), dst_tp_rank=int(msg[7].decode("ascii")), dst_attn_tp_size=int(msg[8].decode("ascii")), @@ -2152,9 +2152,9 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager): num_pages = int(msg[4].decode("ascii")) session_id = msg[5].decode("ascii") handler = self._staging_handler - assert ( - handler is not None - ), "CHUNK_READY received before staging handler initialized" + assert handler is not None, ( + "CHUNK_READY received before staging handler initialized" + ) handler.handle_chunk_arrived( room, chunk_idx, @@ -2339,7 +2339,6 @@ class MooncakeFailureExceptionMixin: class MooncakeKVSender(MooncakeFailureExceptionMixin, CommonKVSender): - def __init__( self, mgr: MooncakeKVManager, diff --git a/python/sglang/srt/disaggregation/mori/conn.py b/python/sglang/srt/disaggregation/mori/conn.py index 2eda6fb81..1069d3bc3 100644 --- a/python/sglang/srt/disaggregation/mori/conn.py +++ b/python/sglang/srt/disaggregation/mori/conn.py @@ -1722,7 +1722,6 @@ class MoriKVSender(CommonKVSender): class MoriKVReceiver(CommonKVReceiver): - def __init__( self, mgr: MoriKVManager, diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index f4f49c3ba..aebc6ad75 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -708,9 +708,9 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager): ) prep_handle = self.agent.prep_xfer_dlist(peer_name, np.vstack(arrays), mem_kind) - assert ( - prep_handle is not None - ), f"prep_xfer_dlist returned None for peer '{peer_name}'" + assert prep_handle is not None, ( + f"prep_xfer_dlist returned None for peer '{peer_name}'" + ) return prep_handle def _init_equal_tp_prep_handle( @@ -859,9 +859,9 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager): ] ) src_handle = self.agent.prep_xfer_dlist("", src_array, src_mem_kind) - assert ( - src_handle is not None - ), f"prep_xfer_dlist returned None for slice src (decode_tp_size={decode_tp_size})" + assert src_handle is not None, ( + f"prep_xfer_dlist returned None for slice src (decode_tp_size={decode_tp_size})" + ) self.prep_handle_slice_src = ( src_handle, num_groups, @@ -896,9 +896,9 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager): ] ) dst_handle = self.agent.prep_xfer_dlist(peer_name, dst_array, dst_mem_kind) - assert ( - dst_handle is not None - ), f"prep_xfer_dlist returned None for slice dst for peer '{peer_name}'" + assert dst_handle is not None, ( + f"prep_xfer_dlist returned None for slice dst for peer '{peer_name}'" + ) self.prep_handles_slice_dst[peer_name] = ( dst_handle, num_slots_dst, @@ -1331,8 +1331,7 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager): or not self.kv_args.kv_data_ptrs ): aux_notif += ( - f"_nokv_{self.transfer_source_rank}" - f"_{kv_chunk.chunk_id}" + f"_nokv_{self.transfer_source_rank}_{kv_chunk.chunk_id}" ) aux_xfer_handle = self.send_aux( req.agent_name, @@ -2097,9 +2096,9 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager): ): """Transfer Mamba states via RDMA.""" assert len(prefill_state_indices) == 1, "Mamba should have single state index" - assert len(dst_state_indices) == len( - prefill_state_indices - ), "State indices count mismatch between Prefill and Decode" + assert len(dst_state_indices) == len(prefill_state_indices), ( + "State indices count mismatch between Prefill and Decode" + ) src_addrs = [] dst_addrs = [] @@ -2556,8 +2555,10 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager): expected = int(components[4]) if len(components) > 4 else 0 self.transfer_statuses[room].expected_kvs_per_pp[pp_rank] = expected if self.transfer_statuses[room].num_pp_ranks_expected is None: - self.transfer_statuses[room].num_pp_ranks_expected = ( - self.required_prefill_response_num_table.get(room, 1) + self.transfer_statuses[ + room + ].num_pp_ranks_expected = self.required_prefill_response_num_table.get( + room, 1 ) if ( self.enable_staging @@ -2574,8 +2575,10 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager): if is_last_chunk: self.transfer_statuses[room].expected_kvs_per_pp[pp_rank] = chunk_id + 1 if self.transfer_statuses[room].num_pp_ranks_expected is None: - self.transfer_statuses[room].num_pp_ranks_expected = ( - self.required_prefill_response_num_table.get(room, 1) + self.transfer_statuses[ + room + ].num_pp_ranks_expected = self.required_prefill_response_num_table.get( + room, 1 ) if ( self.enable_staging @@ -2739,9 +2742,9 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager): if self._handle_abort_notification(waiting_req_bytes): continue - assert ( - waiting_req_bytes[0] == GUARD - ), f"First message should be {GUARD}. Foreign traffic?" + assert waiting_req_bytes[0] == GUARD, ( + f"First message should be {GUARD}. Foreign traffic?" + ) waiting_req_bytes = waiting_req_bytes[1:] room = waiting_req_bytes[0].decode("ascii") agent_name = waiting_req_bytes[3].decode("ascii") diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 404a9d07f..4c8072b4e 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -852,9 +852,9 @@ class SchedulerDisaggregationPrefillMixin: # In non-overlap-mode, KV is sent in process_prefill_chunk # Only send when req's sender is initialized if self.enable_overlap and not req.pending_bootstrap: - assert ( - req.metadata_buffer_index >= 0 - ), f"Req {req.rid} does not have metadata buffer allocated" + assert req.metadata_buffer_index >= 0, ( + f"Req {req.rid} does not have metadata buffer allocated" + ) self.send_kv_chunk(req, last_chunk=False, end_idx=req.tmp_end_idx) req.time_stats.set_last_chunked_prefill_finish_time() diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py index d42fe63a3..6656143f3 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py @@ -63,7 +63,7 @@ def update_environment_variables(envs: Dict[str, str]): for k, v in envs.items(): if k in os.environ and os.environ[k] != v: logger.warning( - "Overwriting environment variable %s " "from '%s' to '%s'", + "Overwriting environment variable %s from '%s' to '%s'", k, os.environ[k], v, @@ -445,9 +445,9 @@ def can_use_custom_all_reduce_with_nvlink( supported_world_size: List[int], cls_name: str, ) -> Optional[bool]: # None if fail; otherwise return whether NVLink is available - assert ( - dist.get_backend(group) != dist.Backend.NCCL - ), f"{cls_name} should be attached to a non-NCCL group." + assert dist.get_backend(group) != dist.Backend.NCCL, ( + f"{cls_name} should be attached to a non-NCCL group." + ) rank = dist.get_rank(group=group) world_size = dist.get_world_size(group=group) @@ -459,7 +459,7 @@ def can_use_custom_all_reduce_with_nvlink( # No need to initialize custom allreduce for multi-node case. if not all(in_the_same_node_as(group, source_rank=0)): logger.warning( - f"{cls_name} is disabled because this process group" " spans across nodes." + f"{cls_name} is disabled because this process group spans across nodes." ) return diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py index 5b3ea5a6d..15b0bbada 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py @@ -404,9 +404,9 @@ class CustomAllReduceV2: yield finally: self._graph_mode_allowed = False - assert ( - not torch.cuda.is_current_stream_capturing() - ), "Cannot register graph inputs while capturing CUDA graph" + assert not torch.cuda.is_current_stream_capturing(), ( + "Cannot register graph inputs while capturing CUDA graph" + ) self._register_graph_inputs() def _register_graph_inputs(self) -> None: diff --git a/python/sglang/srt/distributed/device_communicators/hpu_communicator.py b/python/sglang/srt/distributed/device_communicators/hpu_communicator.py index e733e81e3..cfc0eff9e 100644 --- a/python/sglang/srt/distributed/device_communicators/hpu_communicator.py +++ b/python/sglang/srt/distributed/device_communicators/hpu_communicator.py @@ -13,7 +13,6 @@ if is_hpu(): class HpuCommunicator: - def __init__(self, group: ProcessGroup): if not is_hpu(): self.disabled = True diff --git a/python/sglang/srt/distributed/device_communicators/npu_communicator.py b/python/sglang/srt/distributed/device_communicators/npu_communicator.py index 5518584d4..ec3f28fc6 100644 --- a/python/sglang/srt/distributed/device_communicators/npu_communicator.py +++ b/python/sglang/srt/distributed/device_communicators/npu_communicator.py @@ -11,7 +11,6 @@ if _is_npu: class NpuCommunicator: - def __init__(self, group: ProcessGroup): if not _is_npu: self.disabled = True diff --git a/python/sglang/srt/distributed/device_communicators/pymscclpp.py b/python/sglang/srt/distributed/device_communicators/pymscclpp.py index 95d9411e4..1d5c2cfaa 100644 --- a/python/sglang/srt/distributed/device_communicators/pymscclpp.py +++ b/python/sglang/srt/distributed/device_communicators/pymscclpp.py @@ -267,9 +267,9 @@ class PyMscclppCommunicator: self.available = True self.group = group - assert ( - dist.get_backend(group) != dist.Backend.NCCL - ), "CustomAllreduce should be attached to a non-NCCL group." + assert dist.get_backend(group) != dist.Backend.NCCL, ( + "CustomAllreduce should be attached to a non-NCCL group." + ) rank = dist.get_rank(group=self.group) world_size = dist.get_world_size(group=self.group) diff --git a/python/sglang/srt/distributed/device_communicators/pynccl.py b/python/sglang/srt/distributed/device_communicators/pynccl.py index 9072ec096..46aef5b94 100644 --- a/python/sglang/srt/distributed/device_communicators/pynccl.py +++ b/python/sglang/srt/distributed/device_communicators/pynccl.py @@ -28,7 +28,6 @@ logger = logging.getLogger(__name__) class PyNcclCommunicator: - def __init__( self, group: Union[ProcessGroup, StatelessProcessGroup], @@ -50,9 +49,9 @@ class PyNcclCommunicator: """ if not isinstance(group, StatelessProcessGroup): assert dist.is_initialized() - assert ( - dist.get_backend(group) != dist.Backend.NCCL - ), "PyNcclCommunicator should be attached to a non-NCCL group." + assert dist.get_backend(group) != dist.Backend.NCCL, ( + "PyNcclCommunicator should be attached to a non-NCCL group." + ) # note: this rank is the rank in the group self.rank = dist.get_rank(group) self.world_size = dist.get_world_size(group) diff --git a/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py b/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py index 3574968e1..283a07130 100644 --- a/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py +++ b/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py @@ -266,14 +266,14 @@ class SymmetricMemoryContext: self._comm_ptr = self.group_coordinator.pynccl_comm.comm.value def __enter__(self): - assert ( - self.group_coordinator.pynccl_comm is not None - ), f"Symmetric memory requires pynccl to be enabled in group '{self.group_coordinator.unique_name}'" + assert self.group_coordinator.pynccl_comm is not None, ( + f"Symmetric memory requires pynccl to be enabled in group '{self.group_coordinator.unique_name}'" + ) if self.is_graph_capture: - assert ( - _graph_pool_id is not None - ), "graph_pool_id is not set under graph capture" + assert _graph_pool_id is not None, ( + "graph_pool_id is not set under graph capture" + ) # Pause graph memory pool to use symmetric memory with cuda graph if after_2_8_0: torch._C._cuda_endAllocateToPool(_cur_device, _graph_pool_id) @@ -322,9 +322,9 @@ class SymmetricMemoryContext: # Call C++ API to register all segments with this comm # C++ layer tracks per-comm registration state internally result = _register_func(self._comm_ptr) - assert ( - result == 0 - ), f"nccl_allocator_register_segments_with_comm failed with return code: {result}" + assert result == 0, ( + f"nccl_allocator_register_segments_with_comm failed with return code: {result}" + ) def use_symmetric_memory(group_coordinator: GroupCoordinator, disabled: bool = False): diff --git a/python/sglang/srt/distributed/device_communicators/quick_all_reduce.py b/python/sglang/srt/distributed/device_communicators/quick_all_reduce.py index f9d51246e..4a6835ad0 100644 --- a/python/sglang/srt/distributed/device_communicators/quick_all_reduce.py +++ b/python/sglang/srt/distributed/device_communicators/quick_all_reduce.py @@ -50,7 +50,6 @@ MB = 1024 * 1024 class QuickAllReduce: - _SUPPORTED_WORLD_SIZES = [2, 4, 8] _SUPPORTED_DTYPES = [torch.float16, torch.bfloat16] # The following data is based on kernel tests. @@ -103,9 +102,9 @@ class QuickAllReduce: return self.group = group - assert ( - dist.get_backend(group) != dist.Backend.NCCL - ), "Custom quick allreduce should be attached to a non-NCCL group." + assert dist.get_backend(group) != dist.Backend.NCCL, ( + "Custom quick allreduce should be attached to a non-NCCL group." + ) if not all(in_the_same_node_as(group, source_rank=0)): # No need to initialize custom quick allreduce for # multi-node case. diff --git a/python/sglang/srt/distributed/device_communicators/shm_broadcast.py b/python/sglang/srt/distributed/device_communicators/shm_broadcast.py index 721a07431..38f249790 100644 --- a/python/sglang/srt/distributed/device_communicators/shm_broadcast.py +++ b/python/sglang/srt/distributed/device_communicators/shm_broadcast.py @@ -31,7 +31,6 @@ logger = logging.getLogger(__name__) class ShmRingBuffer: - def __init__( self, n_reader: int, @@ -173,7 +172,6 @@ class Handle: class MessageQueue: - def __init__( self, n_reader, # number of all readers diff --git a/python/sglang/srt/distributed/device_communicators/triton_symm_mem_ag.py b/python/sglang/srt/distributed/device_communicators/triton_symm_mem_ag.py index fa508304c..a7cf40bd7 100644 --- a/python/sglang/srt/distributed/device_communicators/triton_symm_mem_ag.py +++ b/python/sglang/srt/distributed/device_communicators/triton_symm_mem_ag.py @@ -374,9 +374,9 @@ def all_gather_inner( f"hidden_states.data_ptr()={hex(hidden_states.data_ptr())} must be " f"16-byte aligned for 128-bit multimem.st" ) - assert ( - tp_hidden_dim % world_size == 0 - ), f"tp_hidden_dim={tp_hidden_dim} must be divisible by world_size={world_size}" + assert tp_hidden_dim % world_size == 0, ( + f"tp_hidden_dim={tp_hidden_dim} must be divisible by world_size={world_size}" + ) local_hidden = tp_hidden_dim // world_size assert local_hidden % _NUMEL_PER_THREAD == 0, ( f"per-rank hidden shard ({local_hidden}) must be a multiple of " @@ -387,12 +387,12 @@ def all_gather_inner( f"state.hidden_dim={state.hidden_dim}" ) total_tokens, in_hidden = hidden_states.shape - assert ( - in_hidden == local_hidden - ), f"input hidden ({in_hidden}) != this rank's shard ({local_hidden})" - assert ( - total_tokens <= state.max_token_num - ), f"total_tokens={total_tokens} exceeds max_token_num={state.max_token_num}" + assert in_hidden == local_hidden, ( + f"input hidden ({in_hidden}) != this rank's shard ({local_hidden})" + ) + assert total_tokens <= state.max_token_num, ( + f"total_tokens={total_tokens} exceeds max_token_num={state.max_token_num}" + ) hidden_offset = local_hidden * state.rank_in_group symm_mem_hdl = state.symm_mem_hdl diff --git a/python/sglang/srt/distributed/device_communicators/xpu_communicator.py b/python/sglang/srt/distributed/device_communicators/xpu_communicator.py index 78931ac7c..2bca063bf 100644 --- a/python/sglang/srt/distributed/device_communicators/xpu_communicator.py +++ b/python/sglang/srt/distributed/device_communicators/xpu_communicator.py @@ -10,7 +10,6 @@ from sglang.srt.utils import is_xpu class XpuCommunicator: - def __init__(self, group: ProcessGroup): if not is_xpu(): self.disabled = True diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 4d7c64e40..efb3ff07a 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -1036,9 +1036,9 @@ class GroupCoordinator: # Bypass the function if we are using only 1 GPU. if world_size == 1: return input_ - assert ( - -input_.dim() <= dim < input_.dim() - ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + assert -input_.dim() <= dim < input_.dim(), ( + f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + ) if dim < 0: # Convert negative dim to positive. @@ -1178,9 +1178,9 @@ class GroupCoordinator: pynccl_comm = self.pynccl_comm with pynccl_comm.change_state(enable=True): - assert ( - pynccl_comm is not None and not pynccl_comm.disabled - ), "pynccl is required for reduce_scatterv" + assert pynccl_comm is not None and not pynccl_comm.disabled, ( + "pynccl is required for reduce_scatterv" + ) if sizes is not None: assert len(sizes) == world_size @@ -1303,9 +1303,9 @@ class GroupCoordinator: output_tensor_list, input_, group=self.device_group ) - assert ( - -input_.dim() <= dim < input_.dim() - ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + assert -input_.dim() <= dim < input_.dim(), ( + f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + ) # For HPUs, use HPU communicator. hpu_comm = self.hpu_communicator @@ -1369,9 +1369,9 @@ class GroupCoordinator: pynccl_comm = self.pynccl_comm with pynccl_comm.change_state(enable=True): - assert ( - pynccl_comm is not None and not pynccl_comm.disabled - ), "pynccl is required for all_gatherv" + assert pynccl_comm is not None and not pynccl_comm.disabled, ( + "pynccl is required for all_gatherv" + ) def _all_gather_allocate_output( input_: torch.Tensor, @@ -1435,9 +1435,9 @@ class GroupCoordinator: # Bypass the function if we are using only 1 GPU. if world_size == 1: return input_ - assert ( - -input_.dim() <= dim < input_.dim() - ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + assert -input_.dim() <= dim < input_.dim(), ( + f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + ) if dim < 0: # Convert negative dim to positive. dim += input_.dim() @@ -1580,9 +1580,9 @@ class GroupCoordinator: """NOTE: `src` is the local rank of the source rank.""" assert src < self.world_size, f"Invalid src rank ({src})" - assert ( - src != self.rank_in_group - ), "Invalid source rank. Source rank is the same as the current rank." + assert src != self.rank_in_group, ( + "Invalid source rank. Source rank is the same as the current rank." + ) size_tensor = torch.empty(1, dtype=torch.long, device="cpu") @@ -1629,9 +1629,9 @@ class GroupCoordinator: rank_in_group = self.rank_in_group if rank_in_group == src: metadata_list: List[Tuple[Any, Any]] = [] - assert isinstance( - tensor_dict, dict - ), f"Expecting a dictionary, got {type(tensor_dict)}" + assert isinstance(tensor_dict, dict), ( + f"Expecting a dictionary, got {type(tensor_dict)}" + ) metadata_list, tensor_list = _split_tensor_dict(tensor_dict) # `metadata_list` lives in CPU memory. # `broadcast_object_list` has serialization & deserialization, @@ -1716,9 +1716,9 @@ class GroupCoordinator: dst = (self.rank_in_group + 1) % self.world_size assert dst < self.world_size, f"Invalid dst rank ({dst})" - assert isinstance( - tensor_dict, dict - ), f"Expecting a dictionary, got {type(tensor_dict)}" + assert isinstance(tensor_dict, dict), ( + f"Expecting a dictionary, got {type(tensor_dict)}" + ) metadata_list, tensor_list = _split_tensor_dict(tensor_dict) # Note: While switching to Device-to-Device (D2D) would introduce an extra # Device-to-Host (D2H) memory copy overhead for serialization, our benchmarks @@ -1948,25 +1948,25 @@ def set_pdmux_status(enable_prefill_multiplexing: bool): def get_tp_group() -> GroupCoordinator: if _ENABLE_PDMUX_P_TP: - assert ( - _PDMUX_PREFILL_TP_GROUP is not None - ), "tensor model parallel group for PD-Multiplexing Prefill is not initialized" + assert _PDMUX_PREFILL_TP_GROUP is not None, ( + "tensor model parallel group for PD-Multiplexing Prefill is not initialized" + ) return _PDMUX_PREFILL_TP_GROUP assert _TP is not None, "tensor model parallel group is not initialized" return _TP def get_attn_tp_group() -> GroupCoordinator: - assert ( - _ATTN_TP is not None - ), "attention tensor model parallel group is not initialized" + assert _ATTN_TP is not None, ( + "attention tensor model parallel group is not initialized" + ) return _ATTN_TP def get_attn_cp_group() -> GroupCoordinator: - assert ( - _ATTN_CP is not None - ), "attention context model parallel group is not initialized" + assert _ATTN_CP is not None, ( + "attention context model parallel group is not initialized" + ) return _ATTN_CP @@ -1987,9 +1987,9 @@ def _init_attn_cp_overlap_group( """Second communicator over the attention CP ranks; RCCL deadlocks when one communicator is driven from two streams at once.""" global _ATTN_CP_OVERLAP - assert ( - _ATTN_CP_OVERLAP is None - ), "attention context parallel overlap group is already initialized" + assert _ATTN_CP_OVERLAP is None, ( + "attention context parallel overlap group is already initialized" + ) if attn_cp_size <= 1: return @@ -2259,7 +2259,7 @@ def init_distributed_environment( max_world_size: Optional[int] = None, ): logger.debug( - "world_size=%d rank=%d local_rank=%d " "distributed_init_method=%s backend=%s", + "world_size=%d rank=%d local_rank=%d distributed_init_method=%s backend=%s", world_size, rank, local_rank, @@ -2334,9 +2334,9 @@ def init_distributed_environment( ranks, local_rank, backend, recovered_rank=recovered_rank ) else: - assert ( - _WORLD.world_size == torch.distributed.get_world_size() - ), "world group already initialized with a different world size" + assert _WORLD.world_size == torch.distributed.get_world_size(), ( + "world group already initialized with a different world size" + ) def initialize_model_parallel( @@ -2470,9 +2470,9 @@ def initialize_model_parallel( if duplicate_tp_group: global _PDMUX_PREFILL_TP_GROUP - assert ( - _PDMUX_PREFILL_TP_GROUP is None - ), "tensor model parallel group for PD-Multiplexing Prefill is already initialized" + assert _PDMUX_PREFILL_TP_GROUP is None, ( + "tensor model parallel group for PD-Multiplexing Prefill is already initialized" + ) _PDMUX_PREFILL_TP_GROUP = init_model_parallel_group( group_ranks, get_world_group().local_rank, @@ -2526,9 +2526,9 @@ def initialize_model_parallel( attn_tp_size = derived_widths["attn_tp_size"] global _ATTN_CP - assert ( - _ATTN_CP is None - ), "attention context model parallel group is already initialized" + assert _ATTN_CP is None, ( + "attention context model parallel group is already initialized" + ) if attn_cp_size == tensor_model_parallel_size: _ATTN_CP = _TP else: @@ -2573,9 +2573,9 @@ def initialize_model_parallel( from sglang.srt.layers.sampler import SYNC_TOKEN_IDS_ACROSS_TP global _ATTN_TP - assert ( - _ATTN_TP is None - ), "attention tensor model parallel group is already initialized" + assert _ATTN_TP is None, ( + "attention tensor model parallel group is already initialized" + ) if attn_tp_size == tensor_model_parallel_size: _ATTN_TP = _TP else: @@ -2821,9 +2821,9 @@ def ensure_model_parallel_initialized( ) if decode_context_parallel_size > 1: dcp_world_size = get_dcp_group().world_size - assert ( - dcp_world_size == decode_context_parallel_size - ), f"decode context parallel group already initialized, but of unexpected size: {dcp_world_size=} {decode_context_parallel_size=}" + assert dcp_world_size == decode_context_parallel_size, ( + f"decode context parallel group already initialized, but of unexpected size: {dcp_world_size=} {decode_context_parallel_size=}" + ) def model_parallel_is_initialized(): @@ -3084,9 +3084,9 @@ def in_the_same_node_as(pg: ProcessGroup, source_rank: int = 0) -> List[bool]: as the source rank. It tests if processes are attached to the same memory system (shared access to shared memory). """ - assert ( - torch.distributed.get_backend(pg) != torch.distributed.Backend.NCCL - ), "in_the_same_node_as should be tested with a non-NCCL group." + assert torch.distributed.get_backend(pg) != torch.distributed.Backend.NCCL, ( + "in_the_same_node_as should be tested with a non-NCCL group." + ) # local rank inside the group rank = torch.distributed.get_rank(group=pg) world_size = torch.distributed.get_world_size(group=pg) diff --git a/python/sglang/srt/distributed/utils.py b/python/sglang/srt/distributed/utils.py index 0aeead003..8862658fa 100644 --- a/python/sglang/srt/distributed/utils.py +++ b/python/sglang/srt/distributed/utils.py @@ -197,13 +197,13 @@ class StatelessProcessGroup: """ if self.rank == src: self.expire_data() - key = f"broadcast_from/{src}/" f"{self.broadcast_send_counter}" + key = f"broadcast_from/{src}/{self.broadcast_send_counter}" self.store.set(key, pickle.dumps(obj)) self.broadcast_send_counter += 1 self.entries.append((key, time.perf_counter())) return obj else: - key = f"broadcast_from/{src}/" f"{self.broadcast_recv_src_counter[src]}" + key = f"broadcast_from/{src}/{self.broadcast_recv_src_counter[src]}" recv_obj = pickle.loads(self.store.get(key)) self.broadcast_recv_src_counter[src] += 1 return recv_obj diff --git a/python/sglang/srt/dllm/mixin/req.py b/python/sglang/srt/dllm/mixin/req.py index ff72cd4fc..68ea82011 100644 --- a/python/sglang/srt/dllm/mixin/req.py +++ b/python/sglang/srt/dllm/mixin/req.py @@ -86,8 +86,8 @@ class ReqDllmMixin: def _update_block_offset_for_dllm(self): prefix_len = len(self.prefix_indices) - assert ( - prefix_len % self.dllm_config.block_size == 0 - ), f"Unexpected prefix len: {prefix_len}" + assert prefix_len % self.dllm_config.block_size == 0, ( + f"Unexpected prefix len: {prefix_len}" + ) if prefix_len > self.dllm_block_offset: self.dllm_block_offset = prefix_len diff --git a/python/sglang/srt/dllm/mixin/scheduler.py b/python/sglang/srt/dllm/mixin/scheduler.py index 86b1a9500..b127960e2 100644 --- a/python/sglang/srt/dllm/mixin/scheduler.py +++ b/python/sglang/srt/dllm/mixin/scheduler.py @@ -75,9 +75,9 @@ class SchedulerDllmMixin: result.copy_done.synchronize() fdfo_mode = self.dllm_config.first_done_first_out_mode - assert ( - not fdfo_mode or result.accept_length_per_req_cpu is not None - ), "FDFO dLLM result is missing accept lengths." + assert not fdfo_mode or result.accept_length_per_req_cpu is not None, ( + "FDFO dLLM result is missing accept lengths." + ) # FDFO also commits unresolved blocks so their KV can be reused. if fdfo_mode or result.next_token_ids: @@ -317,9 +317,8 @@ class SchedulerDllmMixin: # Try preemption if batch is full if running_batch.batch_is_full: - if ( - not self.enable_priority_preemption - or not adder.preempt_to_schedule(req) + if not self.enable_priority_preemption or not adder.preempt_to_schedule( + req ): break diff --git a/python/sglang/srt/entrypoints/anthropic/protocol.py b/python/sglang/srt/entrypoints/anthropic/protocol.py index 95217d6b3..a79952e3d 100644 --- a/python/sglang/srt/entrypoints/anthropic/protocol.py +++ b/python/sglang/srt/entrypoints/anthropic/protocol.py @@ -284,13 +284,13 @@ class AnthropicThinkingParam(BaseModel): if self.type == "enabled": if self.budget_tokens is None: raise ValueError( - "thinking.budget_tokens is required when " - "thinking.type is 'enabled'" + "thinking.budget_tokens is required when thinking.type is 'enabled'" ) if self.budget_tokens < 1024: raise ValueError( - "thinking.budget_tokens must be >= 1024 " - "(got {})".format(self.budget_tokens) + "thinking.budget_tokens must be >= 1024 (got {})".format( + self.budget_tokens + ) ) elif self.type == "disabled": if self.budget_tokens is not None: @@ -300,8 +300,7 @@ class AnthropicThinkingParam(BaseModel): ) if self.display is not None: raise ValueError( - "thinking.display is not allowed when " - "thinking.type is 'disabled'" + "thinking.display is not allowed when thinking.type is 'disabled'" ) elif self.type == "adaptive": if self.budget_tokens is not None: diff --git a/python/sglang/srt/entrypoints/anthropic/serving.py b/python/sglang/srt/entrypoints/anthropic/serving.py index 37bcdb43a..b20bbdbab 100644 --- a/python/sglang/srt/entrypoints/anthropic/serving.py +++ b/python/sglang/srt/entrypoints/anthropic/serving.py @@ -1069,8 +1069,7 @@ class AnthropicServing: effective_finish = finish_reason or "stop" if effective_finish not in STOP_REASON_MAP: logger.warning( - "Unmapped streaming finish_reason %r; defaulting " - "to end_turn", + "Unmapped streaming finish_reason %r; defaulting to end_turn", effective_finish, ) stop_reason = STOP_REASON_MAP.get(effective_finish, "end_turn") diff --git a/python/sglang/srt/entrypoints/context.py b/python/sglang/srt/entrypoints/context.py index 652bc92a4..c310c3e20 100644 --- a/python/sglang/srt/entrypoints/context.py +++ b/python/sglang/srt/entrypoints/context.py @@ -24,7 +24,6 @@ from sglang.srt.entrypoints.tool import Tool class ConversationContext(ABC): - @abstractmethod def append_output(self, output) -> None: pass @@ -43,7 +42,6 @@ class ConversationContext(ABC): class SimpleContext(ConversationContext): - def __init__(self): self.last_output = None @@ -61,7 +59,6 @@ class SimpleContext(ConversationContext): class HarmonyContext(ConversationContext): - def __init__( self, messages: list, @@ -182,7 +179,6 @@ class HarmonyContext(ConversationContext): class StreamingHarmonyContext(HarmonyContext): - def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.last_output = None diff --git a/python/sglang/srt/entrypoints/grpc_bridge.py b/python/sglang/srt/entrypoints/grpc_bridge.py index 853c1c62e..50870c93c 100644 --- a/python/sglang/srt/entrypoints/grpc_bridge.py +++ b/python/sglang/srt/entrypoints/grpc_bridge.py @@ -545,9 +545,11 @@ class RuntimeHandle: obj = UpdateWeightFromDiskReqInput( model_path=model_path, load_format=load_format ) - success, message, num_paused = ( - await self.tokenizer_manager.update_weights_from_disk(obj, request=None) - ) + ( + success, + message, + num_paused, + ) = await self.tokenizer_manager.update_weights_from_disk(obj, request=None) return { "success": success, "message": message, diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 872d1b169..18c865af4 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -231,9 +231,9 @@ async def init_multi_tokenizer() -> ServerArgs: publish(server_args, role="tokenizer") # API key authentication is not supported in multi-tokenizer mode - assert ( - get_serving().api_key is None - ), "API key is not supported in multi-tokenizer mode" + assert get_serving().api_key is None, ( + "API key is not supported in multi-tokenizer mode" + ) # Create a new ipc name for the current process port_args.tokenizer_ipc_name = ( @@ -819,9 +819,9 @@ async def server_info(): HiCache mirror by `GET /hicache/storage-backend`. """ # Returns internal states per DP. - internal_states: List[Dict[Any, Any]] = ( - await _global_state.tokenizer_manager.get_internal_state() - ) + internal_states: List[ + Dict[Any, Any] + ] = await _global_state.tokenizer_manager.get_internal_state() server_args = _global_state.tokenizer_manager.server_args @@ -1523,9 +1523,12 @@ async def check_weights( ): if obj is None: obj = CheckWeightsReqInput() - success, message, ranks, per_engine_checksum = ( - await _global_state.tokenizer_manager.check_weights(obj, request) - ) + ( + success, + message, + ranks, + per_engine_checksum, + ) = await _global_state.tokenizer_manager.check_weights(obj, request) body = {"success": success, "message": message} if ranks is not None: body["ranks"] = ranks @@ -2395,8 +2398,7 @@ def _wait_and_warmup( skip_elastic_joiner_warmup = server_args.is_ep_scale_joiner if skip_elastic_joiner_warmup: logger.debug( - "[Elastic EP] Skipping server warmup for elastic joiner " - "(ep_join_mode=%s)", + "[Elastic EP] Skipping server warmup for elastic joiner (ep_join_mode=%s)", get_exec().moe.ep_join_mode, ) diff --git a/python/sglang/srt/entrypoints/openai/encoding_dsv32.py b/python/sglang/srt/entrypoints/openai/encoding_dsv32.py index f20f07451..5ae38d6df 100644 --- a/python/sglang/srt/entrypoints/openai/encoding_dsv32.py +++ b/python/sglang/srt/entrypoints/openai/encoding_dsv32.py @@ -46,9 +46,7 @@ user_msg_template: str = "<|User|>{content}<|Assistant|>" assistant_msg_template: str = "{reasoning}{content}{tool_calls}<|end▁of▁sentence|>" thinking_template = "{reasoning_content}" -response_format_template: str = ( - "## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{schema}" -) +response_format_template: str = "## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{schema}" tool_call_template: str = ( '<{dsml_token}invoke name="{name}">\n{arguments}\n' ) diff --git a/python/sglang/srt/entrypoints/openai/encoding_dsv4.py b/python/sglang/srt/entrypoints/openai/encoding_dsv4.py index d6477ba4e..52c8bebd7 100644 --- a/python/sglang/srt/entrypoints/openai/encoding_dsv4.py +++ b/python/sglang/srt/entrypoints/openai/encoding_dsv4.py @@ -47,9 +47,7 @@ assistant_msg_template: str = "{reasoning}{content}{tool_calls}" + eos_token assistant_msg_wo_eos_template: str = "{reasoning}{content}{tool_calls}" thinking_template: str = "{reasoning_content}" -response_format_template: str = ( - "## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{schema}" -) +response_format_template: str = "## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{schema}" tool_call_template: str = ( '<{dsml_token}invoke name="{name}">\n{arguments}\n' ) @@ -447,9 +445,9 @@ def render_message( task = messages[index].get("task") if task is not None: # Task special token for internal classification tasks - assert ( - task in VALID_TASKS - ), f"Invalid task: '{task}'. Valid tasks are: {list(VALID_TASKS)}" + assert task in VALID_TASKS, ( + f"Invalid task: '{task}'. Valid tasks are: {list(VALID_TASKS)}" + ) task_sp_token = DS_TASK_SP_TOKENS[task] if task != "action": @@ -843,9 +841,9 @@ def parse_message_from_completion_text(text: str, thinking_mode: str) -> Dict[st index, text, [thinking_end_token, tool_calls_start_token] ) reasoning_content = content_delta - assert ( - stop_token == thinking_end_token - ), "Invalid thinking format: missing " + assert stop_token == thinking_end_token, ( + "Invalid thinking format: missing " + ) index, content_delta, stop_token = _read_until_stop( index, text, [eos_token, tool_calls_start_token] @@ -874,9 +872,9 @@ def parse_message_from_completion_text(text: str, thinking_mode: str) -> Dict[st thinking_end_token, dsml_token, ]: - assert ( - sp_token not in summary_content and sp_token not in reasoning_content - ), f"Unexpected special token '{sp_token}' in content" + assert sp_token not in summary_content and sp_token not in reasoning_content, ( + f"Unexpected special token '{sp_token}' in content" + ) return { "role": "assistant", diff --git a/python/sglang/srt/entrypoints/openai/serving_completions.py b/python/sglang/srt/entrypoints/openai/serving_completions.py index dbc2d4b6a..1f8ff7d0d 100644 --- a/python/sglang/srt/entrypoints/openai/serving_completions.py +++ b/python/sglang/srt/entrypoints/openai/serving_completions.py @@ -308,9 +308,7 @@ class OpenAIServingCompletion(OpenAIServingBase): output_top_logprobs = content["meta_info"].get( "output_top_logprobs", [] ) - if ( - not self.tokenizer_manager.server_args.incremental_streaming_output - ): + if not self.tokenizer_manager.server_args.incremental_streaming_output: output_token_logprobs = output_token_logprobs[ n_prev_token:total_output_logprobs ] @@ -329,9 +327,7 @@ class OpenAIServingCompletion(OpenAIServingBase): chunk_prompt_token_ids = None if request.return_token_ids: output_ids = content["output_ids"] - if ( - not self.tokenizer_manager.server_args.incremental_streaming_output - ): + if not self.tokenizer_manager.server_args.incremental_streaming_output: n_prev_token_id = n_prev_token_ids.get(index, 0) chunk_token_ids = output_ids[n_prev_token_id:] n_prev_token_ids[index] = len(output_ids) diff --git a/python/sglang/srt/entrypoints/openai/serving_responses.py b/python/sglang/srt/entrypoints/openai/serving_responses.py index 392f6533a..bec853e81 100644 --- a/python/sglang/srt/entrypoints/openai/serving_responses.py +++ b/python/sglang/srt/entrypoints/openai/serving_responses.py @@ -333,8 +333,7 @@ class OpenAIServingResponses(OpenAIServingChat): ) ): return self.create_error_response( - "MCP tool server is not supported in background mode and " - "streaming mode" + "MCP tool server is not supported in background mode and streaming mode" ) # Schedule the request and get the result generator @@ -540,17 +539,17 @@ class OpenAIServingResponses(OpenAIServingChat): require_reasoning=require_reasoning, ) try: - result: Union[ORJSONResponse, ResponsesResponse] = ( - await self.responses_full_generator( - request, - sampling_params, - result_generator, - context, - model_name, - tokenizer, - request_metadata, - require_reasoning=require_reasoning, - ) + result: Union[ + ORJSONResponse, ResponsesResponse + ] = await self.responses_full_generator( + request, + sampling_params, + result_generator, + context, + model_name, + tokenizer, + request_metadata, + require_reasoning=require_reasoning, ) return result except Exception as e: @@ -609,7 +608,7 @@ class OpenAIServingResponses(OpenAIServingChat): ): if request.tool_choice != "auto": raise NotImplementedError( - "Only 'auto' tool_choice is supported in " "response API" + "Only 'auto' tool_choice is supported in response API" ) messages = self._construct_input_messages_with_harmony(request, prev_response) prompt_token_ids = render_for_completion(messages) @@ -1333,9 +1332,7 @@ class OpenAIServingResponses(OpenAIServingChat): recent_turn_msgs = prev_msgs[prev_final_msg_idx + 1 :] del prev_msgs[prev_final_msg_idx + 1 :] for msg in recent_turn_msgs: - if ( - hasattr(msg, "channel") and msg.channel != "analysis" - ): # type: ignore[union-attr] + if hasattr(msg, "channel") and msg.channel != "analysis": # type: ignore[union-attr] prev_msgs.append(msg) messages.extend(prev_msgs) # Append the new input. @@ -1489,8 +1486,7 @@ class OpenAIServingResponses(OpenAIServingChat): # Get event type from the event's type field if it exists event_type = getattr(event, "type", "unknown") return ( - f"event: {event_type}\n" - f"data: {event.model_dump_json(indent=None)}\n\n" + f"event: {event_type}\ndata: {event.model_dump_json(indent=None)}\n\n" ) current_content_index = 0 @@ -1919,8 +1915,7 @@ class OpenAIServingResponses(OpenAIServingChat): sequence_number += 1 event_type = getattr(event, "type", "unknown") return ( - f"event: {event_type}\n" - f"data: {event.model_dump_json(indent=None)}\n\n" + f"event: {event_type}\ndata: {event.model_dump_json(indent=None)}\n\n" ) # The streaming Response* event models echo ``tools`` through a diff --git a/python/sglang/srt/entrypoints/openai/tool_server.py b/python/sglang/srt/entrypoints/openai/tool_server.py index e4bebf647..f51323012 100644 --- a/python/sglang/srt/entrypoints/openai/tool_server.py +++ b/python/sglang/srt/entrypoints/openai/tool_server.py @@ -72,7 +72,6 @@ def post_process_tools_description( class ToolServer(ABC): - @abstractmethod def has_tool(self, tool_name: str): pass @@ -86,7 +85,6 @@ class ToolServer(ABC): class MCPToolServer(ToolServer): - def __init__(self): self.harmony_tool_descriptions = {} @@ -143,7 +141,6 @@ class MCPToolServer(ToolServer): class DemoToolServer(ToolServer): - def __init__(self, *, enable_python: bool = True): from sglang.srt.entrypoints.tool import ( HarmonyBrowserTool, diff --git a/python/sglang/srt/entrypoints/ssl_utils.py b/python/sglang/srt/entrypoints/ssl_utils.py index 87e2b16f2..6885bb4ff 100644 --- a/python/sglang/srt/entrypoints/ssl_utils.py +++ b/python/sglang/srt/entrypoints/ssl_utils.py @@ -46,7 +46,7 @@ class SSLCertRefresher: try: async for _changes in awatch(self._cert_path, self._key_path): logger.info( - "SSL cert/key file change detected, reloading: " "cert=%s key=%s", + "SSL cert/key file change detected, reloading: cert=%s key=%s", self._cert_path, self._key_path, ) diff --git a/python/sglang/srt/entrypoints/tool.py b/python/sglang/srt/entrypoints/tool.py index 00b54cfd2..bac679c29 100644 --- a/python/sglang/srt/entrypoints/tool.py +++ b/python/sglang/srt/entrypoints/tool.py @@ -20,14 +20,12 @@ logger = logging.getLogger(__name__) class Tool(ABC): - @abstractmethod async def get_result(self, context: "ConversationContext") -> Any: pass class HarmonyBrowserTool(Tool): - def __init__(self, client: ExaClient | None = None): self.enabled = True if client is not None: @@ -257,7 +255,6 @@ class HarmonyBrowserTool(Tool): class HarmonyPythonTool(Tool): - def __init__(self): self.enabled = True diff --git a/python/sglang/srt/eplb/eplb_manager.py b/python/sglang/srt/eplb/eplb_manager.py index 25be40dae..e1e6fbd2b 100644 --- a/python/sglang/srt/eplb/eplb_manager.py +++ b/python/sglang/srt/eplb/eplb_manager.py @@ -58,7 +58,9 @@ class EPLBManager: assert ( get_exec().moe.eplb_rebalance_num_iterations >= get_exec().moe.expert_distribution_recorder_buffer_size - ), "eplb_rebalance_num_iterations must be greater than expert_distribution_recorder_buffer_size" + ), ( + "eplb_rebalance_num_iterations must be greater than expert_distribution_recorder_buffer_size" + ) if not get_global_expert_distribution_recorder().recording: get_global_expert_distribution_recorder().start_record() diff --git a/python/sglang/srt/eplb/eplb_simulator/reader.py b/python/sglang/srt/eplb/eplb_simulator/reader.py index 97405c319..2a0b127fe 100644 --- a/python/sglang/srt/eplb/eplb_simulator/reader.py +++ b/python/sglang/srt/eplb/eplb_simulator/reader.py @@ -24,9 +24,9 @@ def read_mode_per_pass(dir_data: Path): for record in data_pack["records"]: forward_pass_id = record["forward_pass_id"] rank = record["rank"] - assert ( - gpc_of_forward_pass_and_rank[forward_pass_id].get(rank) is None - ), f"Duplicated {forward_pass_id=} {rank=}" + assert gpc_of_forward_pass_and_rank[forward_pass_id].get(rank) is None, ( + f"Duplicated {forward_pass_id=} {rank=}" + ) gpc_of_forward_pass_and_rank[forward_pass_id][rank] = record[ "global_physical_count" ] diff --git a/python/sglang/srt/eplb/expert_distribution.py b/python/sglang/srt/eplb/expert_distribution.py index 9ab510d8e..05f5a0e22 100644 --- a/python/sglang/srt/eplb/expert_distribution.py +++ b/python/sglang/srt/eplb/expert_distribution.py @@ -89,9 +89,9 @@ class ExpertDistributionRecorder(ABC): rank: int, ): if get_exec().moe.expert_distribution_recorder_mode is not None: - assert ( - expert_location_metadata is not None - ), "ExpertLocationMetadata is required for expert distribution recording. One possible" + assert expert_location_metadata is not None, ( + "ExpertLocationMetadata is required for expert distribution recording. One possible" + ) "reason is that you are using a model that does not support expert distribution" "recording. Try setting `get_model_config_for_expert_location` in your model." return _ExpertDistributionRecorderReal(expert_location_metadata, rank) @@ -265,9 +265,9 @@ class _ExpertDistributionRecorderReal(ExpertDistributionRecorder): def _reset(self): """Reset the expert distribution recorder.""" logger.info("Resetting ExpertDistributionRecorder...") - assert ( - self._current_layer_idx.value is None - ), f"{self._current_layer_idx.value=}" + assert self._current_layer_idx.value is None, ( + f"{self._current_layer_idx.value=}" + ) for gatherer in self._single_pass_gatherers.values(): gatherer.reset() self._accumulator.reset() @@ -409,9 +409,9 @@ class _DetailSinglePassGatherer(_SinglePassGatherer): device=get_device_namespace().device, ) self._misc_objects: List[Dict[str, Any]] = [] - assert ( - not get_exec().overlap.enable_two_batch_overlap - ), "DetailSinglePassGatherer does not support TBO yet" + assert not get_exec().overlap.enable_two_batch_overlap, ( + "DetailSinglePassGatherer does not support TBO yet" + ) # TODO assert shared experts fusion is disabled, o/w data is wrong def on_forward_pass_start(self, forward_batch: ForwardBatch): @@ -794,7 +794,9 @@ class _UtilizationRateAccumulatorMixin(_Accumulator): assert ( self._expert_location_metadata.ep_size == len(count_of_layer._buckets) - 1 - ), f"{self._expert_location_metadata.ep_size=}, {len(count_of_layer._buckets)=}" + ), ( + f"{self._expert_location_metadata.ep_size=}, {len(count_of_layer._buckets)=}" + ) for gpu_rank in range(self._expert_location_metadata.ep_size): count = gpu_physical_count[layer_idx, gpu_rank] if count > 0: diff --git a/python/sglang/srt/eplb/expert_location.py b/python/sglang/srt/eplb/expert_location.py index c1d16e276..46632b3ff 100644 --- a/python/sglang/srt/eplb/expert_location.py +++ b/python/sglang/srt/eplb/expert_location.py @@ -447,8 +447,7 @@ def format_physical_to_logical_map( row = physical_to_logical_map[layer_id].tolist() if remainder != 0: lines.append( - f"layer={layer_id}: " - f"physical={json.dumps(row, separators=(',', ':'))}" + f"layer={layer_id}: physical={json.dumps(row, separators=(',', ':'))}" ) continue diff --git a/python/sglang/srt/eplb/expert_location_dispatch.py b/python/sglang/srt/eplb/expert_location_dispatch.py index 6325417b1..686ec6ea8 100644 --- a/python/sglang/srt/eplb/expert_location_dispatch.py +++ b/python/sglang/srt/eplb/expert_location_dispatch.py @@ -165,7 +165,7 @@ def _topk_ids_logical_to_physical_probability( """ if not topk_ids.is_cuda: raise RuntimeError( - "LP dispatch requires CUDA tensors; got topk_ids on " f"{topk_ids.device}." + f"LP dispatch requires CUDA tensors; got topk_ids on {topk_ids.device}." ) from sglang.kernels.ops.lplb import cuda_solver diff --git a/python/sglang/srt/function_call/base_format_detector.py b/python/sglang/srt/function_call/base_format_detector.py index dd848402e..e6f97c758 100644 --- a/python/sglang/srt/function_call/base_format_detector.py +++ b/python/sglang/srt/function_call/base_format_detector.py @@ -228,9 +228,9 @@ class BaseFormatDetector(ABC): # Handle parameters/arguments consistency # NOTE: we assume here that the obj is always partial of a single tool call if "parameters" in obj: - assert ( - "arguments" not in obj - ), "model generated both parameters and arguments" + assert "arguments" not in obj, ( + "model generated both parameters and arguments" + ) obj["arguments"] = obj["parameters"] current_tool_call = obj diff --git a/python/sglang/srt/function_call/deepseekv31_detector.py b/python/sglang/srt/function_call/deepseekv31_detector.py index e430cb861..9367bae38 100644 --- a/python/sglang/srt/function_call/deepseekv31_detector.py +++ b/python/sglang/srt/function_call/deepseekv31_detector.py @@ -165,9 +165,9 @@ class DeepSeekV31Detector(BaseFormatDetector): ) ) self._last_arguments += argument_diff - self.streamed_args_for_tool[ - self.current_tool_id - ] += argument_diff + self.streamed_args_for_tool[self.current_tool_id] += ( + argument_diff + ) if _is_complete_json(func_args_raw): # Update the stored arguments diff --git a/python/sglang/srt/function_call/deepseekv3_detector.py b/python/sglang/srt/function_call/deepseekv3_detector.py index 3e744bec9..0a8424ea8 100644 --- a/python/sglang/srt/function_call/deepseekv3_detector.py +++ b/python/sglang/srt/function_call/deepseekv3_detector.py @@ -162,9 +162,9 @@ class DeepSeekV3Detector(BaseFormatDetector): ) ) self._last_arguments += argument_diff - self.streamed_args_for_tool[ - self.current_tool_id - ] += argument_diff + self.streamed_args_for_tool[self.current_tool_id] += ( + argument_diff + ) if _is_complete_json(func_args_raw): # Update the stored arguments diff --git a/python/sglang/srt/function_call/glm47_moe_detector.py b/python/sglang/srt/function_call/glm47_moe_detector.py index d9ddabb65..09f1c3f00 100644 --- a/python/sglang/srt/function_call/glm47_moe_detector.py +++ b/python/sglang/srt/function_call/glm47_moe_detector.py @@ -627,9 +627,9 @@ class Glm47MoeDetector(BaseFormatDetector): pairs = self.func_arg_regex.findall(func_args_raw) if pairs: arguments = self._parse_argument_pairs(pairs, func_name, tools) - self.prev_tool_call_arr[self.current_tool_id][ - "arguments" - ] = arguments + self.prev_tool_call_arr[self.current_tool_id]["arguments"] = ( + arguments + ) except Exception as e: logger.debug(f"Failed to parse arguments: {e}", exc_info=True) diff --git a/python/sglang/srt/function_call/glm4_moe_detector.py b/python/sglang/srt/function_call/glm4_moe_detector.py index 0f065f2a7..baf45711d 100644 --- a/python/sglang/srt/function_call/glm4_moe_detector.py +++ b/python/sglang/srt/function_call/glm4_moe_detector.py @@ -553,9 +553,9 @@ class Glm4MoeDetector(BaseFormatDetector): ) ) self._last_arguments += json_increment - self.streamed_args_for_tool[ - self.current_tool_id - ] += json_increment + self.streamed_args_for_tool[self.current_tool_id] += ( + json_increment + ) if is_tool_end == self.eot_token: if self._is_first_param: @@ -568,9 +568,9 @@ class Glm4MoeDetector(BaseFormatDetector): ) ) self._last_arguments += empty_object - self.streamed_args_for_tool[ - self.current_tool_id - ] += empty_object + self.streamed_args_for_tool[self.current_tool_id] += ( + empty_object + ) else: # The streamed outer `{` is only closed here; a # trailing "}" may belong to a nested object value. @@ -583,9 +583,9 @@ class Glm4MoeDetector(BaseFormatDetector): ) ) self._last_arguments += closing_brace - self.streamed_args_for_tool[ - self.current_tool_id - ] += closing_brace + self.streamed_args_for_tool[self.current_tool_id] += ( + closing_brace + ) try: pairs = self.func_arg_regex.findall(func_args_raw) diff --git a/python/sglang/srt/function_call/internlm_detector.py b/python/sglang/srt/function_call/internlm_detector.py index 1781dbc5e..72c913b85 100644 --- a/python/sglang/srt/function_call/internlm_detector.py +++ b/python/sglang/srt/function_call/internlm_detector.py @@ -113,7 +113,7 @@ class InternlmDetector(BaseFormatDetector): parameters = {} logger.info( - f"[InternLM Tool Call] Parsed tool call #{idx+1}: name={name}, " + f"[InternLM Tool Call] Parsed tool call #{idx + 1}: name={name}, " f"parameters={json.dumps(parameters, ensure_ascii=False)}" ) @@ -136,7 +136,7 @@ class InternlmDetector(BaseFormatDetector): except json.JSONDecodeError as e: logger.error( - f"[InternLM Tool Call] Failed to parse JSON for tool call #{idx+1}: {e}" + f"[InternLM Tool Call] Failed to parse JSON for tool call #{idx + 1}: {e}" ) continue diff --git a/python/sglang/srt/function_call/kimik2_detector.py b/python/sglang/srt/function_call/kimik2_detector.py index 98e0aca70..a8054b587 100644 --- a/python/sglang/srt/function_call/kimik2_detector.py +++ b/python/sglang/srt/function_call/kimik2_detector.py @@ -333,9 +333,9 @@ class KimiK2Detector(BaseFormatDetector): ) if argument_diff: self._last_arguments += argument_diff - self.streamed_args_for_tool[ - self.current_tool_id - ] += argument_diff + self.streamed_args_for_tool[self.current_tool_id] += ( + argument_diff + ) if end_idx == -1: # Args still streaming. diff --git a/python/sglang/srt/function_call/kimik3_structural_tag.py b/python/sglang/srt/function_call/kimik3_structural_tag.py index 04f29a1df..69ba8a35b 100644 --- a/python/sglang/srt/function_call/kimik3_structural_tag.py +++ b/python/sglang/srt/function_call/kimik3_structural_tag.py @@ -511,9 +511,7 @@ def _single_xtml_type( def _nonempty_argument_format(key: str, xtml_type: str) -> Format: # A token-based end keeps the first close token out of both content formats. argument = TagFormat( - begin=( - f'<|open|>argument key="{_escape_attr(key)}" ' f'type="{xtml_type}"<|sep|>' - ), + begin=(f'<|open|>argument key="{_escape_attr(key)}" type="{xtml_type}"<|sep|>'), content=SequenceFormat(elements=[ExcludeTokenFormat(), AnyTokensFormat()]), end=TokenFormat(token=_CLOSE_TOKEN), ) diff --git a/python/sglang/srt/function_call/minimax_m2.py b/python/sglang/srt/function_call/minimax_m2.py index cd5009c99..80cff3f35 100644 --- a/python/sglang/srt/function_call/minimax_m2.py +++ b/python/sglang/srt/function_call/minimax_m2.py @@ -48,9 +48,9 @@ class MinimaxM2Detector(BaseFormatDetector): # Streaming state variables self._current_function_name: str = "" self._current_parameters: Dict[str, Any] = {} - self._streamed_parameters: Dict[str, str] = ( - {} - ) # Track what parameter content we've streamed + self._streamed_parameters: Dict[ + str, str + ] = {} # Track what parameter content we've streamed self._in_tool_call: bool = False self._function_name_sent: bool = False diff --git a/python/sglang/srt/hardware_backend/cpu/quantization/awq_kernels.py b/python/sglang/srt/hardware_backend/cpu/quantization/awq_kernels.py index 4d4c4c483..bd785785c 100644 --- a/python/sglang/srt/hardware_backend/cpu/quantization/awq_kernels.py +++ b/python/sglang/srt/hardware_backend/cpu/quantization/awq_kernels.py @@ -70,9 +70,9 @@ class AWQIntelAMXMoEKernel: ) -> torch.Tensor: from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - assert ( - self.moe_runner_config.activation == "silu" - ), "Only SiLU activation is supported." + assert self.moe_runner_config.activation == "silu", ( + "Only SiLU activation is supported." + ) x = dispatch_output.hidden_states topk_output = dispatch_output.topk_output diff --git a/python/sglang/srt/hardware_backend/cpu/quantization/gptq_kernels.py b/python/sglang/srt/hardware_backend/cpu/quantization/gptq_kernels.py index 53058e905..9dce65900 100644 --- a/python/sglang/srt/hardware_backend/cpu/quantization/gptq_kernels.py +++ b/python/sglang/srt/hardware_backend/cpu/quantization/gptq_kernels.py @@ -70,9 +70,9 @@ class GPTQIntelAMXMoEKernel: ) -> torch.Tensor: from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - assert ( - self.moe_runner_config.activation == "silu" - ), "Only SiLU activation is supported." + assert self.moe_runner_config.activation == "silu", ( + "Only SiLU activation is supported." + ) x = dispatch_output.hidden_states topk_output = dispatch_output.topk_output diff --git a/python/sglang/srt/hardware_backend/gpu/quantization/gptq_kernels.py b/python/sglang/srt/hardware_backend/gpu/quantization/gptq_kernels.py index e3a0c8c52..580a4c22a 100644 --- a/python/sglang/srt/hardware_backend/gpu/quantization/gptq_kernels.py +++ b/python/sglang/srt/hardware_backend/gpu/quantization/gptq_kernels.py @@ -160,7 +160,6 @@ class GPTQMarlinLinearKernel: layer: torch.nn.Module, name: Optional[str], fn: Callable ) -> None: if name is not None and getattr(layer, name, None) is not None: - old_param = getattr(layer, name) new_param = fn(old_param) # replace the parameter with torch.nn.Parameter for TorchDynamo diff --git a/python/sglang/srt/hardware_backend/mlx/aot.py b/python/sglang/srt/hardware_backend/mlx/aot.py index 890ad49eb..dd0235c88 100644 --- a/python/sglang/srt/hardware_backend/mlx/aot.py +++ b/python/sglang/srt/hardware_backend/mlx/aot.py @@ -164,8 +164,7 @@ def _build_rope_kernel(inputs: MlxAOTKernelBuildInputs) -> MlxAOTRoPEKernel: rope_pool_fused = _load_metal_rope_pool_fused() except Exception as exc: # noqa: BLE001 logger.info( - "AOT Metal RoPE kernel not available (%s) - falling back to " - "mx.fast.rope.", + "AOT Metal RoPE kernel not available (%s) - falling back to mx.fast.rope.", exc, ) return MlxAOTRoPEKernel() diff --git a/python/sglang/srt/hardware_backend/mlx/model_runner.py b/python/sglang/srt/hardware_backend/mlx/model_runner.py index e6a296846..6ab8992f2 100644 --- a/python/sglang/srt/hardware_backend/mlx/model_runner.py +++ b/python/sglang/srt/hardware_backend/mlx/model_runner.py @@ -1081,9 +1081,9 @@ class MlxModelRunner: output the scheduler discards; the logit head is skipped when the model exposes a headless trunk. """ - assert ( - req_id in self._req_caches - ), f"extend_start called for unknown request {req_id}" + assert req_id in self._req_caches, ( + f"extend_start called for unknown request {req_id}" + ) cache = self._req_caches[req_id] @@ -1491,8 +1491,7 @@ class MlxModelRunner: cache.state = new_cache.state return raise RuntimeError( - f"Cannot copy {type(new_cache).__name__} state into " - f"{type(cache).__name__}" + f"Cannot copy {type(new_cache).__name__} state into {type(cache).__name__}" ) def _decode_with_native_cache( diff --git a/python/sglang/srt/hardware_backend/mlx/models/muse_glimmer_mlx.py b/python/sglang/srt/hardware_backend/mlx/models/muse_glimmer_mlx.py index 398088fe2..ce5f30c6e 100644 --- a/python/sglang/srt/hardware_backend/mlx/models/muse_glimmer_mlx.py +++ b/python/sglang/srt/hardware_backend/mlx/models/muse_glimmer_mlx.py @@ -182,8 +182,7 @@ def flatten_rc_config(config: dict) -> dict: activation = text.get("hidden_activation", "silu") if activation != "silu": raise ValueError( - f"RC config has hidden_activation={activation!r}; this port " - "hardcodes silu" + f"RC config has hidden_activation={activation!r}; this port hardcodes silu" ) head_dim = int(text.get("head_dim", 128)) diff --git a/python/sglang/srt/hardware_backend/mlx/moe/fused_swiglu.py b/python/sglang/srt/hardware_backend/mlx/moe/fused_swiglu.py index a379c78c3..7f800d2bc 100644 --- a/python/sglang/srt/hardware_backend/mlx/moe/fused_swiglu.py +++ b/python/sglang/srt/hardware_backend/mlx/moe/fused_swiglu.py @@ -285,9 +285,9 @@ def fused_gate_qmv_silu_mul( assert x.shape[-1] == K, f"x last dim {x.shape[-1]} != K={K}" M_tok = x.size // K T = indices.shape[-1] - assert ( - M_tok * T == indices.size - ), f"M_tok({M_tok}) * T({T}) != indices.size({indices.size})" + assert M_tok * T == indices.size, ( + f"M_tok({M_tok}) * T({T}) != indices.size({indices.size})" + ) x_flat = x.reshape(M_tok, K) idx_flat = indices.reshape(M_tok * T) if idx_flat.dtype != mx.uint32: @@ -295,9 +295,9 @@ def fused_gate_qmv_silu_mul( # x_up has N as its last axis and total size M_tok * T * N. The singleton # rank dims (1 or 2 of them) get folded away by reshape. - assert ( - x_up.shape[-1] == N and x_up.size == M_tok * T * N - ), f"x_up shape {x_up.shape} does not match M_tok({M_tok})*T({T})*N({N})" + assert x_up.shape[-1] == N and x_up.size == M_tok * T * N, ( + f"x_up shape {x_up.shape} does not match M_tok({M_tok})*T({T})*N({N})" + ) x_up_flat = x_up.reshape(M_tok * T, N) kernel = _get_kernel(x.dtype) diff --git a/python/sglang/srt/hardware_backend/musa/attention/flashattention_backend.py b/python/sglang/srt/hardware_backend/musa/attention/flashattention_backend.py index 403c7f1f9..fc855da78 100644 --- a/python/sglang/srt/hardware_backend/musa/attention/flashattention_backend.py +++ b/python/sglang/srt/hardware_backend/musa/attention/flashattention_backend.py @@ -927,7 +927,6 @@ class MusaFlashAttentionBackend(FlashAttentionBackend): class MusaFlashAttentionMultiStepBackend(FlashAttentionMultiStepBackend): - def __init__( self, model_runner: ModelRunner, diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py index 0d90b5bc9..b490234de 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -73,7 +73,6 @@ def _reshape_kv_for_fia_nz( @dataclass class ForwardMetadata: - # calculated map for kv positions [bs * maxseqlen] block_tables: Optional[torch.Tensor] = None @@ -301,7 +300,6 @@ def _cp_allgather_and_save_kv_npu( class AscendAttnBackend(AttentionBackend): - def __init__(self, model_runner: ModelRunner, speculative_step_id: int = 0): super().__init__() self.forward_metadata = None diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_dsv4_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_dsv4_backend.py index a662d6cba..6db4fbdf1 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_dsv4_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_dsv4_backend.py @@ -98,7 +98,6 @@ def _build_explicit_state_block_table( class CompressorAscendBackendMixin: - @staticmethod def _to_cpu_int_list(values) -> Optional[list[int]]: if values is None: @@ -451,9 +450,9 @@ class CompressorAscendBackendMixin: coff = 1 + int(compressor.overlap) split = coff * compressor.head_dim w = compressor.wkv_gate.weight - assert ( - w.shape[0] == 2 * split - ), f"wkv_gate.weight rows={w.shape[0]} != 2*coff*head_dim={2*split}" + assert w.shape[0] == 2 * split, ( + f"wkv_gate.weight rows={w.shape[0]} != 2*coff*head_dim={2 * split}" + ) compressor._fused_wkv_w = w[:split] compressor._fused_wgate_w = w[split:] compressor._fused_norm_weight_fp32 = compressor.norm.weight.to(torch.float32) @@ -510,7 +509,6 @@ class CompressorAscendBackendMixin: class C4IndexerAscendBackendMixin: - def init_forward_metadata_indexer(self, core_attn_metadata): # li_quant_metadata is built in _compute_kernel_metadata; None satisfies the mixin contract return None @@ -763,9 +761,9 @@ class C4IndexerAscendBackendMixin: ) -> None: if forward_batch.forward_mode.is_idle(): return - assert ( - not skip_compressor - ), "skip_compressor=True is not supported on the NPU indexer path" + assert not skip_compressor, ( + "skip_compressor=True is not supported on the NPU indexer path" + ) self._ensure_npu_c4_indexer(c4_indexer, x.device) if self._can_use_indexer_multi_stream(): q, weights = self._forward_prepare_multi_stream( @@ -780,7 +778,6 @@ class C4IndexerAscendBackendMixin: class DeepseekV4AscendAttnBackend( AscendAttnBackend, C4IndexerAscendBackendMixin, CompressorAscendBackendMixin ): - def __init__( self, model_runner: ModelRunner, @@ -1074,9 +1071,9 @@ class DeepseekV4AscendAttnBackend( dst.zero_() return n = src.numel() - assert ( - n <= dst.shape[0] - ), f"graph replay 1D metadata overflow: src={n} > dst={dst.shape[0]}" + assert n <= dst.shape[0], ( + f"graph replay 1D metadata overflow: src={n} > dst={dst.shape[0]}" + ) if n > 0: if src.dtype != dst.dtype: src = src.to(dst.dtype) @@ -1924,7 +1921,6 @@ def _get_kv_indices( class DeepseekV4AscendMultiStepDraftBackend: - def __init__( self, model_runner: ModelRunner, diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py index 892066a80..b601a52c2 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py @@ -20,7 +20,6 @@ fused_gdn_gating = fused_gdn_gating_npu class AscendGDNAttnBackend(AscendMambaAttnBackendBase): - def __init__(self, model_runner: ModelRunner): super().__init__(model_runner) self.conv_states_shape = torch.Size( diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_hybrid_linear_attn_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_hybrid_linear_attn_backend.py index 90a52a121..3e60cd9ed 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_hybrid_linear_attn_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_hybrid_linear_attn_backend.py @@ -29,9 +29,9 @@ class AscendMambaAttnBackendBase(MambaAttnBackendBase): self.state_indices_list_gdn = [] def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): - assert ( - max_num_tokens % max_bs == 0 - ), f"max_num_tokens={max_num_tokens} must be divisible by max_bs={max_bs}" + assert max_num_tokens % max_bs == 0, ( + f"max_num_tokens={max_num_tokens} must be divisible by max_bs={max_bs}" + ) draft_token_num = max_num_tokens // max_bs for i in range(max_bs): self.state_indices_list.append( @@ -252,9 +252,7 @@ class AscendHybridLinearAttnBackend(HybridLinearAttnBackend): ] ) - mamba_caches = ( - self.linear_attn_backend.req_to_token_pool.get_speculative_mamba2_params_all_layers() - ) + mamba_caches = self.linear_attn_backend.req_to_token_pool.get_speculative_mamba2_params_all_layers() conv_states = mamba_caches.conv[0] ssm_states = mamba_caches.temporal diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_kda_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_kda_backend.py index 774768569..e00b27ebf 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_kda_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_kda_backend.py @@ -538,9 +538,7 @@ class AscendKDAHybridLinearAttnBackend: ] ) - mamba_caches = ( - self.linear_attn_backend.req_to_token_pool.get_speculative_mamba2_params_all_layers() - ) + mamba_caches = self.linear_attn_backend.req_to_token_pool.get_speculative_mamba2_params_all_layers() conv_states = mamba_caches.conv[0] ssm_states = mamba_caches.temporal diff --git a/python/sglang/srt/hardware_backend/npu/attention/mla_preprocess.py b/python/sglang/srt/hardware_backend/npu/attention/mla_preprocess.py index 1107f11b2..9ec90b5ff 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/mla_preprocess.py +++ b/python/sglang/srt/hardware_backend/npu/attention/mla_preprocess.py @@ -25,9 +25,9 @@ def is_mla_preprocess_enabled() -> bool: def is_fia_nz() -> bool: is_fia_nz_ = get_bool_env_var("SGLANG_USE_FIA_NZ") if is_fia_nz_: - assert ( - is_mla_preprocess_enabled() - ), "SGLANG_USE_FIA_NZ must be enable with SGLANG_NPU_USE_MLAPO" + assert is_mla_preprocess_enabled(), ( + "SGLANG_USE_FIA_NZ must be enable with SGLANG_NPU_USE_MLAPO" + ) return is_fia_nz_ diff --git a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py index eddda0412..8aa279534 100644 --- a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py +++ b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py @@ -549,9 +549,9 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool): if from_indexer: assert ratio == 4, f"indexer only on c4 layers, got ratio={ratio}" if device_type == "npu": - assert ( - self.c4_indexer_kv_pool.has_npu_storage - ), "NPU index buffers not allocated — pool was init'd on CUDA?" + assert self.c4_indexer_kv_pool.has_npu_storage, ( + "NPU index buffers not allocated — pool was init'd on CUDA?" + ) self.c4_indexer_kv_pool.set_index_k_scale( compress_layer_id, loc, kv, kv_scale ) diff --git a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py index 7d8dcb6f4..c91320b24 100644 --- a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py +++ b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py @@ -55,7 +55,6 @@ def _init_npu_conv_state( class NPUMHATokenToKVPool(MHATokenToKVPool): - def __init__( self, size: int, @@ -523,7 +522,6 @@ class NPUMiniMaxSparseKVPool(MiniMaxSparseKVPool): class NPUMLATokenToKVPool(MLATokenToKVPool): - def __init__( self, size: int, diff --git a/python/sglang/srt/hardware_backend/npu/quantization/gptq_kernels.py b/python/sglang/srt/hardware_backend/npu/quantization/gptq_kernels.py index 3a4cb18fa..97d746c33 100644 --- a/python/sglang/srt/hardware_backend/npu/quantization/gptq_kernels.py +++ b/python/sglang/srt/hardware_backend/npu/quantization/gptq_kernels.py @@ -25,12 +25,12 @@ def unpack_from_int32( :param packed_dim: Dimension along which weights are packed (0 or 1), defaults to 1 :return: Unpacked tensor with int8 dtype after applying offset correction """ - assert ( - weight.dtype == torch.int32 - ), f"Expecting `weight.dtype` is torch.int32 but got {weight.dtype}." - assert ( - num_bits <= 8 - ), f"Expecting `num_bits` should not be larger than 8 but got {num_bits}." + assert weight.dtype == torch.int32, ( + f"Expecting `weight.dtype` is torch.int32 but got {weight.dtype}." + ) + assert num_bits <= 8, ( + f"Expecting `num_bits` should not be larger than 8 but got {num_bits}." + ) pack_factor = 32 // num_bits mask = (1 << num_bits) - 1 diff --git a/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py b/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py index aac83c9ca..b761eb99f 100644 --- a/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py +++ b/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py @@ -54,7 +54,6 @@ def _get_float4_e2m1fn_x2_dtype(): class _NPULinearMethodBase(LinearMethodBase): - def __init__( self, quant_config: Optional["QuantizationConfig"] = None, @@ -63,7 +62,6 @@ class _NPULinearMethodBase(LinearMethodBase): class NPUW8A8Int8LinearMethod(_NPULinearMethodBase): - def process_weights_after_loading(self, layer: torch.nn.Module): layer.weight.data = layer.weight.data.transpose(0, 1).contiguous() layer.weight.data = npu_format_cast(layer.weight.data) @@ -121,7 +119,6 @@ class NPUW8A8Int8LinearMethod(_NPULinearMethodBase): class NPUW8A8Int8DynamicLinearMethod(_NPULinearMethodBase): - def process_weights_after_loading(self, layer: torch.nn.Module): layer.weight.data = layer.weight.data.transpose(0, 1).contiguous() layer.weight.data = npu_format_cast(layer.weight.data) @@ -308,7 +305,6 @@ class NPUMXFP8LinearMethod(_NPULinearMethodBase): class NPU_W4A4DynamicLinearMethod(_NPULinearMethodBase): - def process_weights_after_loading(self, layer): layer.weight.data = layer.weight.data.transpose(0, 1).contiguous() layer.weight_scale.data = layer.weight_scale.data.flatten() diff --git a/python/sglang/srt/hardware_backend/xpu/graph_runner/xpu_graph_runner.py b/python/sglang/srt/hardware_backend/xpu/graph_runner/xpu_graph_runner.py index bcc8b4e62..21616c884 100644 --- a/python/sglang/srt/hardware_backend/xpu/graph_runner/xpu_graph_runner.py +++ b/python/sglang/srt/hardware_backend/xpu/graph_runner/xpu_graph_runner.py @@ -116,26 +116,26 @@ class XPUGraphRunner(DecodeCudaGraphRunner): torch._dynamo.config.suppress_errors = True def __init__(self, model_runner: ModelRunner): - assert ( - not get_exec().features.enable_memory_saver - ), "XPUGraphRunner does not support Torch Memory Saver yet." + assert not get_exec().features.enable_memory_saver, ( + "XPUGraphRunner does not support Torch Memory Saver yet." + ) register_fake_ops() self._apply_xpu_compile_config() register_xpu_device_properties_for_dynamo() super().__init__(model_runner) - assert ( - not self.enable_two_batch_overlap - ), "XPUGraphRunner does not support two batch overlap yet." - assert ( - not self.require_mlp_tp_gather - ), "XPUGraphRunner does not support MLP TP gather yet." - assert ( - not self.require_mlp_sync - ), "XPUGraphRunner does not support MLP sync yet." - assert ( - not self.require_gathered_buffer - ), "XPUGraphRunner does not support gathered buffer yet." + assert not self.enable_two_batch_overlap, ( + "XPUGraphRunner does not support two batch overlap yet." + ) + assert not self.require_mlp_tp_gather, ( + "XPUGraphRunner does not support MLP TP gather yet." + ) + assert not self.require_mlp_sync, ( + "XPUGraphRunner does not support MLP sync yet." + ) + assert not self.require_gathered_buffer, ( + "XPUGraphRunner does not support gathered buffer yet." + ) def _init_profile_context_and_memory_record(self): profile_context = profile( diff --git a/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_delta_h.py b/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_delta_h.py index 9f4734a41..2f6a354bc 100644 --- a/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_delta_h.py +++ b/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_delta_h.py @@ -62,9 +62,10 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64_k_loop( i_v, i_nh = tl.program_id(0), tl.program_id(1) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( - cu_seqlens + i_n + 1 - ).to(tl.int32) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) T = eos - bos NT = tl.cdiv(T, BT) boh = tl.load(chunk_offsets + i_n).to(tl.int32) @@ -242,9 +243,9 @@ def chunk_gated_delta_rule_fwd_h( chunk_indices: Optional[torch.LongTensor] = None, use_exp2: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - assert not ( - use_exp2 and g is not None - ), "use_exp2 covers only the per-channel gk path; scalar g stays natural-exp" + assert not (use_exp2 and g is not None), ( + "use_exp2 covers only the per-channel gk path; scalar g stays natural-exp" + ) B, T, Hg, K, V = *k.shape, u.shape[-1] H = u.shape[-2] BT = CHUNK_SIZE diff --git a/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_fwd.py b/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_fwd.py index ead49963c..b327d2578 100644 --- a/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_fwd.py +++ b/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_fwd.py @@ -62,12 +62,14 @@ def chunk_gated_delta_rule_fwd_kkt_solve_kernel_low_reg( i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: - i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load( - chunk_indices + i_t * 2 + 1 - ).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load( - cu_seqlens + i_n + 1 - ).to(tl.int32) + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) T = eos - bos else: bos, eos = i_b * T, i_b * T + T diff --git a/python/sglang/srt/hardware_backend/xpu/kernels/fla/fused_sigmoid_gating_recurrent.py b/python/sglang/srt/hardware_backend/xpu/kernels/fla/fused_sigmoid_gating_recurrent.py index 253e69b2c..4b2499bdd 100644 --- a/python/sglang/srt/hardware_backend/xpu/kernels/fla/fused_sigmoid_gating_recurrent.py +++ b/python/sglang/srt/hardware_backend/xpu/kernels/fla/fused_sigmoid_gating_recurrent.py @@ -53,8 +53,9 @@ def fused_sigmoid_gating_delta_rule_update( stride_a = a.stride()[-2] HV = v.shape[2] N = B if cu_seqlens is None else len(cu_seqlens) - 1 - BK, BV = triton.next_power_of_2(K), min( - triton.next_power_of_2(V), 16 + BK, BV = ( + triton.next_power_of_2(K), + min(triton.next_power_of_2(V), 16), ) # use 16 here to reduce register pressure NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) assert NK == 1, "NK > 1 is not supported yet" diff --git a/python/sglang/srt/kv_canary/api.py b/python/sglang/srt/kv_canary/api.py index ac3cda58d..702518720 100644 --- a/python/sglang/srt/kv_canary/api.py +++ b/python/sglang/srt/kv_canary/api.py @@ -113,9 +113,9 @@ def _patch_model_forward(*, model_runner: ModelRunner, manager: CanaryManager) - return original(*args, **kwargs) forward_batch = _extract_forward_batch(args, kwargs) - assert ( - forward_batch is not None - ), "kv-canary: patched model.forward called without a ForwardBatch" + assert forward_batch is not None, ( + "kv-canary: patched model.forward called without a ForwardBatch" + ) canary_pre_ops_output = manager.pre_ops_maybe_inside_graph(forward_batch) output = original(*args, **kwargs) diff --git a/python/sglang/srt/kv_canary/perturb/utils.py b/python/sglang/srt/kv_canary/perturb/utils.py index d80cbf495..c84ced4cb 100644 --- a/python/sglang/srt/kv_canary/perturb/utils.py +++ b/python/sglang/srt/kv_canary/perturb/utils.py @@ -51,8 +51,7 @@ class WarmupGate: return logger.info( - "kv_canary perturb: disabled during warmup window " - "(first %d forward steps)", + "kv_canary perturb: disabled during warmup window (first %d forward steps)", warmup_steps, ) self._warmup_disable_logged = True diff --git a/python/sglang/srt/kv_canary/runner/canary_manager.py b/python/sglang/srt/kv_canary/runner/canary_manager.py index c04dbb1e8..e0544f840 100644 --- a/python/sglang/srt/kv_canary/runner/canary_manager.py +++ b/python/sglang/srt/kv_canary/runner/canary_manager.py @@ -169,9 +169,9 @@ class CanaryManager: @contextlib.contextmanager def with_active_single_forward_manager(self, index: int) -> Iterator[None]: - assert ( - self._active_single_forward_manager_index is None - ), "kv-canary: nested with_active_single_forward_manager is forbidden" + assert self._active_single_forward_manager_index is None, ( + "kv-canary: nested with_active_single_forward_manager is forbidden" + ) self._active_single_forward_manager_index = index try: yield diff --git a/python/sglang/srt/kv_canary/sweep_plan_builder.py b/python/sglang/srt/kv_canary/sweep_plan_builder.py index e17e9d303..beb5f2f48 100644 --- a/python/sglang/srt/kv_canary/sweep_plan_builder.py +++ b/python/sglang/srt/kv_canary/sweep_plan_builder.py @@ -35,9 +35,9 @@ def build_verify_plan_radix_sweep( prev_slot_indices = walk_result.prev_slot_indices.to(device) if swa_window_size > 0: - assert ( - full_to_swa_index_mapping is not None - ), "full_to_swa_index_mapping is required when SWA is enabled" + assert full_to_swa_index_mapping is not None, ( + "full_to_swa_index_mapping is required when SWA is enabled" + ) slot_indices = _swa_translate( indices=slot_indices, lut=full_to_swa_index_mapping ) diff --git a/python/sglang/srt/layers/activation.py b/python/sglang/srt/layers/activation.py index 8e5d6d048..b620fd701 100644 --- a/python/sglang/srt/layers/activation.py +++ b/python/sglang/srt/layers/activation.py @@ -481,8 +481,7 @@ def get_act_fn( if quant_config is not None and act_fn_name in quant_config.get_scaled_act_names(): if intermediate_size is None: raise ValueError( - "intermediate_size must be specified for scaled " - "activation functions." + "intermediate_size must be specified for scaled activation functions." ) return ScaledActivation( act_fn, intermediate_size, input_is_parallel, params_dtype @@ -495,7 +494,6 @@ def get_cross_encoder_activation_function(config: PretrainedConfig): hasattr(config, "sbert_ce_default_activation_function") and config.sbert_ce_default_activation_function is not None ): - function_name = config.sbert_ce_default_activation_function assert function_name.startswith("torch.nn.modules."), ( "Loading of activation functions is restricted to " diff --git a/python/sglang/srt/layers/amx_utils.py b/python/sglang/srt/layers/amx_utils.py index 485c4e5e9..27677dccf 100755 --- a/python/sglang/srt/layers/amx_utils.py +++ b/python/sglang/srt/layers/amx_utils.py @@ -133,9 +133,9 @@ def _amx_process_weight_after_loading( device = devices.pop() if transpose_dims: - assert len(weight_names) == len( - transpose_dims - ), "len(weight_names) should be equal to len(transpose_dims)" + assert len(weight_names) == len(transpose_dims), ( + "len(weight_names) should be equal to len(transpose_dims)" + ) module.use_intel_amx_backend = ( device == torch.device("cpu") and cpu_has_amx_support() diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index 578fe3152..1f9bc87d3 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -211,7 +211,6 @@ def _asm_context_prefill_gather_indices( class AiterAttnBackend(AttentionBackend): - # kv_indptr/qo_indptr are preallocated at (req pool + 1); an extend batch # can never carry more seqs than the pool. extend_dummy_seqs_capped_by_req_pool: bool = True diff --git a/python/sglang/srt/layers/attention/attention_registry.py b/python/sglang/srt/layers/attention/attention_registry.py index d60c4de0b..bbf45f359 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/python/sglang/srt/layers/attention/attention_registry.py @@ -79,7 +79,6 @@ def create_trtllm_mla_backend(runner): if not runner.use_mla_backend: raise ValueError("trtllm_mla backend can only be used with MLA models.") if get_parallel().dcp_enabled and get_spec().speculative_algorithm is not None: - _, decode_backend = attention_backends_of(resolved_view(runner.server_args)) if decode_backend == "trtllm_mla": raise ValueError( @@ -449,7 +448,9 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac assert ( runner.prefill_attention_backend_str == "ascend" and runner.decode_attention_backend_str == "ascend" - ), "ascend backend is the only supported backend on NPU for hybrid GDN models, use --attention-backend ascend to specify the backend." + ), ( + "ascend backend is the only supported backend on NPU for hybrid GDN models, use --attention-backend ascend to specify the backend." + ) logger.info(f"Using hybrid linear attention backend for hybrid GDN models.") linear_attn_backend = GDNAttnBackend(runner) elif mamba2_config(runner.model_config) is not None: diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 67dc74e18..d7cc363d4 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -294,9 +294,9 @@ class DSV4AttnMetadata: # length; legacy paths use the metadata length. if num_tokens is None: num_tokens = self.seq_lens_casual.shape[0] - assert ( - self.raw_out_loc.shape[0] == num_tokens - ), f"{self.raw_out_loc.shape=}, {num_tokens=}" + assert self.raw_out_loc.shape[0] == num_tokens, ( + f"{self.raw_out_loc.shape=}, {num_tokens=}" + ) ( self.c4_out_loc, @@ -354,9 +354,9 @@ class DSV4AttnMetadata: num_tokens = pre_global_len for field_name in self._CP_REINDEX_FIELDS: val = getattr(self, field_name, None) - assert isinstance( - val, torch.Tensor - ), f"CP reindex: {field_name} is {type(val)}, expected Tensor" + assert isinstance(val, torch.Tensor), ( + f"CP reindex: {field_name} is {type(val)}, expected Tensor" + ) setattr(self, field_name, val[idx].contiguous()) for field_name in self._CP_REINDEX_FIELDS: @@ -542,9 +542,9 @@ class DeepseekV4AttnBackend( self.device = torch.device(model_runner.device) self.max_context_len = model_runner.model_config.context_len head_dim = model_runner.model_config.head_dim - assert ( - head_dim == 512 - ), "DSV4 MQA head_dim = qk_nope_head_dim(448) + qk_rope_head_dim(64) = 512" + assert head_dim == 512, ( + "DSV4 MQA head_dim = qk_nope_head_dim(448) + qk_rope_head_dim(64) = 512" + ) self.softmax_scale: float = head_dim**-0.5 self.head_dim_v: int = model_runner.model_config.v_head_dim self.cuda_int32_kwargs = {"device": self.device, "dtype": torch.int32} @@ -1750,13 +1750,13 @@ class DeepseekV4AttnBackend( flashmla_metadata = core_attn_metadata.get_flashmla_metadata(compress_ratio) - assert ( - swa_page_indices.shape[-1] % 64 == 0 - ), f"{swa_page_indices.shape=}'s last dimension is not aligned to 64" + assert swa_page_indices.shape[-1] % 64 == 0, ( + f"{swa_page_indices.shape=}'s last dimension is not aligned to 64" + ) if extra_indices is not None: - assert ( - extra_indices.shape[-1] % 64 == 0 - ), f"{extra_indices.shape=}'s last dimension is not aligned to 64" + assert extra_indices.shape[-1] % 64 == 0, ( + f"{extra_indices.shape=}'s last dimension is not aligned to 64" + ) # sparse_prefill_fwd does not support SM120. if ( diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index 92c39e57e..d1caca5ca 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py @@ -239,9 +239,9 @@ class DSV4AttnMetadata: def init_compression_metadata(self, unified_swa_pages: int = 0): assert self.page_table.dim() == 2 - assert ( - self.raw_out_loc.shape == self.seq_lens_casual.shape - ), f"{self.raw_out_loc.shape=}, {self.seq_lens_casual.shape=}" + assert self.raw_out_loc.shape == self.seq_lens_casual.shape, ( + f"{self.raw_out_loc.shape=}, {self.seq_lens_casual.shape=}" + ) ( self.c4_out_loc, @@ -302,9 +302,9 @@ class DSV4AttnMetadata: expected_local_len = pre_global_len // cp_size for field_name in self._CP_REINDEX_FIELDS: val = getattr(self, field_name, None) - assert isinstance( - val, torch.Tensor - ), f"CP reindex: {field_name} is {type(val)}, expected Tensor" + assert isinstance(val, torch.Tensor), ( + f"CP reindex: {field_name} is {type(val)}, expected Tensor" + ) setattr(self, field_name, val[idx].contiguous()) for field_name in self._CP_REINDEX_FIELDS: @@ -451,9 +451,9 @@ class DeepseekV4HipRadixBackend( super().__init__() self.device = torch.device(model_runner.device) head_dim = model_runner.model_config.head_dim - assert ( - head_dim == 512 - ), "DSV4 MQA head_dim = qk_nope_head_dim(448) + qk_rope_head_dim(64) = 512" + assert head_dim == 512, ( + "DSV4 MQA head_dim = qk_nope_head_dim(448) + qk_rope_head_dim(64) = 512" + ) self.softmax_scale: float = head_dim**-0.5 self.head_dim_v: int = model_runner.model_config.v_head_dim self.cuda_int32_kwargs = {"device": self.device, "dtype": torch.int32} @@ -1674,13 +1674,13 @@ class DeepseekV4HipRadixBackend( flashmla_metadata = core_attn_metadata.get_flashmla_metadata(compress_ratio) - assert ( - swa_page_indices.shape[-1] % 64 == 0 - ), f"{swa_page_indices.shape=}'s last dimension is not aligned to 64" + assert swa_page_indices.shape[-1] % 64 == 0, ( + f"{swa_page_indices.shape=}'s last dimension is not aligned to 64" + ) if extra_indices is not None: - assert ( - extra_indices.shape[-1] % 64 == 0 - ), f"{extra_indices.shape=}'s last dimension is not aligned to 64" + assert extra_indices.shape[-1] % 64 == 0, ( + f"{extra_indices.shape=}'s last dimension is not aligned to 64" + ) from sglang.srt.layers.attention.hip_flash_mla import ( flash_mla_with_kvcache_entrypoint, diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index 5674b65f2..3969dfbee 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py @@ -193,9 +193,9 @@ def rotate_activation(x: torch.Tensor) -> torch.Tensor: from sglang.kernels.ops.quantization.hadamard import hadamard_transform hidden_size = x.size(-1) - assert ( - hidden_size & (hidden_size - 1) - ) == 0, "Hidden size must be a power of 2 for Hadamard transform." + assert (hidden_size & (hidden_size - 1)) == 0, ( + "Hidden size must be a power of 2 for Hadamard transform." + ) return hadamard_transform(x, scale=hidden_size**-0.5) @@ -806,13 +806,13 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp): # NOTE(dark): blocksize = 64 is hardcoded in deep_gemm if _is_hip: if _use_aiter_preshuffle: - assert ( - page_size % 16 == 0 - ), f"HIP preshuffle requires page_size to be a multiple of 16, got {page_size}" + assert page_size % 16 == 0, ( + f"HIP preshuffle requires page_size to be a multiple of 16, got {page_size}" + ) else: - assert ( - page_size == 1 - ), f"HIP legacy DSA path requires page_size == 1, got {page_size}" + assert page_size == 1, ( + f"HIP legacy DSA path requires page_size == 1, got {page_size}" + ) else: assert page_size == 64, "only support page size 64" # NOTE(dark): this support extend/decode/decode+graph @@ -1041,13 +1041,13 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp): page_size = get_token_to_kv_pool().page_size if _is_hip: if _use_aiter_preshuffle: - assert ( - page_size % 16 == 0 - ), f"HIP preshuffle requires page_size to be a multiple of 16, got {page_size}" + assert page_size % 16 == 0, ( + f"HIP preshuffle requires page_size to be a multiple of 16, got {page_size}" + ) else: - assert ( - page_size == 1 - ), f"HIP legacy DSA path requires page_size == 1, got {page_size}" + assert page_size == 1, ( + f"HIP legacy DSA path requires page_size == 1, got {page_size}" + ) else: assert page_size == 64, "only support page size 64" @@ -1159,13 +1159,13 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp): if global_topk_offset is None: cu_seqlens_q_full = torch.ones(q_offset, dtype=torch.int32, device=device) - assert ( - seq_lens_expanded.shape[0] == q_offset - ), f"seq_lens_expanded length mismatch: {seq_lens_expanded.shape[0]} != {q_offset}" + assert seq_lens_expanded.shape[0] == q_offset, ( + f"seq_lens_expanded length mismatch: {seq_lens_expanded.shape[0]} != {q_offset}" + ) if global_topk_offset is not None: - assert ( - global_topk_offset.shape[0] >= q_offset - ), f"topk_indices_offset too short: {global_topk_offset.shape[0]} < {q_offset}" + assert global_topk_offset.shape[0] >= q_offset, ( + f"topk_indices_offset too short: {global_topk_offset.shape[0]} < {q_offset}" + ) start = 0 while start < q_offset: diff --git a/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py b/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py index 4fae2acbb..2ff9256f9 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py @@ -319,7 +319,9 @@ def _topk_transform_v2_paged( logits.dtype == torch.float32 and logits.stride(1) == 1 and logits.stride(0) % 4 == 0 - ), f"v2 top-k expects fp32 scores with unit row stride and 16B-aligned score_stride, got {logits.dtype=} {logits.stride()=}" + ), ( + f"v2 top-k expects fp32 scores with unit row stride and 16B-aligned score_stride, got {logits.dtype=} {logits.stride()=}" + ) assert 0 < topk <= 2048, f"v2 top-k supports 0 < topk <= 2048, got {topk=}" page_table = attn_metadata.real_page_table @@ -329,9 +331,9 @@ def _topk_transform_v2_paged( # mismatched plan means the caller skipped that preprocessing -- fail loudly # rather than silently recompute it per layer. plan = attn_metadata.topk_v2_plan - assert ( - plan is not None and plan.shape[0] == num_rows + 1 - ), "topk_v2_plan must be preprocessed per forward (see DSAMetadata.topk_v2_plan)" + assert plan is not None and plan.shape[0] == num_rows + 1, ( + "topk_v2_plan must be preprocessed per forward (see DSAMetadata.topk_v2_plan)" + ) page_size = attn_metadata.page_size out = logits.new_empty((num_rows, topk), dtype=torch.int32) diff --git a/python/sglang/srt/layers/attention/dsa/kpool_fp8_index.py b/python/sglang/srt/layers/attention/dsa/kpool_fp8_index.py index f9bf23137..db8d3a88a 100644 --- a/python/sglang/srt/layers/attention/dsa/kpool_fp8_index.py +++ b/python/sglang/srt/layers/attention/dsa/kpool_fp8_index.py @@ -19,9 +19,9 @@ def build_pooled_page_table_64( ) -> torch.Tensor: # Advanced indexing is required: a (1, 1) strided slice can remain non-unit- # stride even after contiguous(), which DeepGEMM rejects. - assert ( - BLOCK_SIZE_K % pool_size == 0 - ), f"pool_size ({pool_size}) must divide page_size ({BLOCK_SIZE_K})" + assert BLOCK_SIZE_K % pool_size == 0, ( + f"pool_size ({pool_size}) must divide page_size ({BLOCK_SIZE_K})" + ) idx = torch.arange( 0, page_table_64.shape[-1], pool_size, device=page_table_64.device ) @@ -218,9 +218,9 @@ def _prep_update_kpool_write_plan_launch( assert write_loc_out.stride(1) == 1, write_loc_out.stride() has_per_q_outputs = pool_seqlens_per_q_out is not None - assert has_per_q_outputs == ( - seqlens_per_q_out is not None - ), "pool_seqlens_per_q_out and seqlens_per_q_out must be both set or both None" + assert has_per_q_outputs == (seqlens_per_q_out is not None), ( + "pool_seqlens_per_q_out and seqlens_per_q_out must be both set or both None" + ) per_q_dummy = ( pool_seqlens_per_q_out if has_per_q_outputs @@ -596,9 +596,9 @@ def topk_from_pooled_history_logits( padded[: result.shape[0]] = result return padded - assert ( - page_table_row_index is None - ), "page_table_row_index requires the fused fast_kpool group_topk path" + assert page_table_row_index is None, ( + "page_table_row_index requires the fused fast_kpool group_topk path" + ) from sgl_kernel import fast_topk_v2 diff --git a/python/sglang/srt/layers/attention/dsa/utils.py b/python/sglang/srt/layers/attention/dsa/utils.py index e469a6116..d9995a94d 100644 --- a/python/sglang/srt/layers/attention/dsa/utils.py +++ b/python/sglang/srt/layers/attention/dsa/utils.py @@ -264,9 +264,9 @@ def can_dsa_cp_split(seq_len: int, cp_size: int, use_dsa: bool, forward_batch): if is_dsa_prefill_cp_round_robin_split(): cur_cp_seq_len = seq_len // cp_size - assert ( - seq_len % cp_size == 0 - ), f"seq_len {seq_len} is not divisible by cp_size {cp_size} when dsa_prefill_cp_mode is round-robin-split" + assert seq_len % cp_size == 0, ( + f"seq_len {seq_len} is not divisible by cp_size {cp_size} when dsa_prefill_cp_mode is round-robin-split" + ) else: # TODO current just support prefill batch=1 and len(input_ids) > self.cp_size * 2 # Note: (self.cp_size * 2) To achieve load balancing for seq computation, diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index 9369ad289..5a5eb5b6b 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -289,7 +289,6 @@ _DSA_IMPL_T: TypeAlias = Literal[ class DeepseekSparseAttnBackend( DeepseekSparseAttnBackendMTPPrecomputeMixin, AttentionBackend ): - # kv_indptr/qo_indptr are preallocated at (req pool + 1); an extend batch # can never carry more seqs than the pool. extend_dummy_seqs_capped_by_req_pool: bool = True @@ -980,9 +979,9 @@ class DeepseekSparseAttnBackend( ) ] ) - assert page_table_1_flattened.shape[0] == sum( - indexer_seq_lens_cpu - ), f"{page_table_1_flattened.shape[0] = } must be the same as {sum(indexer_seq_lens_cpu) = }" + assert page_table_1_flattened.shape[0] == sum(indexer_seq_lens_cpu), ( + f"{page_table_1_flattened.shape[0] = } must be the same as {sum(indexer_seq_lens_cpu) = }" + ) # Validate indices when logical tokens exceed physical capacity # This is likely to be triggered by PP with high kv reuse & parallelism @@ -1938,9 +1937,9 @@ class DeepseekSparseAttnBackend( if self.use_mha: assert k is not None and v is not None assert q_rope is None, "MHA_ONE_SHOT path should not pass q_rope" - assert ( - layer.tp_k_head_num == layer.tp_q_head_num > 1 - ), "MHA_ONE_SHOT requires dense multi-head config" + assert layer.tp_k_head_num == layer.tp_q_head_num > 1, ( + "MHA_ONE_SHOT requires dense multi-head config" + ) return self._forward_standard_mha( q=q, k=k, @@ -2875,8 +2874,8 @@ class DeepseekSparseAttnBackend( # Verify batch sizes match (length of cu_seqlens should be batch_size + 1) assert len(cu_seqlens_q) == len(cu_seqlens_k), ( - f"batch_size mismatch: cu_seqlens_q has {len(cu_seqlens_q)-1} requests, " - f"cu_seqlens_k has {len(cu_seqlens_k)-1} requests" + f"batch_size mismatch: cu_seqlens_q has {len(cu_seqlens_q) - 1} requests, " + f"cu_seqlens_k has {len(cu_seqlens_k) - 1} requests" ) # Use TRTLLm ragged attention for SM100 (Blackwell/B200) to avoid FA4 accuracy issues. @@ -3135,9 +3134,9 @@ class DeepseekSparseAttnBackend( # Note: rope application in deepseek_v2.py:forward_absorb_prepare is skipped for FP8 decode path of this trtllm_mla backend assert q_rope is not None, "For FP8 path q_rope should not be None." assert k_rope is not None, "For FP8 path k_rope should not be None." - assert ( - cos_sin_cache is not None - ), "For FP8 path cos_sin_cache should not be None." + assert cos_sin_cache is not None, ( + "For FP8 path cos_sin_cache should not be None." + ) rope_positions = forward_batch.positions if dsa_use_prefill_cp(forward_batch): @@ -3172,9 +3171,9 @@ class DeepseekSparseAttnBackend( # Save KV cache if requested if save_kv_cache: - assert ( - k is not None and k_rope is not None - ), "For populating trtllm_mla kv cache, both k_nope and k_rope should be not None." + assert k is not None and k_rope is not None, ( + "For populating trtllm_mla kv cache, both k_nope and k_rope should be not None." + ) cache_loc = ( forward_batch.out_cache_loc if not layer.is_cross_attention @@ -3425,7 +3424,6 @@ class DeepseekSparseAttnBackend( class DeepseekSparseAttnMultiStepBackend: - # Per-step draft decode replays from precomputed GPU metadata; opt out so # decide_needs_cpu_seq_lens' OR over the backends stays False. needs_cpu_seq_lens: bool = False diff --git a/python/sglang/srt/layers/attention/dsv4/compress_hip.py b/python/sglang/srt/layers/attention/dsv4/compress_hip.py index caf37976d..225a008a3 100644 --- a/python/sglang/srt/layers/attention/dsv4/compress_hip.py +++ b/python/sglang/srt/layers/attention/dsv4/compress_hip.py @@ -223,9 +223,9 @@ class CompressorHip(_CompressorBase): assert kv_compressed.dtype == torch.float32 freqs_cis = self.freqs_cis[beg_idx : end_idx : self.ratio] - assert freqs_cis.size(0) == kv_compressed.size( - 0 - ), f"{freqs_cis.shape=} {kv_compressed.shape=}" + assert freqs_cis.size(0) == kv_compressed.size(0), ( + f"{freqs_cis.shape=} {kv_compressed.shape=}" + ) fused_norm_rope_inplace_triton( kv_compressed, self.norm.weight, self.norm.eps, freqs_cis ) diff --git a/python/sglang/srt/layers/attention/dsv4/indexer.py b/python/sglang/srt/layers/attention/dsv4/indexer.py index ea4608eb8..6f21371e7 100644 --- a/python/sglang/srt/layers/attention/dsv4/indexer.py +++ b/python/sglang/srt/layers/attention/dsv4/indexer.py @@ -214,9 +214,9 @@ def fp8_paged_mqa_logits_torch_sm120( ) assert head_dim == 128, "Vectorized torch impl hardcodes DSV4 indexer head_dim=128" - assert ( - block_size == 64 - ), "Vectorized torch impl hardcodes block_size=64 cache layout" + assert block_size == 64, ( + "Vectorized torch impl hardcodes block_size=64 cache layout" + ) assert q_fp8.shape == (batch_size, 1, num_heads, head_dim) assert kvcache_fp8.shape[1:] == (block_size, 1, head_dim + 4) assert weight.shape == (batch_size, num_heads) diff --git a/python/sglang/srt/layers/attention/dsv4/metadata.py b/python/sglang/srt/layers/attention/dsv4/metadata.py index 4c6c3d663..bf7ce9670 100644 --- a/python/sglang/srt/layers/attention/dsv4/metadata.py +++ b/python/sglang/srt/layers/attention/dsv4/metadata.py @@ -89,14 +89,14 @@ def copy_metadata( provided_fields = check_eq_fields + copy_fields + assign_fields provided_fields_unique = set(provided_fields) - assert len(provided_fields) == len( - provided_fields_unique - ), f"{provided_fields=} has dup" + assert len(provided_fields) == len(provided_fields_unique), ( + f"{provided_fields=} has dup" + ) all_fields = {f.name for f in fields(src)} provided_fields = set(provided_fields) - assert ( - provided_fields == all_fields - ), f"{provided_fields - all_fields=}, {all_fields - provided_fields=}" + assert provided_fields == all_fields, ( + f"{provided_fields - all_fields=}, {all_fields - provided_fields=}" + ) @dataclass diff --git a/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py b/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py index 3ca85e5f2..80b9fa897 100644 --- a/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py +++ b/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py @@ -166,9 +166,9 @@ def combine_topk_swa_indices( assert compressed_base.dtype == torch.int32 assert swa_base.dtype == torch.int32 assert compress_ratio >= 1, "compress_ratio must be >= 1 (use topk=0 for SWA-only)" - assert ( - topk_indices.shape[-1] >= topk - ), f"topk_indices width {topk_indices.shape[-1]} must be >= topk {topk}" + assert topk_indices.shape[-1] >= topk, ( + f"topk_indices width {topk_indices.shape[-1]} must be >= topk {topk}" + ) num_tokens = topk_indices.shape[0] num_reqs = seq_lens.shape[0] @@ -473,9 +473,9 @@ class SparsePrefillChunkCache: device = self.seq_lens.device c4_max = max(self.max_seq_len // 4, 1) c4_capacity = page_table.shape[-1] * c4_page_size - assert ( - c4_max <= c4_capacity - ), f"live c4 extent {c4_max} exceeds metadata capacity {c4_capacity}" + assert c4_max <= c4_capacity, ( + f"live c4 extent {c4_max} exceeds metadata capacity {c4_capacity}" + ) first_q_per_req = self.query_start_loc[:-1].long() num_blocks = (c4_max + c4_page_size - 1) // c4_page_size assert num_blocks <= page_table.shape[1] diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index cb962e6ed..d228b4048 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -235,9 +235,9 @@ class FlashAttentionBackend(AttentionBackend): "on the fa3/fa4 backend." ) if self.has_local_attention: - assert ( - model_runner.attention_chunk_size is not None - ), "Attention chunk size is required for local attention" + assert model_runner.attention_chunk_size is not None, ( + "Attention chunk size is required for local attention" + ) self.attention_chunk_size = model_runner.attention_chunk_size # For each layer, the sliding_window_size can be different. This is only used for preparing SWA metadata. @@ -930,9 +930,7 @@ class FlashAttentionBackend(AttentionBackend): # create expand page table offsets = torch.arange( self.speculative_num_draft_tokens, device=device - ).unsqueeze( - 0 - ) # shape: (1, self.speculative_num_draft_tokens) + ).unsqueeze(0) # shape: (1, self.speculative_num_draft_tokens) cols = offsets.expand( forward_batch.seq_lens.numel(), -1 ) + forward_batch.seq_lens.unsqueeze(1) @@ -1094,9 +1092,7 @@ class FlashAttentionBackend(AttentionBackend): ) text_col = forward_batch.encoder_lens.long().unsqueeze( 1 - ) + arange_text.unsqueeze( - 0 - ) # (bs, max_seq_len_k) + ) + arange_text.unsqueeze(0) # (bs, max_seq_len_k) text_row = forward_batch.req_pool_indices.unsqueeze(1).expand(-1, text_max) metadata.page_table = self.req_to_token_pool.req_to_token[ text_row, text_col @@ -1717,9 +1713,9 @@ class FlashAttentionBackend(AttentionBackend): # Concat q_nope + q_rope along dim=-1 so the wrapper's # chunk(2, dim=0) keeps their alignment; split back # inside the closure. - assert ( - not use_cascade_attn - ), "Cascade attention under MLA CP is not supported in v1." + assert not use_cascade_attn, ( + "Cascade attention under MLA CP is not supported in v1." + ) q_fused = torch.cat([q_nope, q_rope], dim=-1) def _mla_cp_attn( @@ -2987,9 +2983,7 @@ class FlashAttentionBackend(AttentionBackend): # metadata_expand.cu_seqlens_q already set in capture offsets = torch.arange( self.speculative_num_draft_tokens, device=device - ).unsqueeze( - 0 - ) # shape: (1, self.speculative_num_draft_tokens) + ).unsqueeze(0) # shape: (1, self.speculative_num_draft_tokens) cols = offsets.expand(seq_lens.numel(), -1) + seq_lens.unsqueeze(1) cum_len = torch.nn.functional.pad( @@ -3334,9 +3328,9 @@ class FlashAttentionBackend(AttentionBackend): metadata_swa: Optional[FlashAttentionMetadata] = None, ): # TODO: support page_size > 1 for swa spec - assert ( - self.page_size == 1 - ), "FlashAttention backend doesn't support topk > 1 speculative decoding with page size > 1 sliding window attention" + assert self.page_size == 1, ( + "FlashAttention backend doesn't support topk > 1 speculative decoding with page size > 1 sliding window attention" + ) cache_seqlens_int32 = ( metadata.cache_seqlens_int32.repeat_interleave( @@ -3640,8 +3634,7 @@ def make_local_attention_virtual_batches( block_starts = k_seqstarts_absolute // page_size assert attn_chunk_size % page_size == 0, ( - f"attn_chunk_size {attn_chunk_size} is not " - f"divisible by page_size {page_size}" + f"attn_chunk_size {attn_chunk_size} is not divisible by page_size {page_size}" ) pages_per_local_batch = attn_chunk_size // page_size diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index 504456fd9..40b0eb35c 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_backend.py @@ -218,12 +218,12 @@ def fast_prefill_plan( is identical to plan()'s. """ assert self.is_cuda_graph_enabled, "fast_prefill_plan is cuda-graph only" - assert ( - getattr(self, "_backend", None) == "fa2" - ), "fast_prefill_plan supports the fa2 backend only" - assert ( - getattr(self, "_cached_module", None) is not None - ), "fast_prefill_plan requires _cached_module from a prior real plan() (capture)" + assert getattr(self, "_backend", None) == "fa2", ( + "fast_prefill_plan supports the fa2 backend only" + ) + assert getattr(self, "_cached_module", None) is not None, ( + "fast_prefill_plan requires _cached_module from a prior real plan() (capture)" + ) if head_dim_vo is None: head_dim_vo = head_dim_qk @@ -1309,9 +1309,9 @@ class FlashInferAttnBackend(AttentionBackend): q = q.contiguous() - assert not ( - self.prefill_uses_dequant_workspace and layer.is_cross_attention - ), "FP4 dequant KV cache is not supported for cross-attention" + assert not (self.prefill_uses_dequant_workspace and layer.is_cross_attention), ( + "FP4 dequant KV cache is not supported for cross-attention" + ) # We perform dequant for chunk prefill/cache reuse. pool = self.token_to_kv_pool @@ -1377,9 +1377,9 @@ class FlashInferAttnBackend(AttentionBackend): # previously cached context without re-materializing KV tensors (e.g., the # IQuestLoopCoder path uses token_to_kv_pool as the KV source). if k is None and v is None: - assert ( - not self.prefill_uses_dequant_workspace - ), "KV cache must be provided for ragged attention when using FP4 dequant KV cache" + assert not self.prefill_uses_dequant_workspace, ( + "KV cache must be provided for ragged attention when using FP4 dequant KV cache" + ) k = self.token_to_kv_pool.get_kv_buffer(layer.layer_id)[0] v = self.token_to_kv_pool.get_kv_buffer(layer.layer_id)[1] causal = True @@ -2241,15 +2241,15 @@ class FlashInferIndicesUpdaterPrefill: and wrapper_paged.begin_forward.func is fast_prefill_plan ) if uses_fast_prefill: - assert ( - seq_lens_cpu is not None - ), "fast_prefill_plan replay requires host-known seq_lens_cpu (got None)" - assert ( - num_tokens_per_req is not None and num_tokens_per_req > 0 - ), f"fast_prefill_plan replay requires num_tokens_per_req > 0 (got {num_tokens_per_req})" - assert ( - use_custom_mask is None - ), "fast_prefill_plan does not support custom_mask; keep the plain plan()" + assert seq_lens_cpu is not None, ( + "fast_prefill_plan replay requires host-known seq_lens_cpu (got None)" + ) + assert num_tokens_per_req is not None and num_tokens_per_req > 0, ( + f"fast_prefill_plan replay requires num_tokens_per_req > 0 (got {num_tokens_per_req})" + ) + assert use_custom_mask is None, ( + "fast_prefill_plan does not support custom_mask; keep the plain plan()" + ) seq_lens_cpu_i32 = seq_lens_cpu.to(torch.int32) qo_indptr_host = torch.arange( 0, diff --git a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py index 1a7527a6d..a58424d7a 100644 --- a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py @@ -564,9 +564,9 @@ class FlashInferMLAAttnBackend(AttentionBackend): ) elif forward_mode.is_target_verify(): # build host indptr/len arrays for target-verify fast plan path - assert ( - seq_lens_cpu is not None and spec_info is not None - ), "target-verify cuda-graph replay requires host-resident seq_lens_cpu" + assert seq_lens_cpu is not None and spec_info is not None, ( + "target-verify cuda-graph replay requires host-resident seq_lens_cpu" + ) ndt = spec_info.draft_token_num self.fast_plan_qo_indptr_cpu[: bs + 1] = torch.arange( 0, (bs + 1) * ndt, ndt, dtype=torch.int32 diff --git a/python/sglang/srt/layers/attention/flashmla_backend.py b/python/sglang/srt/layers/attention/flashmla_backend.py index 301a06d79..d39171052 100644 --- a/python/sglang/srt/layers/attention/flashmla_backend.py +++ b/python/sglang/srt/layers/attention/flashmla_backend.py @@ -440,9 +440,9 @@ class FlashMLABackend(FlashInferMLAAttnBackend): reshape_q = q.view(bs, -1, layer.tp_q_head_num, layer.head_dim) if self.is_fp8_kvcache: - assert ( - self.dcp_world_size == 1 - ), "FlashMLA does not support DCP for FP8 kv cache" + assert self.dcp_world_size == 1, ( + "FlashMLA does not support DCP for FP8 kv cache" + ) if layer.k_scale is not None: q_scale = layer.k_scale descale_q = layer.k_scale.reshape(1) diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index 126e48bdd..ce625166b 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -438,9 +438,9 @@ class MambaAttnBackendBase(AttentionBackend): return mask.cpu() def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): - assert ( - max_num_tokens % max_bs == 0 - ), f"max_num_tokens={max_num_tokens} must be divisible by max_bs={max_bs}" + assert max_num_tokens % max_bs == 0, ( + f"max_num_tokens={max_num_tokens} must be divisible by max_bs={max_bs}" + ) draft_token_num = max_num_tokens // max_bs # Per-bs static write-cursor / force-flush buffers, captured by pointer + # refreshed in-place each replay; sized like state_indices_list. None when off. @@ -495,9 +495,9 @@ class MambaAttnBackendBase(AttentionBackend): ) def init_cpu_graph_state(self, max_bs: int, max_num_tokens: int): - assert ( - max_num_tokens % max_bs == 0 - ), f"max_num_tokens={max_num_tokens} must be divisible by max_bs={max_bs}" + assert max_num_tokens % max_bs == 0, ( + f"max_num_tokens={max_num_tokens} must be divisible by max_bs={max_bs}" + ) for i in range(max_bs): self.state_indices_list.append( torch.full( @@ -628,9 +628,9 @@ class MambaAttnBackendBase(AttentionBackend): # [-num_decodes:], which on the full max_bs buffer binds the stale tail. track_buf = None if mamba_track_indices is not None: - assert ( - len(mamba_track_indices) >= bs - ), f"{len(mamba_track_indices)=} < {bs=}" + assert len(mamba_track_indices) >= bs, ( + f"{len(mamba_track_indices)=} < {bs=}" + ) track_buf = self.mamba_track_indices_buf[:bs] track_buf.copy_(self._translate_mamba_indices(mamba_track_indices[:bs])) # Refresh the static write cursor in-place (mirrors the eager @@ -877,12 +877,14 @@ class Mamba2AttnBackend(MambaAttnBackendBase): ) if model_runner.server_args.enable_mamba_extra_buffer(): - assert ( - self.conv_states_shape[-1] < self.mamba_chunk_size - ), f"{self.conv_states_shape[-1]=} should be less than {self.mamba_chunk_size}" + assert self.conv_states_shape[-1] < self.mamba_chunk_size, ( + f"{self.conv_states_shape[-1]=} should be less than {self.mamba_chunk_size}" + ) assert ( model_runner.server_args.mamba_track_interval >= self.mamba_chunk_size - ), f"mamba_track_interval ({model_runner.server_args.mamba_track_interval}) must be >= mamba_chunk_size ({self.mamba_chunk_size})" + ), ( + f"mamba_track_interval ({model_runner.server_args.mamba_track_interval}) must be >= mamba_chunk_size ({self.mamba_chunk_size})" + ) def init_forward_metadata_out_graph( self, @@ -910,9 +912,9 @@ class Mamba2AttnBackend(MambaAttnBackendBase): draft_token_num=draft_token_num, ) # `forward` slices the track destinations from ([-num_decodes:]) - assert ( - self.forward_metadata.num_decodes == forward_batch.batch_size - ), f"{self.forward_metadata.num_decodes=} != {forward_batch.batch_size=}" + assert self.forward_metadata.num_decodes == forward_batch.batch_size, ( + f"{self.forward_metadata.num_decodes=} != {forward_batch.batch_size=}" + ) def init_forward_metadata(self, forward_batch: ForwardBatch): metadata = self._forward_metadata(forward_batch) diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py index ac0338511..d9be3424c 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py @@ -399,9 +399,9 @@ class GDNAttnBackend(MambaAttnBackendBase): model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape ) if not is_cpu() and not is_npu(): - assert ( - self.conv_states_shape[-1] < FLA_CHUNK_SIZE - ), f"{self.conv_states_shape[-1]=} should be less than {FLA_CHUNK_SIZE}" + assert self.conv_states_shape[-1] < FLA_CHUNK_SIZE, ( + f"{self.conv_states_shape[-1]=} should be less than {FLA_CHUNK_SIZE}" + ) backends = model_runner.linear_attn_backends self.linear_attn_backends = backends diff --git a/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py b/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py index a3af28686..27b4c948b 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py +++ b/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py @@ -109,7 +109,12 @@ def _get_flashinfer_gdn_kernels(): Returns (available, prefill_fn, mtp_fn, decode_fn, mtp_bf16_fn). """ - global _flashinfer_gdn_available, _flashinfer_chunk_gated_delta_rule, _flashinfer_gated_delta_rule_mtp, _flashinfer_gated_delta_rule_decode, _flashinfer_gated_delta_rule_mtp_bf16 + global \ + _flashinfer_gdn_available, \ + _flashinfer_chunk_gated_delta_rule, \ + _flashinfer_gated_delta_rule_mtp, \ + _flashinfer_gated_delta_rule_decode, \ + _flashinfer_gated_delta_rule_mtp_bf16 if _flashinfer_gdn_available is None: try: os.environ.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1") diff --git a/python/sglang/srt/layers/attention/mamba/mamba.py b/python/sglang/srt/layers/attention/mamba/mamba.py index 4e21ae80a..13c0f35fc 100644 --- a/python/sglang/srt/layers/attention/mamba/mamba.py +++ b/python/sglang/srt/layers/attention/mamba/mamba.py @@ -108,9 +108,9 @@ def mamba_v2_sharded_weight_loader( weight_full_dim_list.append( int(full_dim / full_dim_sum * loaded_weight.size(0)) ) - assert sum(weight_full_dim_list) == loaded_weight.size( - 0 - ), f"Padding the loaded weight failed due to sizes are not divisible cleanly from {weight_full_dim_list} to {loaded_weight.size(0)}" + assert sum(weight_full_dim_list) == loaded_weight.size(0), ( + f"Padding the loaded weight failed due to sizes are not divisible cleanly from {weight_full_dim_list} to {loaded_weight.size(0)}" + ) if loaded_weight.size(0) < full_dim_sum and tp_rank == 0: logger.warning( f"[ZERO-PADDING] Loaded_weight.dim(0) size:{loaded_weight.size(0)} is padding to {full_dim_sum}" @@ -177,7 +177,9 @@ def mamba_v2_sharded_weight_loader( param.data[ boundary : (boundary + take), ... # type: ignore[misc] ] = loaded_weight[ - loaded_start_idx : (loaded_start_idx + take) # type: ignore[misc] + loaded_start_idx : ( + loaded_start_idx + take + ) # type: ignore[misc] ] # type: ignore[misc] # move indexing boundaries @@ -237,9 +239,9 @@ class MambaMixer2(torch.nn.Module): self.num_heads = num_heads = cache_params.shape.num_heads self.head_dim = cache_params.shape.head_dim - assert ( - num_heads % self.tp_size == 0 - ), "Tensor parallel world size must divide num heads." + assert num_heads % self.tp_size == 0, ( + "Tensor parallel world size must divide num heads." + ) assert (n_groups % self.tp_size) == 0 or n_groups == 1, ( "If tensor parallel world size does not divide num_groups, " @@ -637,12 +639,12 @@ class MambaMixer2(torch.nn.Module): # 2. Convolution sequence transformation if is_target_verify: - assert ( - use_triton_causal_conv - ), "Speculative decoding requires use_triton_causal_conv=True for intermediate state support" - assert isinstance( - layer_cache, MambaPool.SpeculativeState - ), "layer_cache must be SpeculativeState for speculative decoding" + assert use_triton_causal_conv, ( + "Speculative decoding requires use_triton_causal_conv=True for intermediate state support" + ) + assert isinstance(layer_cache, MambaPool.SpeculativeState), ( + "layer_cache must be SpeculativeState for speculative decoding" + ) draft_token_num = metadata.draft_token_num self.intermediate_state_indices = torch.arange( num_decodes, dtype=torch.int32, device=state_indices_tensor_d.device diff --git a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py b/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py index 5ebd6f367..60a9da352 100644 --- a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py +++ b/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py @@ -166,7 +166,6 @@ class Mamba2Metadata(ForwardMetadata): p = 0 # num of insertions for s, e in zip(cu_seqlens[:-1], cu_seqlens[1:]): - # if does not divide chunk_size, then there is one chunk insertion p += s % chunk_size > 0 diff --git a/python/sglang/srt/layers/attention/mamba/mixer2_rms_norm_gated.py b/python/sglang/srt/layers/attention/mamba/mixer2_rms_norm_gated.py index 045e4fec3..c1b1fad4e 100644 --- a/python/sglang/srt/layers/attention/mamba/mixer2_rms_norm_gated.py +++ b/python/sglang/srt/layers/attention/mamba/mixer2_rms_norm_gated.py @@ -47,9 +47,9 @@ class Mixer2RMSNormGated(BaseFusedOp): else: # Avoid checkpoint mismatch by skipping unused parameter self.register_parameter("weight", None) - assert ( - self.full_hidden_size % self.tp_size == 0 - ), "Tensor parallel world size must divide hidden size." + assert self.full_hidden_size % self.tp_size == 0, ( + "Tensor parallel world size must divide hidden size." + ) def forward_native( self, diff --git a/python/sglang/srt/layers/attention/minicpm/backend.py b/python/sglang/srt/layers/attention/minicpm/backend.py index 126c9e859..4526f8e17 100644 --- a/python/sglang/srt/layers/attention/minicpm/backend.py +++ b/python/sglang/srt/layers/attention/minicpm/backend.py @@ -480,9 +480,12 @@ class MiniCPMSparseBackend(AttentionBackend): _gather_compressed_keys(full_compressed_k, level, sparse_bs) ) - (compressed_k, compressed_cu_seqlens), ( - compressed_k2, - compressed_cu_seqlens2, + ( + (compressed_k, compressed_cu_seqlens), + ( + compressed_k2, + compressed_cu_seqlens2, + ), ) = compressed ret = self.sparse_get_topk_impl( diff --git a/python/sglang/srt/layers/attention/minimax_sparse_backend.py b/python/sglang/srt/layers/attention/minimax_sparse_backend.py index a49cd02c7..2a84669e6 100644 --- a/python/sglang/srt/layers/attention/minimax_sparse_backend.py +++ b/python/sglang/srt/layers/attention/minimax_sparse_backend.py @@ -226,7 +226,6 @@ class MiniMaxSparseAttnBackend(AttentionBackend): self._msa_dec_meta = None if self.use_msa: - self.num_q_heads = ( runner.model_config.num_attention_heads // get_parallel().attn_tp_size ) @@ -996,9 +995,7 @@ class MiniMaxSparseAttnBackend(AttentionBackend): per_query_prefix = prefix_lens_l.repeat_interleave(extend_lens) # [total_q] per_query_within = torch.arange( total_q, device=device, dtype=torch.long - ) - cu_q[:-1].repeat_interleave( - extend_lens - ) # 0-indexed within each request + ) - cu_q[:-1].repeat_interleave(extend_lens) # 0-indexed within each request per_query_seq_lens = (per_query_prefix + per_query_within + 1).to(torch.int32) max_seqlen = ( diff --git a/python/sglang/srt/layers/attention/minimax_sparse_ops/msa.py b/python/sglang/srt/layers/attention/minimax_sparse_ops/msa.py index 164710cb2..c083bc13c 100644 --- a/python/sglang/srt/layers/attention/minimax_sparse_ops/msa.py +++ b/python/sglang/srt/layers/attention/minimax_sparse_ops/msa.py @@ -31,15 +31,15 @@ def _check_msa_dtypes(q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Ten # on q.dtype alone and casts the k/v pointers to the same element type, so # a mismatched cache would be silently reinterpreted. if q.dtype == torch.bfloat16: - assert ( - k_cache.dtype == torch.bfloat16 - ), f"MSA bf16 requires a bf16 K cache, got {k_cache.dtype}" + assert k_cache.dtype == torch.bfloat16, ( + f"MSA bf16 requires a bf16 K cache, got {k_cache.dtype}" + ) elif q.dtype == torch.float8_e4m3fn: # e5m2 is rejected here too: fmha_sm100's variant lookup falls back to # the e4m3 kernel for unknown dtype codes. - assert ( - k_cache.dtype == torch.float8_e4m3fn - ), f"MSA fp8 requires an fp8_e4m3fn K cache, got {k_cache.dtype}" + assert k_cache.dtype == torch.float8_e4m3fn, ( + f"MSA fp8 requires an fp8_e4m3fn K cache, got {k_cache.dtype}" + ) else: raise AssertionError(f"MSA supports bf16 or fp8_e4m3fn Q, got {q.dtype}") assert v_cache.dtype == k_cache.dtype diff --git a/python/sglang/srt/layers/attention/minimax_sparse_ops/naive/flash_with_topk_idx.py b/python/sglang/srt/layers/attention/minimax_sparse_ops/naive/flash_with_topk_idx.py index 379ffa31e..d473cef75 100644 --- a/python/sglang/srt/layers/attention/minimax_sparse_ops/naive/flash_with_topk_idx.py +++ b/python/sglang/srt/layers/attention/minimax_sparse_ops/naive/flash_with_topk_idx.py @@ -19,9 +19,9 @@ def naive_flash_decode_with_topk_idx( init_blocks: int = 0, local_blocks: int = 0, ): - assert ( - kv_cache.shape[2] % block_size == 0 - ), "max cache len must be divisible by block size" + assert kv_cache.shape[2] % block_size == 0, ( + "max cache len must be divisible by block size" + ) if sm_scale is None: sm_scale = q.shape[-1] ** -0.5 original_dtype = q.dtype diff --git a/python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_flash_with_topk_idx.py b/python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_flash_with_topk_idx.py index 837aa221d..8f3b88bb4 100644 --- a/python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_flash_with_topk_idx.py +++ b/python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_flash_with_topk_idx.py @@ -266,9 +266,9 @@ def test_flash_decode_with_topk_idx( actual_k = min(tk, num_blocks) set_new = set(topk_new[h, b, :actual_k].tolist()) set_ref = set(topk_ref[h, b, :actual_k].tolist()) - assert ( - set_new == set_ref - ), f"topk mismatch at h={h} b={b}: kernel={set_new} ref={set_ref}" + assert set_new == set_ref, ( + f"topk mismatch at h={h} b={b}: kernel={set_new} ref={set_ref}" + ) # --- topk sentinel: invalid positions must be -1 --- for b in range(bs): @@ -277,9 +277,9 @@ def test_flash_decode_with_topk_idx( actual_k = min(tk, num_blocks) if actual_k < tk: invalid = topk_new[:, b, actual_k:] - assert ( - invalid == -1 - ).all(), f"sentinel fail at b={b}: expected -1, got {invalid[invalid != -1].tolist()}" + assert (invalid == -1).all(), ( + f"sentinel fail at b={b}: expected -1, got {invalid[invalid != -1].tolist()}" + ) @pytest.mark.parametrize("score_type", ["max", "lse"]) @@ -343,9 +343,9 @@ def test_flash_decode_score_only( actual_k = min(tk, num_blocks) set_new = set(topk_new[h, b, :actual_k].tolist()) set_ref = set(topk_ref[h, b, :actual_k].tolist()) - assert ( - set_new == set_ref - ), f"topk mismatch at h={h} b={b}: kernel={set_new} ref={set_ref}" + assert set_new == set_ref, ( + f"topk mismatch at h={h} b={b}: kernel={set_new} ref={set_ref}" + ) for b in range(bs): sl = seq_lens[b] @@ -353,9 +353,9 @@ def test_flash_decode_score_only( actual_k = min(tk, num_blocks) if actual_k < tk: invalid = topk_new[:, b, actual_k:] - assert ( - invalid == -1 - ).all(), f"sentinel fail at b={b}: expected -1, got {invalid[invalid != -1].tolist()}" + assert (invalid == -1).all(), ( + f"sentinel fail at b={b}: expected -1, got {invalid[invalid != -1].tolist()}" + ) def test_flash_decode_jit_topk_trivial_rows_skip_score_writes(): diff --git a/python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_msa_fp8_parity.py b/python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_msa_fp8_parity.py index fd08af80b..593960663 100644 --- a/python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_msa_fp8_parity.py +++ b/python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_msa_fp8_parity.py @@ -131,7 +131,7 @@ def test_msa_fp8_decode_vs_triton_fp8(): ) err = (o_msa.float() - o_bf16.float()).abs().mean() ref = o_bf16.float().abs().mean() - assert err / ref < 0.06, f"mean rel err {err/ref:.4f} too high vs bf16 MSA" + assert err / ref < 0.06, f"mean rel err {err / ref:.4f} too high vs bf16 MSA" def test_msa_fp8_decode_scales(): diff --git a/python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_sparse_gqa.py b/python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_sparse_gqa.py index 92d6dd21a..39adece33 100644 --- a/python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_sparse_gqa.py +++ b/python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_sparse_gqa.py @@ -245,9 +245,9 @@ def test_sparse_gqa_vs_reference(bs, nqh, nkh, hd, blk, tk, with_sink, seq_pat, ) o_ref = o_ref.to(o_kernel.dtype) - assert torch.allclose( - o_kernel.float(), o_ref.float(), rtol=RTOL, atol=ATOL - ), f"max abs diff {(o_kernel.float() - o_ref.float()).abs().max().item():.4e}" + assert torch.allclose(o_kernel.float(), o_ref.float(), rtol=RTOL, atol=ATOL), ( + f"max abs diff {(o_kernel.float() - o_ref.float()).abs().max().item():.4e}" + ) @pytest.mark.parametrize( @@ -300,9 +300,9 @@ def test_sparse_gqa_topk_exceeds_blocks( ) o_ref = o_ref.to(o_kernel.dtype) - assert torch.allclose( - o_kernel.float(), o_ref.float(), rtol=RTOL, atol=ATOL - ), f"max abs diff {(o_kernel.float() - o_ref.float()).abs().max().item():.4e}" + assert torch.allclose(o_kernel.float(), o_ref.float(), rtol=RTOL, atol=ATOL), ( + f"max abs diff {(o_kernel.float() - o_ref.float()).abs().max().item():.4e}" + ) @pytest.mark.parametrize( @@ -353,9 +353,9 @@ def test_sparse_gqa_deterministic(bs, nqh, nkh, hd, blk, tk, with_sink, seq_pat, topk_idx, ) - assert torch.equal( - o1, o2 - ), f"non-deterministic: max diff {(o1.float() - o2.float()).abs().max().item():.4e}" + assert torch.equal(o1, o2), ( + f"non-deterministic: max diff {(o1.float() - o2.float()).abs().max().item():.4e}" + ) if __name__ == "__main__": diff --git a/python/sglang/srt/layers/attention/tbo_backend.py b/python/sglang/srt/layers/attention/tbo_backend.py index 77615cf50..94001ba56 100644 --- a/python/sglang/srt/layers/attention/tbo_backend.py +++ b/python/sglang/srt/layers/attention/tbo_backend.py @@ -220,9 +220,9 @@ def _build_tbo_child_replay_fb_view( capture-time buffers are sliced per child, spec_info is split, and seq_lens_sum is recomputed from the sliced ``seq_lens_cpu``. """ - assert ( - getattr(fb_view, "encoder_lens", None) is None - ), "TBO replay split does not support encoder_lens yet" + assert getattr(fb_view, "encoder_lens", None) is None, ( + "TBO replay split does not support encoder_lens yet" + ) spec_info = getattr(fb_view, "spec_info", None) if spec_info is not None: start_seq = seq_slice.start or 0 diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index 0d736d7e7..0e04dd0eb 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -382,9 +382,9 @@ class TritonAttnBackend(AttentionBackend): # And the real_num_token is num_seq in decoding phase. num_group = num_token // num_seq - assert ( - num_group * num_seq == num_token - ), f"num_seq({num_seq}), num_token({num_token}), something goes wrong!" + assert num_group * num_seq == num_token, ( + f"num_seq({num_seq}), num_token({num_token}), something goes wrong!" + ) if ( self.static_kv_splits or self.device_core_count <= 0 diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index e5f8359ae..357239c61 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -242,7 +242,6 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): # "Missing TRTLLM-GEN kernel" error during CUDA-graph capture. # XQA (SM90/SM120 decode) has native page-128 kernels; no check needed. if self.page_size >= 128 and not self.is_xqa_impl: - attn_tp_size = get_parallel().attn_tp_size num_q_heads = config.num_attention_heads // attn_tp_size num_kv_heads = config.get_num_kv_heads(attn_tp_size) @@ -1214,7 +1213,9 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): ) return output_by_request.view(-1, query.shape[-2], query.shape[-1]) - def _get_nvfp4_decode_kv_cache(self, layer: RadixAttention) -> tuple[ + def _get_nvfp4_decode_kv_cache( + self, layer: RadixAttention + ) -> tuple[ tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor], ]: @@ -1426,8 +1427,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): # run bs*L single-token rows over the full window instead (the # window's K/V are already in the pool). assert not self.forward_metadata.is_ragged_verify, ( - "ENCODER_ONLY target_verify does not support ragged " - "verify layouts" + "ENCODER_ONLY target_verify does not support ragged verify layouts" ) assert self.forward_metadata.encoder_cache_seqlens is not None, ( "ENCODER_ONLY target_verify requires the expanded decode " diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index 8e8fc49ed..efb7bfe85 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -1290,9 +1290,9 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): # Save KV cache if requested (the fused fp8 path already wrote it) query = fused_fp8_query if query is None and save_kv_cache: - assert ( - k is not None and k_rope is not None - ), "For populating trtllm_mla kv cache, both k_nope and k_rope should be not None." + assert k is not None and k_rope is not None, ( + "For populating trtllm_mla kv cache, both k_nope and k_rope should be not None." + ) if self._decode_kernel_loc is not None: if merge_query and self._fused_set_kv_concat_q: # Fused: KV scatter + [q_nope | q_rope] concat in one @@ -1494,9 +1494,9 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): # Save KV cache if requested if save_kv_cache: - assert ( - k is not None and k_rope is not None - ), "For populating trtllm_mla kv cache, both k_nope and k_rope should be not None." + assert k is not None and k_rope is not None, ( + "For populating trtllm_mla kv cache, both k_nope and k_rope should be not None." + ) if self._decode_kernel_loc is not None: self.token_to_kv_pool.set_mla_kv_buffer( layer, self._decode_kernel_loc, k, k_rope diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index fbc0f4efd..bff80a1f1 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/python/sglang/srt/layers/attention/vision.py @@ -137,7 +137,6 @@ class SingletonCache: @dataclasses.dataclass class VisionAttentionMetadata: - cu_seqlens: torch.Tensor seq_lens: torch.Tensor max_seqlen: int @@ -242,9 +241,9 @@ def resolve_seqlens( resolved_seqlens = cu_seqlens.get_data() else: resolved_seqlens = cu_seqlens - assert isinstance( - resolved_seqlens, torch.Tensor - ), "cu_seqlens must be a torch.Tensor" + assert isinstance(resolved_seqlens, torch.Tensor), ( + "cu_seqlens must be a torch.Tensor" + ) return resolved_seqlens @@ -820,7 +819,6 @@ class VisionAiterAttention(nn.Module): class VisionAscendAttention(nn.Module): - def __init__( self, **kwargs, @@ -1505,7 +1503,6 @@ class VisionAttention(nn.Module): if self.qk_normalization and not self.qk_normalization_by_head_size: # jit kernel if can_use_jit_qk_norm(self.head_size, q.dtype): - # q: [tokens, head, head_size] -> [tokens, embed_dim] head_dim_for_norm = head * self.head_size diff --git a/python/sglang/srt/layers/attention/wave_backend.py b/python/sglang/srt/layers/attention/wave_backend.py index 68a8b7f62..91c742f4e 100644 --- a/python/sglang/srt/layers/attention/wave_backend.py +++ b/python/sglang/srt/layers/attention/wave_backend.py @@ -42,7 +42,6 @@ class ForwardMetadata: class WaveAttnBackend(AttentionBackend): - # kv_indptr/qo_indptr are preallocated at (req pool + 1); an extend batch # can never carry more seqs than the pool. extend_dummy_seqs_capped_by_req_pool: bool = True @@ -131,9 +130,9 @@ class WaveAttnBackend(AttentionBackend): num_token, num_seq = num_kv_splits.shape[0], seq_lens.shape[0] num_group = num_token // num_seq - assert ( - num_group * num_seq == num_token - ), f"num_seq({num_seq}), num_token({num_token}), something goes wrong!" + assert num_group * num_seq == num_token, ( + f"num_seq({num_seq}), num_token({num_token}), something goes wrong!" + ) if self.static_kv_splits or self.device_core_count <= 0: num_kv_splits.fill_(self.max_kv_splits) diff --git a/python/sglang/srt/layers/attention/xpu_backend.py b/python/sglang/srt/layers/attention/xpu_backend.py index 136580d14..d56459962 100644 --- a/python/sglang/srt/layers/attention/xpu_backend.py +++ b/python/sglang/srt/layers/attention/xpu_backend.py @@ -121,9 +121,9 @@ class XPUAttentionBackend(AttentionBackend): if forward_batch.forward_mode.is_decode_or_idle(): # Draft Decode if forward_batch.spec_info is not None: - assert ( - False - ), "XPUAttentionBackend doesn't support speculative decoding yet, please use --attention-backend triton instead." + assert False, ( + "XPUAttentionBackend doesn't support speculative decoding yet, please use --attention-backend triton instead." + ) if self.topk <= 1: metadata.cache_seqlens_int32 = ( seqlens_in_batch + (self.speculative_step_id + 1) @@ -273,9 +273,7 @@ class XPUAttentionBackend(AttentionBackend): # create expand page table offsets = torch.arange( self.speculative_num_draft_tokens, device=device - ).unsqueeze( - 0 - ) # shape: (1, self.speculative_num_draft_tokens) + ).unsqueeze(0) # shape: (1, self.speculative_num_draft_tokens) cols = offsets.expand( forward_batch.seq_lens.numel(), -1 ) + forward_batch.seq_lens.unsqueeze(1) @@ -363,9 +361,9 @@ class XPUAttentionBackend(AttentionBackend): # Encoder metadata for cross attention if forward_batch.encoder_lens is not None: - assert ( - forward_batch.encoder_lens.numel() == 1 - ), "Only encoder size 1 is supported for now" + assert forward_batch.encoder_lens.numel() == 1, ( + "Only encoder size 1 is supported for now" + ) metadata.encoder_lens_int32 = forward_batch.encoder_lens.to(torch.int32) metadata.encoder_cu_seqlens_k = torch.nn.functional.pad( @@ -1084,12 +1082,12 @@ class XPUAttentionBackend(AttentionBackend): forward_mode = forward_batch.forward_mode spec_info = forward_batch.spec_info - assert ( - spec_info is None - ), "XPUAttentionBackend does not support speculative decoding in XPU graph" - assert ( - forward_mode.is_decode_or_idle() - ), "XPUAttentionBackend XPU graph only supports decode mode" + assert spec_info is None, ( + "XPUAttentionBackend does not support speculative decoding in XPU graph" + ) + assert forward_mode.is_decode_or_idle(), ( + "XPUAttentionBackend XPU graph only supports decode mode" + ) if in_capture: # Bind static-shape slices of the pre-allocated buffers so the @@ -1290,9 +1288,9 @@ class XPUAttentionBackend(AttentionBackend): metadata_swa: Optional[FlashAttentionMetadata] = None, ): # TODO: support page_size > 1 for swa spec - assert ( - self.page_size == 1 - ), "FlashAttention backend doesn't support topk > 1 speculative decoding with page size > 1 sliding window attention" + assert self.page_size == 1, ( + "FlashAttention backend doesn't support topk > 1 speculative decoding with page size > 1 sliding window attention" + ) cache_seqlens_int32 = ( metadata.cache_seqlens_int32.repeat_interleave( diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index c5c35016a..c395f96bf 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -813,9 +813,9 @@ class LayerCommunicator: ) -> Tuple[torch.Tensor, torch.Tensor]: if hidden_states.shape[0] == 0: return hidden_states, hidden_states - assert ( - hidden_states.shape[0] % self._context.tp_size == 0 - ), f"Expected total tokens {hidden_states.shape[0]} % tp_size {self._context.tp_size} to be 0" + assert hidden_states.shape[0] % self._context.tp_size == 0, ( + f"Expected total tokens {hidden_states.shape[0]} % tp_size {self._context.tp_size} to be 0" + ) local_tokens = hidden_states.shape[0] // self._context.tp_size output = hidden_states.new_empty(local_tokens, *hidden_states.shape[1:]) get_tp_group().reduce_scatter_tensor(output, hidden_states) diff --git a/python/sglang/srt/layers/communicator_dsa_cp.py b/python/sglang/srt/layers/communicator_dsa_cp.py index c55d64b3b..ee661b90f 100644 --- a/python/sglang/srt/layers/communicator_dsa_cp.py +++ b/python/sglang/srt/layers/communicator_dsa_cp.py @@ -115,9 +115,9 @@ class DSACPLayerCommunicator(LayerCommunicator): def _post_init_communicate(self): # SCATTERED in attn tp is different from SCATTERED in global tp when dp_size > 1 if self.layer_scatter_modes.mlp_mode != ScatterMode.SCATTERED: - assert ( - self._context.attn_dp_size == 1 - ), f"dp_size should be 1 when moe_runner_backend is none" + assert self._context.attn_dp_size == 1, ( + f"dp_size should be 1 when moe_runner_backend is none" + ) self._communicate_simple_fn = DSACPCommunicateSimpleFn.get_fn( input_mode=ScatterMode.SCATTERED, output_mode=ScatterMode.SCATTERED, diff --git a/python/sglang/srt/layers/cp/base.py b/python/sglang/srt/layers/cp/base.py index 308d06add..6328151fa 100644 --- a/python/sglang/srt/layers/cp/base.py +++ b/python/sglang/srt/layers/cp/base.py @@ -49,8 +49,7 @@ class ContextParallelStrategyKind(IntEnum): if value == "interleave": return cls.INTERLEAVE raise ValueError( - f"Unknown cp_strategy={value!r}; expected one of " - "{'zigzag', 'interleave'}" + f"Unknown cp_strategy={value!r}; expected one of {{'zigzag', 'interleave'}}" ) @property diff --git a/python/sglang/srt/layers/cp/zigzag.py b/python/sglang/srt/layers/cp/zigzag.py index ce2da4e97..868c527c6 100644 --- a/python/sglang/srt/layers/cp/zigzag.py +++ b/python/sglang/srt/layers/cp/zigzag.py @@ -348,9 +348,9 @@ class ZigzagCPStrategy(ContextParallelStrategy): attn_fn, attention_backend: CPAttentionBackendKind = CPAttentionBackendKind.FLASH_ATTENTION, ) -> Any: - assert ( - attention_backend in self.get_supported_attention_backend() - ), f"{self.name} CP does not support {attention_backend=}" + assert attention_backend in self.get_supported_attention_backend(), ( + f"{self.name} CP does not support {attention_backend=}" + ) meta = forward_batch.attn_cp_metadata q_prev = q[: meta.total_q_prev_tokens] diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py index fce7d735a..00f756c39 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py @@ -72,7 +72,6 @@ def grouped_gemm_nt_f8f8bf16_masked( with configure_deep_gemm_num_sms( overlap_args.num_sms if overlap_args is not None else None ): - fp4_kwargs = {} if recipe_a is not None: fp4_kwargs["recipe_a"] = recipe_a diff --git a/python/sglang/srt/layers/dp_attention.py b/python/sglang/srt/layers/dp_attention.py index 87117547a..86e2131f2 100644 --- a/python/sglang/srt/layers/dp_attention.py +++ b/python/sglang/srt/layers/dp_attention.py @@ -83,7 +83,6 @@ _is_cpu = is_cpu() class DpPaddingMode(IntEnum): - # Padding tokens to max length and then gather tokens using `all_gather_into_tensor` MAX_LEN = auto() # Padding tokens to sum length and then gather tokens using `all_reduce` @@ -512,9 +511,9 @@ def _dp_gather_via_all_reduce( if local_tokens.shape[0] > 0 and ( is_partial or get_attn_tensor_model_parallel_rank() == 0 ): - assert ( - local_tokens.untyped_storage() is not global_tokens.untyped_storage() - ), "aliasing between global_tokens and local_tokens not allowed" + assert local_tokens.untyped_storage() is not global_tokens.untyped_storage(), ( + "aliasing between global_tokens and local_tokens not allowed" + ) memcpy(global_tokens, local_tokens, 0, local_start_pos, local_num_tokens, False) @@ -865,9 +864,9 @@ def dp_scatter( assert local_tokens.is_contiguous() assert global_tokens.is_contiguous() if local_tokens.shape[0] > 0: - assert ( - local_tokens.untyped_storage() is not global_tokens.untyped_storage() - ), "aliasing between local_tokens and global_tokens not allowed" + assert local_tokens.untyped_storage() is not global_tokens.untyped_storage(), ( + "aliasing between local_tokens and global_tokens not allowed" + ) memcpy(local_tokens, global_tokens, 0, local_start_pos, local_num_tokens, True) diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py index d84019fbd..2fb64f1ab 100644 --- a/python/sglang/srt/layers/linear.py +++ b/python/sglang/srt/layers/linear.py @@ -285,13 +285,13 @@ class ReplicatedLinear(LinearBase): raise ValueError(f"{loaded_weight} are not all equal") if param.dtype == torch.int8 or loaded_weight.dtype == torch.int8: - assert ( - param.dtype == loaded_weight.dtype - ), "init para dtype and loaded weight dtype should be the same" + assert param.dtype == loaded_weight.dtype, ( + "init para dtype and loaded weight dtype should be the same" + ) - assert ( - param.size() == loaded_weight.size() - ), f"{param.shape=} {param.dtype=} {loaded_weight.shape=} {loaded_weight.dtype=}" + assert param.size() == loaded_weight.size(), ( + f"{param.shape=} {param.dtype=} {loaded_weight.shape=} {loaded_weight.dtype=}" + ) param.data.copy_(loaded_weight) def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: @@ -458,9 +458,9 @@ class ColumnParallelLinear(LinearBase): if len(loaded_weight.shape) == 0: loaded_weight = loaded_weight.reshape(1) - assert ( - param_data.shape == loaded_weight.shape - ), f"param_data.shape={param_data.shape} != loaded_weight.shape={loaded_weight.shape}" + assert param_data.shape == loaded_weight.shape, ( + f"param_data.shape={param_data.shape} != loaded_weight.shape={loaded_weight.shape}" + ) param_data.copy_(loaded_weight) def weight_loader_v2(self, param: Parameter, loaded_weight: torch.Tensor): @@ -1415,9 +1415,9 @@ class QKVParallelLinear(ColumnParallelLinear): "for all partitions." ) - assert ( - param_data.shape == loaded_weight.shape - ), f"{param_data.shape=} {loaded_weight.shape=}" + assert param_data.shape == loaded_weight.shape, ( + f"{param_data.shape=} {loaded_weight.shape=}" + ) param_data.copy_(loaded_weight) @@ -1575,9 +1575,9 @@ class RowParallelLinear(LinearBase): if len(loaded_weight.shape) == 0: loaded_weight = loaded_weight.reshape(1) - assert ( - param_data.shape == loaded_weight.shape - ), f"{param_data.shape=} {loaded_weight.shape=}" + assert param_data.shape == loaded_weight.shape, ( + f"{param_data.shape=} {loaded_weight.shape=}" + ) param_data.copy_(loaded_weight) def weight_loader_v2(self, param: BasevLLMParameter, loaded_weight: torch.Tensor): diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index ef0ad8b5a..77feb0dba 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -88,7 +88,7 @@ def _trace_e2e_logits(stage: str, **fields) -> None: return try: parallel = get_parallel() - rank = f"dp={parallel.attn_dp_rank} " f"tp={parallel.tp_rank}" + rank = f"dp={parallel.attn_dp_rank} tp={parallel.tp_rank}" except Exception: rank = "rank=unknown" details = " ".join(f"{key}={value}" for key, value in fields.items()) diff --git a/python/sglang/srt/layers/logprob_processor.py b/python/sglang/srt/layers/logprob_processor.py index 1536030ae..032da6ace 100644 --- a/python/sglang/srt/layers/logprob_processor.py +++ b/python/sglang/srt/layers/logprob_processor.py @@ -357,9 +357,9 @@ def compute_spec_logprobs( accept_index: Optional[torch.Tensor] = None, chain_stride: Optional[int] = None, ): - assert (accept_index is None) != ( - chain_stride is None - ), "pass exactly one of accept_index / chain_stride" + assert (accept_index is None) != (chain_stride is None), ( + "pass exactly one of accept_index / chain_stride" + ) bs = len(batch.seq_lens) next_token_logits = logits_output.next_token_logits @@ -667,9 +667,9 @@ class InputLogprobProcessor: # Restore the full-pruned lm_head batch_info after chunk iteration. if num_chunks > 1 and hasattr(lm_head, "reset_lm_head_pass"): - assert hasattr( - lm_head, "set_lm_head_pass" - ), "lm_head must have set_lm_head_pass method and reset_lm_head_pass method at the same time" + assert hasattr(lm_head, "set_lm_head_pass"), ( + "lm_head must have set_lm_head_pass method and reset_lm_head_pass method at the same time" + ) lm_head.reset_lm_head_pass() # Concatenate the results diff --git a/python/sglang/srt/layers/moe/cutlass_moe.py b/python/sglang/srt/layers/moe/cutlass_moe.py index ebc738c70..11bcad7dd 100755 --- a/python/sglang/srt/layers/moe/cutlass_moe.py +++ b/python/sglang/srt/layers/moe/cutlass_moe.py @@ -161,12 +161,12 @@ def cutlass_fused_experts_fp8( w2_q.shape[1] // 32, w2_q.shape[2], ) - assert ( - w1_scale.shape == expected_w1_scale_shape - ), f"MXFP8 w1_scale must be {expected_w1_scale_shape}, got {w1_scale.shape}" - assert ( - w2_scale.shape == expected_w2_scale_shape - ), f"MXFP8 w2_scale must be {expected_w2_scale_shape}, got {w2_scale.shape}" + assert w1_scale.shape == expected_w1_scale_shape, ( + f"MXFP8 w1_scale must be {expected_w1_scale_shape}, got {w1_scale.shape}" + ) + assert w2_scale.shape == expected_w2_scale_shape, ( + f"MXFP8 w2_scale must be {expected_w2_scale_shape}, got {w2_scale.shape}" + ) mxfp8_blockscale_align = 128 total_tokens = m * topk diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index ba81c9f1d..d445bfc71 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/python/sglang/srt/layers/moe/ep_moe/layer.py @@ -143,9 +143,9 @@ class DeepEPMoE(FusedMoE): and not _is_npu and not _is_hip ): - assert ( - deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM - ), "Unquantized DeepEP MoE requires DeepGEMM BF16" + assert deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM, ( + "Unquantized DeepEP MoE requires DeepGEMM BF16" + ) self.deprecate_flag = True else: self.deprecate_flag = False @@ -175,9 +175,9 @@ class DeepEPMoE(FusedMoE): and quant_config is not None ): # AMD HIP and NPU support low_latency DeepEP without DeepGEMM. - assert ( - deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM - ), f"DeepEP {self.deepep_mode} mode requires deep_gemm" + assert deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM, ( + f"DeepEP {self.deepep_mode} mode requires deep_gemm" + ) def _a2a_forward_with_output_impl( self, @@ -223,9 +223,9 @@ class DeepEPMoE(FusedMoE): ): # DeepEP NORMAL mode is not capturable; run it as an eager node. if is_in_breakable_cuda_graph(): - assert TopKOutputChecker.format_is_standard( - topk_output - ), "Only standard topk output is supported for breakable cuda graph" + assert TopKOutputChecker.format_is_standard(topk_output), ( + "Only standard topk output is supported for breakable cuda graph" + ) output = torch.empty_like(hidden_states) self.a2a_forward_with_output( hidden_states, @@ -236,9 +236,9 @@ class DeepEPMoE(FusedMoE): ) return output if is_in_tc_piecewise_cuda_graph(): - assert TopKOutputChecker.format_is_standard( - topk_output - ), "Only standard topk output is supported for piecewise cuda graph" + assert TopKOutputChecker.format_is_standard(topk_output), ( + "Only standard topk output is supported for piecewise cuda graph" + ) return moe_forward_piecewise_cuda_graph_impl( hidden_states, topk_output.topk_weights, diff --git a/python/sglang/srt/layers/moe/flashinfer_cutedsl_moe.py b/python/sglang/srt/layers/moe/flashinfer_cutedsl_moe.py index b2887b3c8..6eea348bc 100644 --- a/python/sglang/srt/layers/moe/flashinfer_cutedsl_moe.py +++ b/python/sglang/srt/layers/moe/flashinfer_cutedsl_moe.py @@ -59,31 +59,30 @@ def flashinfer_cutedsl_moe_masked( # === Assertions on dtypes === assert w1.dtype == torch.uint8, f"w1 must be uint8 (fp4 packed), got {w1.dtype}" - assert ( - w1_blockscale.dtype == torch.float8_e4m3fn - ), f"w1_blockscale must be float8_e4m3fn, got {w1_blockscale.dtype}" - assert ( - w1_alpha.dtype == torch.float32 - ), f"w1_alpha must be float32, got {w1_alpha.dtype}" + assert w1_blockscale.dtype == torch.float8_e4m3fn, ( + f"w1_blockscale must be float8_e4m3fn, got {w1_blockscale.dtype}" + ) + assert w1_alpha.dtype == torch.float32, ( + f"w1_alpha must be float32, got {w1_alpha.dtype}" + ) assert w2.dtype == torch.uint8, f"w2 must be uint8 (fp4 packed), got {w2.dtype}" - assert ( - a2_global_scale.dtype == torch.float32 - ), f"a2_global_scale must be float32, got {a2_global_scale.dtype}" - assert ( - w2_blockscale.dtype == torch.float8_e4m3fn - ), f"w2_blockscale must be float8_e4m3fn, got {w2_blockscale.dtype}" - assert ( - w2_alpha.dtype == torch.float32 - ), f"w2_alpha must be float32, got {w2_alpha.dtype}" - assert ( - len(hidden_states) == 2 - ), f"hidden_states must be a tuple of length 2, got {len(hidden_states)}" + assert a2_global_scale.dtype == torch.float32, ( + f"a2_global_scale must be float32, got {a2_global_scale.dtype}" + ) + assert w2_blockscale.dtype == torch.float8_e4m3fn, ( + f"w2_blockscale must be float8_e4m3fn, got {w2_blockscale.dtype}" + ) + assert w2_alpha.dtype == torch.float32, ( + f"w2_alpha must be float32, got {w2_alpha.dtype}" + ) + assert len(hidden_states) == 2, ( + f"hidden_states must be a tuple of length 2, got {len(hidden_states)}" + ) # === Assertions on shapes === n = w2.shape[-1] * 2 # intermediate dimension if hidden_states[1] is not None: - a_q = hidden_states[0].view(torch.uint8) a_q_sf = hidden_states[1].view(torch.float8_e4m3fn) m, k_by_2, num_experts = a_q.shape @@ -91,12 +90,12 @@ def flashinfer_cutedsl_moe_masked( else: num_experts, m, k = hidden_states[0].shape - assert ( - input_global_scale.dtype == torch.float32 - ), f"input_global_scale must be float32, got {input_global_scale.dtype}" - assert input_global_scale.shape == ( - num_experts, - ), f"input_global_scale must be (l,), got {input_global_scale.shape}" + assert input_global_scale.dtype == torch.float32, ( + f"input_global_scale must be float32, got {input_global_scale.dtype}" + ) + assert input_global_scale.shape == (num_experts,), ( + f"input_global_scale must be (l,), got {input_global_scale.shape}" + ) a_q, a_q_sf = scaled_fp4_grouped_quantize( hidden_states[0], @@ -116,25 +115,25 @@ def flashinfer_cutedsl_moe_masked( # Gated (silu_and_mul) GEMM1 emits [gate, up] so w1 has 2*n rows; non-gated # relu2 emits a single projection of n rows. gemm1_out_dim = 2 * n if gated else n - assert ( - w1.shape[-2] == gemm1_out_dim - ), f"w1 last-2 dim must be {gemm1_out_dim} (gated={gated}), got {w1.shape}" - assert ( - w1.shape[-1] * 2 == k - ), f"w1 last dim * 2 must equal k, got {w1.shape[-1]} vs k={k}" + assert w1.shape[-2] == gemm1_out_dim, ( + f"w1 last-2 dim must be {gemm1_out_dim} (gated={gated}), got {w1.shape}" + ) + assert w1.shape[-1] * 2 == k, ( + f"w1 last dim * 2 must equal k, got {w1.shape[-1]} vs k={k}" + ) assert w2.shape[-2:] == ( k, n // 2, - ), f"w2 shape mismatch, got {w2.shape[-2:]}, expected {(k, n//2)}" - assert w1_alpha.shape == ( - num_experts, - ), f"w1_alpha must be (l,), got {w1_alpha.shape}" - assert a2_global_scale.shape == ( - num_experts, - ), f"a2_global_scale must be (l,), got {a2_global_scale.shape}" - assert w2_alpha.shape == ( - num_experts, - ), f"w2_alpha must be (l,), got {w2_alpha.shape}" + ), f"w2 shape mismatch, got {w2.shape[-2:]}, expected {(k, n // 2)}" + assert w1_alpha.shape == (num_experts,), ( + f"w1_alpha must be (l,), got {w1_alpha.shape}" + ) + assert a2_global_scale.shape == (num_experts,), ( + f"a2_global_scale must be (l,), got {a2_global_scale.shape}" + ) + assert w2_alpha.shape == (num_experts,), ( + f"w2_alpha must be (l,), got {w2_alpha.shape}" + ) # TODO(kaixih@nvidia): dtype should be based on inputs. gateup_output = torch.empty( diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py index fd37548ac..7f76b2616 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py @@ -194,9 +194,9 @@ def fused_marlin_moe( assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch" assert hidden_states.shape[1] == w1.shape[1] * 16, "Hidden size mismatch w1" - assert hidden_states.shape[1] == w2.shape[2] // ( - num_bits // 2 - ), "Hidden size mismatch w2" + assert hidden_states.shape[1] == w2.shape[2] // (num_bits // 2), ( + "Hidden size mismatch w2" + ) assert hidden_states.is_contiguous(), "Hidden_states must be contiguous" assert w1.is_contiguous(), "Expert weights1 must be contiguous" assert w2.is_contiguous(), "Expert weights2 must be contiguous" @@ -221,12 +221,12 @@ def fused_marlin_moe( f"activations, got {hidden_states.dtype}" ) elif not is_nvfp4_marlin: - assert ( - hidden_states.dtype == w1_scale.dtype - ), f"moe_wna16_marlin_gemm assumes hidden_states.dtype ({hidden_states.dtype}) == w1_scale.dtype ({w1_scale.dtype})" - assert ( - hidden_states.dtype == w2_scale.dtype - ), f"moe_wna16_marlin_gemm assumes hidden_states.dtype ({hidden_states.dtype}) == w2_scale.dtype ({w2_scale.dtype})" + assert hidden_states.dtype == w1_scale.dtype, ( + f"moe_wna16_marlin_gemm assumes hidden_states.dtype ({hidden_states.dtype}) == w1_scale.dtype ({w1_scale.dtype})" + ) + assert hidden_states.dtype == w2_scale.dtype, ( + f"moe_wna16_marlin_gemm assumes hidden_states.dtype ({hidden_states.dtype}) == w2_scale.dtype ({w2_scale.dtype})" + ) assert num_bits in [4, 8] M, K = hidden_states.shape diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/triton_kernels_moe.py b/python/sglang/srt/layers/moe/fused_moe_triton/triton_kernels_moe.py index d410fde58..397eae1a6 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/triton_kernels_moe.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/triton_kernels_moe.py @@ -153,12 +153,12 @@ def triton_kernel_fused_experts( # Shape check assert hidden_states.ndim == 2, "hidden_states must be 2D" - assert ( - hidden_states.shape[-1] == w1.shape[-2] - ), f"hidden_states shape[-1] {hidden_states.shape} must be equal to w1 shape[-2] {w1.shape}" - assert ( - w2.shape[-1] == w1.shape[1] - ), f"w2 shape[-1] {w2.shape[-1]} must be equal to w1 shape[1] {w1.shape[1]}" + assert hidden_states.shape[-1] == w1.shape[-2], ( + f"hidden_states shape[-1] {hidden_states.shape} must be equal to w1 shape[-2] {w1.shape}" + ) + assert w2.shape[-1] == w1.shape[1], ( + f"w2 shape[-1] {w2.shape[-1]} must be equal to w1 shape[1] {w1.shape[1]}" + ) # feature check assert inplace is False, "Inplace is not supported in new triton MoE kernel" @@ -309,12 +309,12 @@ def triton_kernel_fused_experts_with_bias( # Shape check assert hidden_states.ndim == 2, "hidden_states must be 2D" - assert ( - hidden_states.shape[-1] == w1.shape[-2] - ), f"hidden_states shape[-1] {hidden_states.shape} must be equal to w1 shape[-2] {w1.shape}" - assert ( - w2.shape[-1] == w1.shape[1] - ), f"w2 shape[-1] {w2.shape[-1]} must be equal to w1 shape[1] {w1.shape[1]}" + assert hidden_states.shape[-1] == w1.shape[-2], ( + f"hidden_states shape[-1] {hidden_states.shape} must be equal to w1 shape[-2] {w1.shape}" + ) + assert w2.shape[-1] == w1.shape[1], ( + f"w2 shape[-1] {w2.shape[-1]} must be equal to w1 shape[1] {w1.shape[1]}" + ) # feature check assert inplace is False, "Inplace is not supported in new triton MoE kernel" diff --git a/python/sglang/srt/layers/moe/hash_topk.py b/python/sglang/srt/layers/moe/hash_topk.py index 271344732..72f707211 100644 --- a/python/sglang/srt/layers/moe/hash_topk.py +++ b/python/sglang/srt/layers/moe/hash_topk.py @@ -218,9 +218,9 @@ class HashTopK(nn.Module): num_token_non_padded: Optional[torch.Tensor] = None, expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, ): - assert ( - input_ids.shape[0] == hidden_states.shape[0] == router_logits.shape[0] - ), f"{input_ids.shape=} {hidden_states.shape=} {router_logits.shape=}" + assert input_ids.shape[0] == hidden_states.shape[0] == router_logits.shape[0], ( + f"{input_ids.shape=} {hidden_states.shape=} {router_logits.shape=}" + ) if _is_xpu: topk_weights, topk_ids = self._forward_xpu(router_logits, input_ids) diff --git a/python/sglang/srt/layers/moe/kt_ep_wrapper.py b/python/sglang/srt/layers/moe/kt_ep_wrapper.py index 8eb3d0fc8..36c15cf13 100644 --- a/python/sglang/srt/layers/moe/kt_ep_wrapper.py +++ b/python/sglang/srt/layers/moe/kt_ep_wrapper.py @@ -286,9 +286,9 @@ class KTEPWrapperMethod(FusedMoEMethodBase): layer: The MoE layer module dispatch_output: Dispatched tokens and routing information """ - assert ( - self.moe_runner_config.activation == "silu" - ), "Only SiLU activation is supported." + assert self.moe_runner_config.activation == "silu", ( + "Only SiLU activation is supported." + ) if self.tp_rank != 0 or self.wrapper is None: return diff --git a/python/sglang/srt/layers/moe/mega_moe_sm90.py b/python/sglang/srt/layers/moe/mega_moe_sm90.py index 8929a0a64..2da005dff 100644 --- a/python/sglang/srt/layers/moe/mega_moe_sm90.py +++ b/python/sglang/srt/layers/moe/mega_moe_sm90.py @@ -114,8 +114,7 @@ def build_sm90_mega_moe_experts_weights(experts) -> None: scale_group_mn, scale_group_k = 128, 128 assert k1 % scale_group_k == 0 and k2 % scale_group_k == 0, ( - f"invalid SM90 mega-moe K/group_size: k1={k1}, k2={k2}, " - f"group_k={scale_group_k}" + f"invalid SM90 mega-moe K/group_size: k1={k1}, k2={k2}, group_k={scale_group_k}" ) expected_n_groups_1 = (n1 + scale_group_mn - 1) // scale_group_mn expected_n_groups_2 = (n2 + scale_group_mn - 1) // scale_group_mn diff --git a/python/sglang/srt/layers/moe/moe_runner/aiter.py b/python/sglang/srt/layers/moe/moe_runner/aiter.py index 66311cc4f..6159fbca6 100644 --- a/python/sglang/srt/layers/moe/moe_runner/aiter.py +++ b/python/sglang/srt/layers/moe/moe_runner/aiter.py @@ -336,9 +336,9 @@ def pre_permute_standard_to_aiter( if runner_config.apply_router_weight_on_input and not quant_info.doweight_stage1: # Pre-scale at the Python level for kernels that don't honor doweight_stage1. - assert ( - topk_weights.dim() == 2 and topk_weights.shape[-1] == 1 - ), "apply_router_weight_on_input requires topk=1" + assert topk_weights.dim() == 2 and topk_weights.shape[-1] == 1, ( + "apply_router_weight_on_input requires topk=1" + ) hidden_states = hidden_states * topk_weights.to(hidden_states.dtype) topk_weights = torch.ones_like(topk_weights) diff --git a/python/sglang/srt/layers/moe/moe_runner/base.py b/python/sglang/srt/layers/moe/moe_runner/base.py index e5a2b23d9..f76acf89c 100644 --- a/python/sglang/srt/layers/moe/moe_runner/base.py +++ b/python/sglang/srt/layers/moe/moe_runner/base.py @@ -146,12 +146,12 @@ class FusedOpPool: raise ValueError( f"Fused function for {a2a_backend_name} to {runner_backend_name} is already registered." ) - assert MoeA2ABackend( - a2a_backend_name - ), f"Invalid dispatch name: {a2a_backend_name}" - assert MoeRunnerBackend( - runner_backend_name - ), f"Invalid runner name: {runner_backend_name}" + assert MoeA2ABackend(a2a_backend_name), ( + f"Invalid dispatch name: {a2a_backend_name}" + ) + assert MoeRunnerBackend(runner_backend_name), ( + f"Invalid runner name: {runner_backend_name}" + ) cls._fused_funcs[key] = fused_func @classmethod @@ -228,9 +228,9 @@ class PermuteMethodPool: """ key = (dispatch_output_format, runner_input_format) pre_permute_func = cls._pre_permute_methods.get(key) - assert ( - pre_permute_func is not None - ), f"Pre-permute function for {dispatch_output_format} to {runner_input_format} is not registered" + assert pre_permute_func is not None, ( + f"Pre-permute function for {dispatch_output_format} to {runner_input_format} is not registered" + ) return pre_permute_func @classmethod @@ -248,9 +248,9 @@ class PermuteMethodPool: """ key = (runner_output_format, combine_input_format) post_permute_func = cls._post_permute_methods.get(key) - assert ( - post_permute_func is not None - ), f"Post-permute function for {runner_output_format} to {combine_input_format} is not registered" + assert post_permute_func is not None, ( + f"Post-permute function for {runner_output_format} to {combine_input_format} is not registered" + ) return post_permute_func diff --git a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py index 8eff8224f..bc25cbaf2 100644 --- a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py +++ b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py @@ -260,9 +260,9 @@ class DeepGemmMoeQuantInfo(MoeQuantInfo): 1, 32, ], f"MXFP8 requires block_shape [1, 32], got {self.block_shape}" - assert ( - deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 - ), "MXFP8 requires DEEPGEMM_SCALE_UE8M0=True" + assert deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, ( + "MXFP8 requires DEEPGEMM_SCALE_UE8M0=True" + ) class DeepGemmRunnerCore(MoeRunnerCore): @@ -657,9 +657,9 @@ class DeepGemmRunnerCore(MoeRunnerCore): if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0: if hidden_states_scale.dtype != torch.int: b, s_mn, s_k = hidden_states_scale.shape - assert ( - s_mn % 4 == 0 and s_k % 4 == 0 - ), f"scales must be aligned to 4, but got ({b}, {s_mn}, {s_k})" + assert s_mn % 4 == 0 and s_k % 4 == 0, ( + f"scales must be aligned to 4, but got ({b}, {s_mn}, {s_k})" + ) hidden_states_scale = _cast_to_e8m0_with_rounding_up( hidden_states_scale ) @@ -1453,9 +1453,9 @@ def _varlen_deep_gemm_silu_mul_quant( # int32 UE8M0 (no follow-up transform; needs G % 4 == 0 and the # num_real_tokens grid bound) when eligible, row-major fp32 otherwise. if gemm1_alpha is not None: - assert ( - swiglu_limit is None - ), "swiglu_limit and gemm1_alpha are mutually exclusive" + assert swiglu_limit is None, ( + "swiglu_limit and gemm1_alpha are mutually exclusive" + ) assert not swizzle, "swizzle is not supported with gemm1_alpha" from sglang.kernels.ops.moe.ep_moe_kernels import ( silu_and_mul_masked_post_quant_fwd, diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py index 30de87353..efdb3cc36 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py @@ -448,7 +448,9 @@ def fused_experts_none_to_flashinfer_cutedsl_fp4( assert runner_config.activation in ( "silu", "relu2", - ), f"CuteDSL MoE supports 'silu' (gated) or 'relu2' (non-gated), got {runner_config.activation!r}." + ), ( + f"CuteDSL MoE supports 'silu' (gated) or 'relu2' (non-gated), got {runner_config.activation!r}." + ) assert quant_info.wrapper is not None, "CuteDSL v2 path requires CuteDslMoEWrapper." hidden_states = dispatch_output.hidden_states @@ -535,7 +537,9 @@ def fused_experts_flashinfer_to_flashinfer_cutedsl_fp4( assert runner_config.activation in ( "silu", "relu2", - ), f"CuteDSL MoE supports 'silu' (gated) or 'relu2' (non-gated), got {runner_config.activation!r}." + ), ( + f"CuteDSL MoE supports 'silu' (gated) or 'relu2' (non-gated), got {runner_config.activation!r}." + ) assert quant_info.wrapper is not None, "CuteDSL v2 path requires CuteDslMoEWrapper." hidden_states = dispatch_output.hidden_states @@ -630,10 +634,12 @@ def fused_experts_deepep_to_flashinfer_cutedsl_fp4( assert runner_config.activation in ( "silu", "relu2", - ), f"CuteDSL masked MoE supports 'silu' or 'relu2', got {runner_config.activation!r}." - assert ( - not runner_config.apply_router_weight_on_input - ), "apply_router_weight_on_input is not supported for Flashinfer" + ), ( + f"CuteDSL masked MoE supports 'silu' or 'relu2', got {runner_config.activation!r}." + ) + assert not runner_config.apply_router_weight_on_input, ( + "apply_router_weight_on_input is not supported for Flashinfer" + ) hidden_states, hidden_states_scale, _, _, masked_m, _ = dispatch_output diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py index ac5d1c78b..950591ca1 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py @@ -264,12 +264,12 @@ def fused_experts_none_to_flashinfer_cutlass( ) -> StandardCombineInput: from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput - assert isinstance( - quant_info, FlashInferCutlassMoeQuantInfo - ), f"Unexpected quant_info type for flashinfer_cutlass: {type(quant_info)}" - assert ( - not runner_config.apply_router_weight_on_input - ), "apply_router_weight_on_input is not supported for FlashInfer CUTLASS" + assert isinstance(quant_info, FlashInferCutlassMoeQuantInfo), ( + f"Unexpected quant_info type for flashinfer_cutlass: {type(quant_info)}" + ) + assert not runner_config.apply_router_weight_on_input, ( + "apply_router_weight_on_input is not supported for FlashInfer CUTLASS" + ) output = _run_flashinfer_cutlass( dispatch_output=dispatch_output, @@ -289,12 +289,12 @@ def fused_experts_flashinfer_to_flashinfer_cutlass( FlashinferCombineInput, ) - assert isinstance( - quant_info, FlashInferCutlassMoeQuantInfo - ), f"Unexpected quant_info type for flashinfer_cutlass: {type(quant_info)}" - assert ( - not runner_config.apply_router_weight_on_input - ), "apply_router_weight_on_input is not supported for FlashInfer CUTLASS" + assert isinstance(quant_info, FlashInferCutlassMoeQuantInfo), ( + f"Unexpected quant_info type for flashinfer_cutlass: {type(quant_info)}" + ) + assert not runner_config.apply_router_weight_on_input, ( + "apply_router_weight_on_input is not supported for FlashInfer CUTLASS" + ) output = _run_flashinfer_cutlass( dispatch_output=dispatch_output, @@ -316,9 +316,9 @@ def fused_experts_none_to_flashinfer_mxfp4( from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput from sglang.srt.layers.moe.topk import TopKOutputChecker - assert isinstance( - quant_info, FlashInferCutlassMxfp4MoeQuantInfo - ), f"Unexpected quant_info type for flashinfer_mxfp4: {type(quant_info)}" + assert isinstance(quant_info, FlashInferCutlassMxfp4MoeQuantInfo), ( + f"Unexpected quant_info type for flashinfer_mxfp4: {type(quant_info)}" + ) flashinfer_cutlass_fused_moe, ActivationType = _flashinfer_cutlass_fused_moe() diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index 22f426fda..caf6b6fe9 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -799,9 +799,9 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( # during torch.compile for piecewise cuda graph. # Use custom op wrapper for torch.compile compatibility. if use_routed_topk: - assert ( - runner_config.top_k is not None - ), "runner_config.top_k is required for flashinfer_trtllm_routed." + assert runner_config.top_k is not None, ( + "runner_config.top_k is required for flashinfer_trtllm_routed." + ) packed_topk_ids = _get_packed_topk_ids_for_flashinfer_routed(topk_output) trtllm_fp8_block_scale_routed_moe_out_wrapper( @@ -1286,12 +1286,12 @@ def fused_experts_none_to_flashinfer_trtllm_bf16( f"got '{runner_config.activation}'." ) if not use_routed_topk: - assert ( - dispatch_output.topk_output.topk_config.renormalize - ), "Renormalize is required for flashinfer trtllm moe" - assert ( - runner_config.num_fused_shared_experts == 0 - ), "Fused shared experts are not supported for flashinfer trtllm moe" + assert dispatch_output.topk_output.topk_config.renormalize, ( + "Renormalize is required for flashinfer trtllm moe" + ) + assert runner_config.num_fused_shared_experts == 0, ( + "Fused shared experts are not supported for flashinfer trtllm moe" + ) activation_type = get_activation_type( runner_config.activation, is_gated=runner_config.is_gated ) @@ -1301,9 +1301,9 @@ def fused_experts_none_to_flashinfer_trtllm_bf16( with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()): if use_routed_topk: - assert ( - runner_config.top_k is not None - ), "runner_config.top_k is required for flashinfer_trtllm_routed." + assert runner_config.top_k is not None, ( + "runner_config.top_k is required for flashinfer_trtllm_routed." + ) routing_method_type = runner_config.routing_method_type if routing_method_type is None: routing_method_type = RoutingMethodType.Default diff --git a/python/sglang/srt/layers/moe/moe_runner/hpc_ops.py b/python/sglang/srt/layers/moe/moe_runner/hpc_ops.py index 102d6cebd..6cfad99d4 100644 --- a/python/sglang/srt/layers/moe/moe_runner/hpc_ops.py +++ b/python/sglang/srt/layers/moe/moe_runner/hpc_ops.py @@ -146,12 +146,12 @@ def fused_experts_none_to_hpc_ops( "this backend also expects global top-k ids, so other quant " "methods must not run with --moe-runner-backend hpc_ops." ) - assert ( - quant_info.w13_weight.dtype == torch.float8_e4m3fn - ), f"expected fp8 w13_weight, got {quant_info.w13_weight.dtype}" - assert ( - quant_info.w2_weight.dtype == torch.float8_e4m3fn - ), f"expected fp8 w2_weight, got {quant_info.w2_weight.dtype}" + assert quant_info.w13_weight.dtype == torch.float8_e4m3fn, ( + f"expected fp8 w13_weight, got {quant_info.w13_weight.dtype}" + ) + assert quant_info.w2_weight.dtype == torch.float8_e4m3fn, ( + f"expected fp8 w2_weight, got {quant_info.w2_weight.dtype}" + ) _check_runner_config_supported(runner_config) x = dispatch_output.hidden_states diff --git a/python/sglang/srt/layers/moe/moe_runner/triton.py b/python/sglang/srt/layers/moe/moe_runner/triton.py index b634d4c4f..961b581d2 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton.py @@ -27,7 +27,6 @@ if TYPE_CHECKING: @dataclass class TritonRunnerInput(RunnerInput): - hidden_states: torch.Tensor topk_weights: torch.Tensor topk_ids: torch.Tensor @@ -42,7 +41,6 @@ class TritonRunnerInput(RunnerInput): @dataclass class TritonRunnerOutput(RunnerOutput): - hidden_states: torch.Tensor @property @@ -75,7 +73,6 @@ class TritonMoeQuantInfo(MoeQuantInfo): class TritonRunnerCore(MoeRunnerCore): - def __init__(self, config: MoeRunnerConfig): super().__init__(config) diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_kernels.py b/python/sglang/srt/layers/moe/moe_runner/triton_kernels.py index 5dd162569..7eb389fa6 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton_kernels.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton_kernels.py @@ -93,9 +93,9 @@ class TritonKernelsRunnerCore(MoeRunnerCore): triton_kernel_fused_experts_with_bias, ) - assert ( - self.config.is_gated - ), "Only gated MoEs are supported for Triton Kernels runner" + assert self.config.is_gated, ( + "Only gated MoEs are supported for Triton Kernels runner" + ) hidden_states = runner_input.hidden_states @@ -114,9 +114,9 @@ class TritonKernelsRunnerCore(MoeRunnerCore): has_bias = quant_info.w13_bias is not None or quant_info.w2_bias is not None if has_bias: - assert ( - quant_info.w13_bias is not None and quant_info.w2_bias is not None - ), "Bias execution requires both w13_bias and w2_bias" + assert quant_info.w13_bias is not None and quant_info.w2_bias is not None, ( + "Bias execution requires both w13_bias and w2_bias" + ) output = triton_kernel_fused_experts_with_bias( hidden_states=hidden_states, w1=quant_info.w13_weight, @@ -170,9 +170,9 @@ def pre_permute_standard_to_triton_kernels( hidden_states = dispatch_output.hidden_states topk_output = dispatch_output.topk_output - assert TopKOutputChecker.format_is_triton_kernels( - topk_output - ), "Triton-kernel runner expects TritonKernelTopKOutput" + assert TopKOutputChecker.format_is_triton_kernels(topk_output), ( + "Triton-kernel runner expects TritonKernelTopKOutput" + ) a_ragged_metadata, gather_indx, scatter_indx, gate_scal, n_expts_act = topk_output diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py index 758c2eacf..ad550eba3 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py @@ -699,9 +699,9 @@ def _fused_moe_kernel_sequence( # fusion=False: explicit clamp_ on intermediate_cache1 (path checker) assert swiglu_limit == 10 assert intermediate_cache1.shape == (total_tokens, N) - assert ( - _is_cuda or _is_hip or _is_xpu - ), "DeepSeek V4 only supports CUDA/HIP/XPU downstream" + assert _is_cuda or _is_hip or _is_xpu, ( + "DeepSeek V4 only supports CUDA/HIP/XPU downstream" + ) swiglu_limit_for_triton: Optional[float] = None swiglu_limit_for_silu_and_mul_clamp: Optional[float] = None @@ -709,9 +709,9 @@ def _fused_moe_kernel_sequence( if filter_expert: swiglu_limit_for_triton = swiglu_limit else: - assert ( - _is_cuda or _is_xpu - ), "fused silu_and_mul_clamp kernel is CUDA/XPU only; HIP must disable SWIGLU_CLAMP_FUSION" + assert _is_cuda or _is_xpu, ( + "fused silu_and_mul_clamp kernel is CUDA/XPU only; HIP must disable SWIGLU_CLAMP_FUSION" + ) swiglu_limit_for_silu_and_mul_clamp = swiglu_limit if not filter_expert: @@ -997,9 +997,9 @@ def fused_experts_impl( if use_int4_w4a16: assert hidden_states.shape[1] // 2 == w1.shape[2], "Hidden size mismatch" else: - assert ( - hidden_states.shape[1] == w1.shape[2] - padded_size - ), "Hidden size mismatch" + assert hidden_states.shape[1] == w1.shape[2] - padded_size, ( + "Hidden size mismatch" + ) assert topk_weights.shape == topk_ids.shape, "topk shape mismatch" assert hidden_states.is_contiguous(), "Hidden_states must be contiguous" assert w1.is_contiguous(), "Expert weights1 must be contiguous" diff --git a/python/sglang/srt/layers/moe/token_dispatcher/base.py b/python/sglang/srt/layers/moe/token_dispatcher/base.py index e718c0085..1f52798d9 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/base.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/base.py @@ -41,7 +41,6 @@ if TYPE_CHECKING: class _RemovableDispatcherHandle: - next_id = 0 # Global counter for unique IDs def __init__(self, hooks_dict: OrderedDict): @@ -56,7 +55,6 @@ class _RemovableDispatcherHandle: class DispatcherBaseHooks: - def __init__(self): self.hook_dict = OrderedDict[int, Callable]() @@ -70,7 +68,6 @@ class DispatcherBaseHooks: class _PreDispatchHooks(DispatcherBaseHooks): - def __call__( self, dispatcher: BaseDispatcher, @@ -85,7 +82,6 @@ class _PreDispatchHooks(DispatcherBaseHooks): class _PostDispatchHooks(DispatcherBaseHooks): - def __call__( self, dispatcher: BaseDispatcher, dispatch_output: DispatchOutput ) -> Optional[DispatchOutput]: @@ -97,7 +93,6 @@ class _PostDispatchHooks(DispatcherBaseHooks): class _PreCombineHooks(DispatcherBaseHooks): - def __call__( self, dispatcher: BaseDispatcher, combine_input: CombineInput ) -> Optional[CombineInput]: @@ -109,7 +104,6 @@ class _PreCombineHooks(DispatcherBaseHooks): class _PostCombineHooks(DispatcherBaseHooks): - def __call__( self, dispatcher: BaseDispatcher, hidden_states: torch.Tensor ) -> Optional[torch.Tensor]: @@ -124,7 +118,6 @@ class _PostCombineHooks(DispatcherBaseHooks): class DispatchOutputChecker: - @staticmethod def format_is_standard( dispatch_output: DispatchOutput, @@ -175,7 +168,6 @@ class DispatchOutputChecker: class DispatchOutputFormat(Enum): - STANDARD = "standard" DEEPEP_NORMAL = "deepep_normal" DEEPEP_LL = "deepep_ll" diff --git a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py b/python/sglang/srt/layers/moe/token_dispatcher/moriep.py index 2616c148f..f97221278 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/moriep.py @@ -92,7 +92,6 @@ def _aiter_supports_mxfp8_dispatch() -> bool: class MoriEPPDispatchHooks(DeepEPPDispatchHooks): - def __call__(self, dispatcher: BaseDispatcher): for hook_fun in self.hook_dict.values(): hook_fun(dispatcher) @@ -782,7 +781,6 @@ class _MoriEPDispatcherImplNormal(_MoriEPDispatcherImplBase): compute_stream.wait_stream(comm_stream) else: - ( packed_recv_hidden, recv_topk_weights, diff --git a/python/sglang/srt/layers/moe/token_dispatcher/standard.py b/python/sglang/srt/layers/moe/token_dispatcher/standard.py index ccee81c62..7bd27ab3f 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/standard.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/standard.py @@ -194,8 +194,9 @@ class StandardDispatcher(BaseDispatcher): (self.num_experts,), -1, dtype=torch.int32, device=device ) self.local_expert_mapping[ - self.moe_ep_rank - * self.num_local_routed_experts : (self.moe_ep_rank + 1) + self.moe_ep_rank * self.num_local_routed_experts : ( + self.moe_ep_rank + 1 + ) * self.num_local_routed_experts ] = torch.arange( 0, self.num_local_routed_experts, dtype=torch.int32, device=device diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 7fa5c8f9e..c6019f32c 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -244,7 +244,6 @@ class TopKConfig: class TopKOutputChecker: - @staticmethod def format_is_standard(topk_output: TopKOutput) -> TypeGuard[StandardTopKOutput]: # ===== TO BE REFACTORED ==== @@ -618,7 +617,6 @@ class TopK(BaseFusedOp): # ===== TO BE REFACTORED ==== elif get_moe_runner_backend().is_experimental_sgl_trtllm(): try: - use_standard_for_lora = bool(get_lora().enable_lora) except ValueError: use_standard_for_lora = False @@ -817,9 +815,9 @@ def fused_topk_torch_native( topk_ids = torch.topk(scores_for_choice, k=topk, dim=-1, sorted=False)[1] topk_weights = scores.gather(1, topk_ids) else: - assert ( - hidden_states.shape[0] == gating_output.shape[0] - ), f"Number of tokens mismatch, {hidden_states.shape=} vs {gating_output.shape=}" + assert hidden_states.shape[0] == gating_output.shape[0], ( + f"Number of tokens mismatch, {hidden_states.shape=} vs {gating_output.shape=}" + ) M, _ = hidden_states.shape topk_weights = torch.empty( M, topk, dtype=torch.float32, device=hidden_states.device @@ -845,9 +843,9 @@ def fused_topk_softmax_torch_raw_logits( topk: int, renormalize: bool, ): - assert ( - hidden_states.shape[0] == gating_output.shape[0] - ), f"Number of tokens mismatch, {hidden_states.shape=} vs {gating_output.shape=}" + assert hidden_states.shape[0] == gating_output.shape[0], ( + f"Number of tokens mismatch, {hidden_states.shape=} vs {gating_output.shape=}" + ) _, topk_ids = torch.topk(gating_output, k=topk, dim=-1, sorted=False) logits = gating_output.float() @@ -953,7 +951,6 @@ def fused_topk( if scoring_func == "softmax": if _use_aiter: - # Use fused_topk instead of topk_softmax to auto dispatch to the correct kernel topk_weights, topk_ids = aiter_fused_topk( hidden_states, @@ -1385,8 +1382,9 @@ def biased_topk_jit_kernel_impl( routed_scaling_factor=routed_scaling_factor, apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output, ) - topk_weights, topk_ids = topk_weights.to(torch.float32), topk_ids.to( - torch.int32 + topk_weights, topk_ids = ( + topk_weights.to(torch.float32), + topk_ids.to(torch.int32), ) return topk_weights, topk_ids @@ -1712,9 +1710,9 @@ def biased_grouped_topk_gpu( assert not apply_routed_scaling_factor_on_output, "Not implemented" token = gating_output.shape[0] device = gating_output.device - assert ( - hidden_states.shape[0] == gating_output.shape[0] - ), f"Number of tokens mismatch: hidden_states.shape[0] = {hidden_states.shape[0]}, gating_output.shape[0] = {gating_output.shape[0]}" + assert hidden_states.shape[0] == gating_output.shape[0], ( + f"Number of tokens mismatch: hidden_states.shape[0] = {hidden_states.shape[0]}, gating_output.shape[0] = {gating_output.shape[0]}" + ) bias = correction_bias.to(dtype=gating_output.dtype) scaling = routed_scaling_factor if routed_scaling_factor is not None else 1.0 @@ -1874,7 +1872,6 @@ def biased_grouped_topk_gpu( and num_experts <= 256 and topk <= 8 ): - scale = ( routed_scaling_factor if ( @@ -2259,9 +2256,9 @@ def _post_process_topk_ids( ): # Standard EP all-reduces the per-rank shared-expert outputs; without the # supplied 1/ep_size factor the shared contribution is summed ep_size times. - topk_weights[ - :, -num_fused_shared_experts: - ] *= fused_shared_experts_scaling_factor + topk_weights[:, -num_fused_shared_experts:] *= ( + fused_shared_experts_scaling_factor + ) if _is_hip and not _skip_hip_pad_mask: # Shared-expert append/remap can introduce non-zero weights after the @@ -2354,9 +2351,9 @@ def select_experts( apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output, ) elif torch_native and custom_routing_function is None: - assert ( - num_token_non_padded is None - ), "num_token_non_padded is not yet supported in fused_topk_native" + assert num_token_non_padded is None, ( + "num_token_non_padded is not yet supported in fused_topk_native" + ) assert expert_location_dispatch_info is None assert not apply_routed_scaling_factor_on_output, "Not implemented" topk_weights, topk_ids = fused_topk_native( diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index e2b4635d0..88c07c1e3 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/python/sglang/srt/layers/moe/utils.py @@ -32,7 +32,6 @@ logger = logging.getLogger(__name__) class MoeA2ABackend(Enum): - NONE = "none" DEEPEP = "deepep" MOONCAKE = "mooncake" @@ -168,7 +167,6 @@ class _MoeRunnerBackendPredicates: class MoeRunnerBackend(_MoeRunnerBackendPredicates, Enum): - AUTO = "auto" DEEP_GEMM = "deep_gemm" TRITON = "triton" @@ -241,7 +239,6 @@ class DeepEPv2Fp8ScaleFormat(NamedTuple): class DeepEPMode(Enum): - NORMAL = "normal" LOW_LATENCY = "low_latency" AUTO = "auto" @@ -486,7 +483,6 @@ def is_shared_experts_fusion_disabled() -> bool: ) moe = get_flags().moe if moe.disable_shared_experts_fusion is None: - return get_exec().moe.disable_shared_experts_fusion return moe.disable_shared_experts_fusion diff --git a/python/sglang/srt/layers/n_gram_embedding.py b/python/sglang/srt/layers/n_gram_embedding.py index 699e7c326..8d6fe33ef 100644 --- a/python/sglang/srt/layers/n_gram_embedding.py +++ b/python/sglang/srt/layers/n_gram_embedding.py @@ -9,7 +9,6 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch class NgramEmbedding(torch.nn.Module): - def __init__( self, num_embeddings: int, @@ -20,9 +19,9 @@ class NgramEmbedding(torch.nn.Module): eos_token_id: int, ): super().__init__() - assert ( - over_embedding_n > 1 - ), f"over_embedding_n must be > 1, got {over_embedding_n}" + assert over_embedding_n > 1, ( + f"over_embedding_n must be > 1, got {over_embedding_n}" + ) self.num_embeddings = num_embeddings self.embedding_dim = embedding_dim self.over_embedding_m = over_embedding_m @@ -105,9 +104,9 @@ class NgramEmbedding(torch.nn.Module): ) oe_weight_start = self.exclusive_oe_embedder_size_sums[index] oe_weight_end = self.exclusive_oe_embedder_size_sums[index + 1] - assert ( - oe_weight_end - oe_weight_start == loaded_weight.shape[0] - ), f"{oe_weight_end - oe_weight_start=} {loaded_weight.shape[0]=}" + assert oe_weight_end - oe_weight_start == loaded_weight.shape[0], ( + f"{oe_weight_end - oe_weight_start=} {loaded_weight.shape[0]=}" + ) tp_start = self.oe_embeder.shard_indices.org_vocab_start_index tp_end = self.oe_embeder.shard_indices.org_vocab_end_index to_load_start = max(oe_weight_start, tp_start) diff --git a/python/sglang/srt/layers/parameter.py b/python/sglang/srt/layers/parameter.py index fed7348c8..eef850085 100644 --- a/python/sglang/srt/layers/parameter.py +++ b/python/sglang/srt/layers/parameter.py @@ -54,9 +54,9 @@ def copy_with_check(target: torch.Tensor, loaded_weight: torch.Tensor): bf16/fp16 share the same rank, and all fp8 variants share the same rank. """ - assert ( - target.shape == loaded_weight.shape - ), f"{target.shape=}, {loaded_weight.shape=}" + assert target.shape == loaded_weight.shape, ( + f"{target.shape=}, {loaded_weight.shape=}" + ) if target.dtype == loaded_weight.dtype: target.copy_(loaded_weight) @@ -268,9 +268,9 @@ class _ColumnvLLMParameter(BasevLLMParameter): self.output_dim, shard_id * shard_size, shard_size ) - assert ( - param_data.shape == loaded_weight.shape - ), f"{param_data.shape=}, {loaded_weight.shape=}" + assert param_data.shape == loaded_weight.shape, ( + f"{param_data.shape=}, {loaded_weight.shape=}" + ) param_data.copy_(loaded_weight) diff --git a/python/sglang/srt/layers/quantization/awq/awq.py b/python/sglang/srt/layers/quantization/awq/awq.py index 2bac4a10a..6199ae8f1 100644 --- a/python/sglang/srt/layers/quantization/awq/awq.py +++ b/python/sglang/srt/layers/quantization/awq/awq.py @@ -468,7 +468,6 @@ class AWQLinearMethod(LinearMethodBase): class AWQMoEMethod(FusedMoEMethodBase): - def __init__(self, quant_config: AWQMarlinConfig): self.quant_config = quant_config self.quant_type = scalar_types.uint4 diff --git a/python/sglang/srt/layers/quantization/base_config.py b/python/sglang/srt/layers/quantization/base_config.py index c597e388e..2877c8eb4 100644 --- a/python/sglang/srt/layers/quantization/base_config.py +++ b/python/sglang/srt/layers/quantization/base_config.py @@ -234,7 +234,7 @@ class QuantizationConfig(ABC): if key in config: return config[key] raise ValueError( - f"Cannot find any of {keys} in the model's " "quantization config." + f"Cannot find any of {keys} in the model's quantization config." ) @staticmethod @@ -268,9 +268,7 @@ class QuantizationConfig(ABC): """ raise NotImplementedError() - def apply_weight_name_mapper( - self, hf_to_sglang_mapper: WeightsMapper - ): # noqa: B027 + def apply_weight_name_mapper(self, hf_to_sglang_mapper: WeightsMapper): # noqa: B027 """ Interface for models to update module names referenced in quantization configs in order to reflect the sglang model structure diff --git a/python/sglang/srt/layers/quantization/bitsandbytes.py b/python/sglang/srt/layers/quantization/bitsandbytes.py index d1d35e777..09cef9708 100644 --- a/python/sglang/srt/layers/quantization/bitsandbytes.py +++ b/python/sglang/srt/layers/quantization/bitsandbytes.py @@ -320,9 +320,9 @@ class BitsAndBytesLinearMethod(LinearMethodBase): matmul_states[i].CB = qweight[offsets[i] : offsets[i + 1]] matmul_states[i].SCB = quant_states[i].to(x.device) matmul_states[i].threshold = self.quant_config.llm_int8_threshold - matmul_states[i].has_fp16_weights = ( - self.quant_config.llm_int8_has_fp16_weight - ) + matmul_states[ + i + ].has_fp16_weights = self.quant_config.llm_int8_has_fp16_weight matmul_states[i].is_training = False if ( matmul_states[i].threshold > 0.0 diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxint4_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxint4_moe.py index cc3624a7b..3756ac26e 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxint4_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxint4_moe.py @@ -66,12 +66,12 @@ class CompressedTensorsMxInt4MoE(CompressedTensorsMoEScheme): and config.num_bits == 4 ), "MxInt4 only supports group strategy with group size 32" assert config.symmetric, "Only symmetric quantization is supported for MoE" - assert ( - get_moe_runner_backend().is_flashinfer_trtllm() - ), "MxInt4 only supports flashinfer_trtllm backend" - assert ( - not config.actorder - ), "Actorder is not supported by flashinfer_trtllm backend" + assert get_moe_runner_backend().is_flashinfer_trtllm(), ( + "MxInt4 only supports flashinfer_trtllm backend" + ) + assert not config.actorder, ( + "Actorder is not supported by flashinfer_trtllm backend" + ) self.moe_ep_rank = get_parallel().moe_ep_rank if self.quant_config.quant_format != CompressionFormat.pack_quantized.value: @@ -95,9 +95,9 @@ class CompressedTensorsMxInt4MoE(CompressedTensorsMoEScheme): params_dtype: torch.dtype, **extra_weight_attrs, ): - assert ( - params_dtype == torch.bfloat16 - ), f"Params dtype should be torch.bfloat16, but got: {params_dtype}" + assert params_dtype == torch.bfloat16, ( + f"Params dtype should be torch.bfloat16, but got: {params_dtype}" + ) extra_weight_attrs.update({"quant_method": self.strategy}) w13_weight = torch.nn.Parameter( @@ -301,9 +301,9 @@ class CompressedTensorsMxInt4MoE(CompressedTensorsMoEScheme): ) -> CombineInput: from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - assert ( - self.moe_runner_config.is_gated - ), "Only gated MoEs are supported for flashinfer mxint4" + assert self.moe_runner_config.is_gated, ( + "Only gated MoEs are supported for flashinfer mxint4" + ) x = dispatch_output.hidden_states topk_output = dispatch_output.topk_output diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4_moe.py index 245bc3f19..c72443183 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4_moe.py @@ -31,7 +31,6 @@ if TYPE_CHECKING: class CompressedTensorsW4A4Nvfp4MoE(CompressedTensorsMoEScheme): - def __init__(self): if not get_platform().is_blackwell: raise ValueError( @@ -337,9 +336,9 @@ class CompressedTensorsW4A4Nvfp4MoE(CompressedTensorsMoEScheme): FlashInferCutlassMoeQuantInfo, ) - assert ( - not self.moe_runner_config.apply_router_weight_on_input - ), "apply_router_weight_on_input is not supported for Flashinfer" + assert not self.moe_runner_config.apply_router_weight_on_input, ( + "apply_router_weight_on_input is not supported for Flashinfer" + ) quant_info = FlashInferCutlassMoeQuantInfo( quant_type="fp4", diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8_moe.py index ef8d56abe..59bb910a3 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8_moe.py @@ -94,7 +94,9 @@ class CompressedTensorsW4AFP8MoE(CompressedTensorsMoEScheme): assert config.symmetric, "Only symmetric quantization is supported" assert ( self.quant_config.quant_format == CompressionFormat.pack_quantized.value - ), f"W4AFP8MoE requires pack-quantized format, got {self.quant_config.quant_format}" + ), ( + f"W4AFP8MoE requires pack-quantized format, got {self.quant_config.quant_format}" + ) @classmethod def get_min_capability(cls) -> int: @@ -288,9 +290,9 @@ class CompressedTensorsW4AFP8MoE(CompressedTensorsMoEScheme): from sglang.srt.layers.moe.cutlass_w4a8_moe import cutlass_w4a8_moe from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - assert ( - self.moe_runner_config.activation == "silu" - ), "Only SiLU activation is supported." + assert self.moe_runner_config.activation == "silu", ( + "Only SiLU activation is supported." + ) x = dispatch_output.hidden_states topk_output = dispatch_output.topk_output diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_int8_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_int8_moe.py index 45a25c6f0..835d50b36 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_int8_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_int8_moe.py @@ -28,7 +28,6 @@ logger = logging.getLogger(__name__) class NPUCompressedTensorsW4A8Int8DynamicMoE(CompressedTensorsMoEScheme): - ### TODO: Get rid of code duplication with python/sglang/srt/modelslim/modelslim_moe.py @OrangeRedeng @TamirBaydasov def __init__(self, quantization_config) -> None: self.group_size = 0 diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py index 6a2e23cb3..1e0ac8788 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py @@ -49,7 +49,6 @@ logger = logging.getLogger(__name__) class CompressedTensorsW8A8Fp8MoE(CompressedTensorsMoEScheme): - def __init__(self, weight_quant, input_quant): self.weight_quant = weight_quant self.input_quant = input_quant @@ -220,9 +219,9 @@ class CompressedTensorsW8A8Fp8MoE(CompressedTensorsMoEScheme): # INPUT_SCALES if self.static_input_scales: - assert ( - self.input_quant.strategy == QuantizationStrategy.TENSOR - ), "Only per-tensor quantization is supported for static input scales" + assert self.input_quant.strategy == QuantizationStrategy.TENSOR, ( + "Only per-tensor quantization is supported for static input scales" + ) w13_input_scale = torch.nn.Parameter( torch.ones(num_experts, dtype=torch.float32), requires_grad=False ) diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py index 7634e8197..e6c4fae67 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py @@ -31,7 +31,6 @@ if _is_cuda: class CompressedTensorsW8A8Int8(CompressedTensorsLinearScheme): - def __init__( self, strategy: str, is_static_input_scheme: bool, input_symmetric: bool ): @@ -180,7 +179,6 @@ class CompressedTensorsW8A8Int8(CompressedTensorsLinearScheme): class NPUCompressedTensorsW8A8Int8(CompressedTensorsW8A8Int8): - def __init__( self, strategy: str, is_static_input_scheme: bool, input_symmetric: bool ): diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8_moe.py index 1669899ee..497a6a3a5 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8_moe.py @@ -28,7 +28,6 @@ logger = logging.getLogger(__name__) class NPUCompressedTensorsW8A8Int8DynamicMoE(CompressedTensorsMoEScheme): - def __init__(self, weight_quant, input_quant): self.weight_quant = weight_quant self.input_quant = input_quant diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py index c06834bf9..9d182e066 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py @@ -63,7 +63,6 @@ class GPTQMarlinState(Enum): class CompressedTensorsWNA16MoE(CompressedTensorsMoEScheme): - def __init__( self, quant_config: CompressedTensorsConfig, @@ -452,9 +451,9 @@ class CompressedTensorsWNA16MoE(CompressedTensorsMoEScheme): ) from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - assert ( - self.moe_runner_config.activation == "silu" - ), "Only SiLU activation is supported." + assert self.moe_runner_config.activation == "silu", ( + "Only SiLU activation is supported." + ) x = dispatch_output.hidden_states topk_output = dispatch_output.topk_output @@ -556,16 +555,15 @@ class CompressedTensorsWNA16TritonMoE(CompressedTensorsWNA16MoE): layer: torch.nn.Module, dispatch_output: StandardDispatchOutput, ) -> CombineInput: - assert ( - self.moe_runner_config.activation == "silu" - ), "Only SiLU activation is supported." + assert self.moe_runner_config.activation == "silu", ( + "Only SiLU activation is supported." + ) quant_info = self.get_triton_quant_info(layer) return self.runner.run(dispatch_output, quant_info) class NPUCompressedTensorsW4A16Int4DynamicMoE(CompressedTensorsMoEScheme): - def __init__(self, quantization_config) -> None: self.pack_factor = 8 # weight dtype is int4, but use int32 to create target = ( diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 72935cffb..1ed1d21c8 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -390,9 +390,9 @@ class Fp8Config(QuantizationConfig): fp8_method = Fp8MoEMethod(self) if self.is_fp4_experts and self.dequant_fp4_to_fp8: - assert ( - get_moe_runner_backend().is_auto() - ), f"{get_moe_runner_backend()} is not compatible with SGLANG_DSV4_FP4_DEQUANT=1" + assert get_moe_runner_backend().is_auto(), ( + f"{get_moe_runner_backend()} is not compatible with SGLANG_DSV4_FP4_DEQUANT=1" + ) return fp8_method if self.is_fp4_experts and get_moe_runner_backend().is_marlin(): @@ -697,9 +697,9 @@ class Fp8LinearMethod(LinearMethodBase): ) layer.input_scale = None elif _is_cpu: - assert ( - _is_cpu_amx_available - ), "Fp8LinearMethod on CPU requires that CPU has AMX support" + assert _is_cpu_amx_available, ( + "Fp8LinearMethod on CPU requires that CPU has AMX support" + ) _amx_process_weight_after_loading(layer, ["weight"]) layer.weight_scale_inv = torch.nn.Parameter( layer.weight_scale_inv.data, requires_grad=False @@ -1098,9 +1098,9 @@ class Fp8MoEMethod(FusedMoEMethodBase): # they never call create_moe_runner, so moe_runner_config is unset. self._owns_moe_runner = False if get_moe_runner_backend().is_cutlass(): - assert ( - cutlass_fp8_supported() - ), "cutlass_fp8 MoE requires CUDA 12.0+ with SM90 or CUDA 12.4+ with SM89" + assert cutlass_fp8_supported(), ( + "cutlass_fp8 MoE requires CUDA 12.0+ with SM90 or CUDA 12.4+ with SM89" + ) assert self.block_quant, "cutlass_fp8 MoE requires block quantization" assert ( get_platform().is_sm100 @@ -1666,9 +1666,9 @@ class Fp8MoEMethod(FusedMoEMethodBase): layer.w2_weight.copy_(t) del t elif _is_cpu: - assert ( - _is_cpu_amx_available - ), "Fp8MoEMethod on CPU requires that CPU has AMX support" + assert _is_cpu_amx_available, ( + "Fp8MoEMethod on CPU requires that CPU has AMX support" + ) _amx_process_weight_after_loading(layer, ["w13_weight", "w2_weight"]) else: # For fp8 moe run with deepgemm, the expert weights and scales need be requantized to ue8m0 @@ -2319,9 +2319,9 @@ class Fp8MoEMethod(FusedMoEMethodBase): int4_rescale = ( layer.w13_weight_scale[expert_id][shard_id] / max_w13_scale_fp8 ) - layer.w13_weight_scale1[expert_id][ - start : start + shard_size - ] *= int4_rescale + layer.w13_weight_scale1[expert_id][start : start + shard_size] *= ( + int4_rescale + ) start += shard_size layer.w13_weight_scale = torch.nn.Parameter(max_w13_scales, requires_grad=False) @@ -2546,7 +2546,6 @@ class Fp8MoEMethod(FusedMoEMethodBase): return StandardCombineInput(hidden_states=output) if self.runner.runner_backend.is_deep_gemm(): - w13_weight = layer.w13_weight w2_weight = layer.w2_weight diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 7f4145353..4801ae65b 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -1576,9 +1576,9 @@ def quant_weight_ue8m0( weight_block_size: List[int], ): assert weight_block_size == [128, 128] - assert ( - weight_dequant.dtype == torch.bfloat16 - ), f"{weight_dequant.dtype=} {weight_dequant.shape=}" + assert weight_dequant.dtype == torch.bfloat16, ( + f"{weight_dequant.dtype=} {weight_dequant.shape=}" + ) *batch_dims, n, k = weight_dequant.shape @@ -1665,9 +1665,9 @@ def inverse_transform_scale_ue8m0(sf_packed, mn): sf_fp32 = _inverse_transform_scale_ue8m0_impl(sf_packed) # Can call consistency check every time since this is only called on startup sf_packed_recreated = transform_scale_ue8m0(sf_fp32, mn=mn, use_torch_impl=True) - assert torch.all( - sf_packed == sf_packed_recreated - ), f"{sf_packed=} {sf_packed_recreated=} {sf_fp32=}" + assert torch.all(sf_packed == sf_packed_recreated), ( + f"{sf_packed=} {sf_packed_recreated=} {sf_fp32=}" + ) return sf_fp32 diff --git a/python/sglang/srt/layers/quantization/gguf.py b/python/sglang/srt/layers/quantization/gguf.py index cc4ed6e1e..1d8b4ba3b 100644 --- a/python/sglang/srt/layers/quantization/gguf.py +++ b/python/sglang/srt/layers/quantization/gguf.py @@ -569,9 +569,9 @@ class GGUFMoEMethod(FusedMoEMethodBase): from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - assert ( - self.moe_runner_config.activation == "silu" - ), "Only SiLU activation is supported." + assert self.moe_runner_config.activation == "silu", ( + "Only SiLU activation is supported." + ) x = dispatch_output.hidden_states topk_output = dispatch_output.topk_output diff --git a/python/sglang/srt/layers/quantization/gptq/gptq.py b/python/sglang/srt/layers/quantization/gptq/gptq.py index e6a6c65cc..123789bad 100644 --- a/python/sglang/srt/layers/quantization/gptq/gptq.py +++ b/python/sglang/srt/layers/quantization/gptq/gptq.py @@ -352,7 +352,7 @@ class GPTQMarlinConfig(QuantizationConfig): if (weight_bits, is_sym) not in self.TYPE_MAP: raise ValueError( - "Unsupported quantization config: " f"bits={weight_bits}, sym={is_sym}" + f"Unsupported quantization config: bits={weight_bits}, sym={is_sym}" ) # (num_bits, is_sym) -> quant_type @@ -527,7 +527,6 @@ class GPTQLinearMethod(LinearMethodBase): class GPTQMoEMethod(FusedMoEMethodBase): - def __init__(self, quant_config: GPTQConfig): super().__init__() self.quant_config = quant_config diff --git a/python/sglang/srt/layers/quantization/kv_cache.py b/python/sglang/srt/layers/quantization/kv_cache.py index 853f0764e..c8e1f0053 100644 --- a/python/sglang/srt/layers/quantization/kv_cache.py +++ b/python/sglang/srt/layers/quantization/kv_cache.py @@ -74,9 +74,7 @@ class BaseKVCacheMethod(QuantizeMethodBase): v_scale *= 2 if not isinstance(k_scale, float) or not isinstance(v_scale, float): - raise ValueError( - "Only support per-tensor scaling factor " "for fp8 KV cache" - ) + raise ValueError("Only support per-tensor scaling factor for fp8 KV cache") # These are used in the final Attention.forward() layer.k_scale.copy_(k_scale) diff --git a/python/sglang/srt/layers/quantization/kvfp4_tensor.py b/python/sglang/srt/layers/quantization/kvfp4_tensor.py index 4412b33fa..3191280b9 100644 --- a/python/sglang/srt/layers/quantization/kvfp4_tensor.py +++ b/python/sglang/srt/layers/quantization/kvfp4_tensor.py @@ -273,9 +273,9 @@ class NVFP4KVQuantizeUtil: ) return output_2d.reshape(b, m, -1) else: - assert ( - get_platform().is_sm90 - ), "NVFP4 KV cache dequantize requires SM100/SM120 or SM90 fallback GPU" + assert get_platform().is_sm90, ( + "NVFP4 KV cache dequantize requires SM100/SM120 or SM90 fallback GPU" + ) # Pure PyTorch fallback for SM90 n = n_half * 2 fp4_vals = torch.empty( diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index ee232f8f2..7df1969e7 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -344,9 +344,7 @@ class ModelOptQuantConfig(QuantizationConfig): def get_scaled_act_names(self) -> List[str]: return [] - def apply_weight_name_mapper( - self, hf_to_sglang_mapper: WeightsMapper - ): # noqa: B027 + def apply_weight_name_mapper(self, hf_to_sglang_mapper: WeightsMapper): # noqa: B027 # Map excluded module patterns from HF layout to sglang layout. # Ref: HF hf_quant_config.json for nvidia/Kimi-K2.5-NVFP4 # https://huggingface.co/nvidia/Kimi-K2.5-NVFP4/blob/main/hf_quant_config.json @@ -2558,8 +2556,9 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): assert w.shape == (layer.num_experts,) assert layer.moe_ep_size * layer.num_local_experts == layer.num_experts return w[ - layer.moe_ep_rank - * layer.num_local_experts : (layer.moe_ep_rank + 1) + layer.moe_ep_rank * layer.num_local_experts : ( + layer.moe_ep_rank + 1 + ) * layer.num_local_experts ] @@ -2675,9 +2674,9 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): "w13": layer.w13_weight.shape[2] * 2 // block_size, "w2": layer.w2_weight.shape[2] * 2 // block_size, } - assert ( - weight_scale.shape[-1] == expected_blocks[name] - ), f"Expected {name}_weight_scale.dim(2) == {expected_blocks[name]}, got {weight_scale.shape[-1]}" + assert weight_scale.shape[-1] == expected_blocks[name], ( + f"Expected {name}_weight_scale.dim(2) == {expected_blocks[name]}, got {weight_scale.shape[-1]}" + ) else: if weight_scale.shape[assert_dim] % 4 != 0: logger.warning( @@ -2686,9 +2685,9 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): tuple(weight_scale.shape), getattr(self.quant_config, "group_size", None), ) - assert ( - weight_scale.dtype == torch.float8_e4m3fn - ), f"{name} Weight Blockscale must be represented as FP8-E4M3" + assert weight_scale.dtype == torch.float8_e4m3fn, ( + f"{name} Weight Blockscale must be represented as FP8-E4M3" + ) # Weight processing based on strategy if ( @@ -3021,9 +3020,9 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): FlashInferCutlassMoeQuantInfo, ) - assert ( - not moe_runner_config.apply_router_weight_on_input - ), "apply_router_weight_on_input is not supported for Flashinfer" + assert not moe_runner_config.apply_router_weight_on_input, ( + "apply_router_weight_on_input is not supported for Flashinfer" + ) quant_info = FlashInferCutlassMoeQuantInfo( quant_type="fp4", w13_weight=layer.w13_weight, diff --git a/python/sglang/srt/layers/quantization/modelslim/modelslim.py b/python/sglang/srt/layers/quantization/modelslim/modelslim.py index 83ed583b0..3e1500b06 100644 --- a/python/sglang/srt/layers/quantization/modelslim/modelslim.py +++ b/python/sglang/srt/layers/quantization/modelslim/modelslim.py @@ -164,9 +164,9 @@ class ModelSlimConfig(QuantizationConfig): if rest.startswith(("embed.", "embed_tokens.", "head.", "lm_head.")): continue if rest.startswith("markov_head."): - alias = f"markov_head.{rest[len('markov_head.'):]}" + alias = f"markov_head.{rest[len('markov_head.') :]}" elif rest.startswith("confidence_head."): - alias = f"confidence_head.{rest[len('confidence_head.'):]}" + alias = f"confidence_head.{rest[len('confidence_head.') :]}" else: mapped_rest = rest if mapped_rest.startswith("attn."): @@ -326,8 +326,7 @@ class ModelSlimConfig(QuantizationConfig): return scheme_class(quant_config=self.quant_description, prefix=prefix) logger.warning( - f"Unsupported Linear modelslim scheme: " - f"{quant_schemes} in layer: {prefix}" + f"Unsupported Linear modelslim scheme: {quant_schemes} in layer: {prefix}" ) return None @@ -473,7 +472,6 @@ class ModelSlimConfig(QuantizationConfig): class ModelSlimLinearMethod(_NPULinearMethodBase): - def __init__(self, quantization_config: ModelSlimConfig): self.quantization_config = quantization_config diff --git a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_mxfp8.py b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_mxfp8.py index 1b747630e..8c82ed932 100644 --- a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_mxfp8.py +++ b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_mxfp8.py @@ -24,7 +24,6 @@ MXFP8_BLOCK_SIZE = 32 class ModelSlimMXFP8Scheme(ModelSlimLinearScheme): - def __init__( self, quant_config: Optional[Dict[str, any]] = None, diff --git a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_int4.py b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_int4.py index c55d71668..18bafb50f 100644 --- a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_int4.py +++ b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_int4.py @@ -16,7 +16,6 @@ from sglang.srt.utils import set_weight_attrs class ModelSlimW4A4Int4(ModelSlimLinearScheme): - def __init__( self, quant_config: Dict[str, any], diff --git a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w8a8_int8.py b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w8a8_int8.py index ff9bc3e8e..a8a9baf34 100644 --- a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w8a8_int8.py +++ b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w8a8_int8.py @@ -19,7 +19,6 @@ from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimLinearSche class ModelSlimW8A8Int8(ModelSlimLinearScheme): - def __init__( self, quant_config: Dict[str, any], diff --git a/python/sglang/srt/layers/quantization/moe_wna16.py b/python/sglang/srt/layers/quantization/moe_wna16.py index e5221c8a9..711f361bb 100644 --- a/python/sglang/srt/layers/quantization/moe_wna16.py +++ b/python/sglang/srt/layers/quantization/moe_wna16.py @@ -203,7 +203,6 @@ class MoeWNA16Config(QuantizationConfig): return UnquantizedFusedMoEMethod() return UnquantizedLinearMethod() elif isinstance(layer, LinearBase): - if self.linear_quant_method == "gptq": if self.use_marlin: return GPTQMarlinConfig.from_config( @@ -386,9 +385,9 @@ class MoeWNA16Method(FusedMoEMethodBase): layer: torch.nn.Module, dispatch_output: StandardDispatchOutput, ) -> CombineInput: - assert ( - self.moe_runner_config.activation == "silu" - ), "Only SiLU activation is supported." + assert self.moe_runner_config.activation == "silu", ( + "Only SiLU activation is supported." + ) quant_info = self.get_triton_quant_info(layer) return self.runner.run(dispatch_output, quant_info) diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index 2f9155527..d1921569e 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/python/sglang/srt/layers/quantization/mxfp4.py @@ -305,7 +305,6 @@ def quant_dequant_mxfp4( class Mxfp4Config(QuantizationConfig): - def __init__( self, ignored_layers: Optional[list[str]] = None, @@ -327,7 +326,6 @@ class Mxfp4Config(QuantizationConfig): is_checkpoint_mxfp4_serialized=is_checkpoint_mxfp4_serialized ) else: - platform = torch.cuda.get_device_properties(0).gcnArchName raise ValueError( f"Current platform {platform} not support mxfp4 computation" @@ -386,7 +384,6 @@ class Mxfp4Config(QuantizationConfig): class Mxfp4MoEMethod(FusedMoEMethodBase): - def __init__( self, prefix: str, @@ -1012,7 +1009,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): return if self.use_triton_kernels: - from triton_kernels.matmul import FlexCtx, PrecisionConfig w13_weight_bias = layer.w13_weight_bias.to(torch.float32) @@ -1628,13 +1624,13 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): is_standard = TopKOutputChecker.format_is_standard(topk_output) # The situ path accepts precomputed (standard) routing; the # public path below is logits-only. - assert is_standard or TopKOutputChecker.format_is_bypassed( - topk_output - ), f"unsupported topk format: {topk_output.format}" + assert is_standard or TopKOutputChecker.format_is_bypassed(topk_output), ( + f"unsupported topk format: {topk_output.format}" + ) if is_standard: - assert ( - self.moe_runner_config.activation == "situ" - ), "standard topk output only wired for the situ path" + assert self.moe_runner_config.activation == "situ", ( + "standard topk output only wired for the situ path" + ) top_k = topk_output.topk_ids.shape[1] router_logits = None else: @@ -1828,9 +1824,9 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): TritonKernelsQuantInfo, ) - assert ( - layer.moe_ep_size == 1 - ), "Expert parallel is not supported when using triton kernels" + assert layer.moe_ep_size == 1, ( + "Expert parallel is not supported when using triton kernels" + ) quant_info = TritonKernelsQuantInfo( w13_weight=( self.w13_weight_triton_tensor diff --git a/python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py b/python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py index 585e9d422..52503455e 100644 --- a/python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py +++ b/python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py @@ -159,7 +159,7 @@ class Mxfp4MarlinMoEMethod: log_info_on_rank0( logger, - f"Preparing MXFP4 experts for Marlin backend " f"(layer: {self.prefix})...", + f"Preparing MXFP4 experts for Marlin backend (layer: {self.prefix})...", ) if self.runner.config.gemm1_alpha is not None: deinterleave_moe_mxfp4_w13_for_marlin(layer) diff --git a/python/sglang/srt/layers/quantization/mxfp4_tensor.py b/python/sglang/srt/layers/quantization/mxfp4_tensor.py index 76cb92c54..781f0979e 100644 --- a/python/sglang/srt/layers/quantization/mxfp4_tensor.py +++ b/python/sglang/srt/layers/quantization/mxfp4_tensor.py @@ -50,9 +50,9 @@ class MXFP4QuantizeUtil: new_data = ( right_side.clone() << 4 ) # Put odd indices (higher addresses) in high bits - new_data[ - ..., : left_side.shape[-1] - ] += left_side # Put even indices in low bits + new_data[..., : left_side.shape[-1]] += ( + left_side # Put even indices in low bits + ) return new_data if block_size is None: diff --git a/python/sglang/srt/layers/quantization/petit.py b/python/sglang/srt/layers/quantization/petit.py index a087a9eb8..62238dd0a 100644 --- a/python/sglang/srt/layers/quantization/petit.py +++ b/python/sglang/srt/layers/quantization/petit.py @@ -178,7 +178,7 @@ class PetitNvFp4LinearMethod(LinearMethodBase): layer.output_size_per_partition = output_size_per_partition if input_size_per_partition % 16 != 0: raise ValueError( - "Unsupported model when in features size is " "not multiple of 16" + "Unsupported model when in features size is not multiple of 16" ) weight_dtype = ( diff --git a/python/sglang/srt/layers/quantization/petit_utils.py b/python/sglang/srt/layers/quantization/petit_utils.py index 529869f24..d0e723c70 100644 --- a/python/sglang/srt/layers/quantization/petit_utils.py +++ b/python/sglang/srt/layers/quantization/petit_utils.py @@ -47,7 +47,7 @@ def _check_petit_nvfp4_supported( if group_size is not None and group_size != 16: return ( False, - "Petit currently only supports: group_size=16" " quantizations.", + "Petit currently only supports: group_size=16 quantizations.", ) return (True, None) diff --git a/python/sglang/srt/layers/quantization/quark/quark.py b/python/sglang/srt/layers/quantization/quark/quark.py index 60db3f52b..c37101056 100644 --- a/python/sglang/srt/layers/quantization/quark/quark.py +++ b/python/sglang/srt/layers/quantization/quark/quark.py @@ -287,7 +287,6 @@ _SHARED_EXPERT_BODY_PROJ_SUFFIXES: tuple[str, ...] = ( class QuarkConfig(QuantizationConfig): - def __init__( self, quant_config: dict[str, Any] | None = None, @@ -963,7 +962,6 @@ class QuarkConfig(QuantizationConfig): class QuarkLinearMethod(LinearMethodBase): - def __init__(self, quantization_config: QuarkConfig): self.quantization_config = quantization_config self.quant_config = quantization_config @@ -1016,7 +1014,6 @@ class QuarkLinearMethod(LinearMethodBase): class QuarkFusedMoEMethod(FusedMoEMethodBase): - def __init__(self, quantization_config: QuarkConfig): self.quantization_config = quantization_config diff --git a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py index c388d23ea..f860a3fbe 100644 --- a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py +++ b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py @@ -227,7 +227,6 @@ OCP_MX_BLOCK_SIZE = 32 class QuarkW4A4MXFP4(QuarkLinearScheme): - # PackedvLLMParameter / ModelWeightParameter (online and NVFP4->MXFP4 # paths) only implement the v2 loader API. requires_weight_loader_v2 = True @@ -595,7 +594,6 @@ class QuarkW4A4MXFP4(QuarkLinearScheme): # Materialize FP8 parameters on first load on device (there may be several shards for a single layer parameter, e.g. q_proj, k_proj, v_proj). if is_weight_or_weight_scale and not layer._fp8_weight_materialized: - # Sanity check. assert layer.weight.device.type == "meta" @@ -702,7 +700,9 @@ class QuarkW4A4MXFP4(QuarkLinearScheme): 2, 3, 5, - ], "For tuple input, only (x, x_s), (x, x_s, y), or (x, y, S1, S2, out_dtype) formats are accepted" + ], ( + "For tuple input, only (x, x_s), (x, x_s, y), or (x, y, S1, S2, out_dtype) formats are accepted" + ) if len(x) == 2: x, x_s = x elif len(x) == 3: diff --git a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py index cd86f1965..12243eb5c 100644 --- a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py +++ b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py @@ -78,7 +78,6 @@ OCP_MX_BLOCK_SIZE = 32 class QuarkW4A4MXFp4MoE(QuarkMoEScheme): - def __init__( self, weight_config: dict[str, Any], @@ -594,9 +593,9 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme): else: raise ValueError("Expected w13 or w2.") - assert ( - current_loaded <= target_loaded_numel - ), f"target_loaded_numel={target_loaded_numel}, current_loaded={current_loaded}" + assert current_loaded <= target_loaded_numel, ( + f"target_loaded_numel={target_loaded_numel}, current_loaded={current_loaded}" + ) # Delay online quantization until all tensor shards (e.g. w1 and w3) are loaded, to avoid having to re-quantize later on. if is_w13 and layer._w13_loaded_numel == target_loaded_numel: @@ -634,7 +633,6 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme): # Materialize FP8 parameters on first load from meta device. Adds a small but manageable overhead compared to materializing one by one - but weights are loaded in order layer by layer so it is fine. with layer._fp8_loading_lock: - if not layer._fp8_materialized: # w13_weight assert layer.w13_weight.device.type == "meta" diff --git a/python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8.py b/python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8.py index 2cfd9d423..fc16e4466 100644 --- a/python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8.py +++ b/python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8.py @@ -30,7 +30,6 @@ if _use_aiter: class QuarkW8A8Fp8(QuarkLinearScheme): - def __init__( self, weight_config: dict[str, Any], input_config: Optional[dict[str, Any]] ): diff --git a/python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8_moe.py b/python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8_moe.py index 4583e3dab..673ab108b 100644 --- a/python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8_moe.py +++ b/python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8_moe.py @@ -35,7 +35,6 @@ if _use_aiter: class QuarkW8A8FP8MoE(QuarkMoEScheme): - def __init__(self, weight_config: dict[str, Any], input_config: dict[str, Any]): self.is_static_input_scheme: bool = False self.input_qscheme = None @@ -135,9 +134,9 @@ class QuarkW8A8FP8MoE(QuarkMoEScheme): # INPUT_SCALES if self.is_static_input_scheme: - assert ( - self.input_qscheme == "per_tensor" - ), "Only per-tensor quantization is supported for static input scales" + assert self.input_qscheme == "per_tensor", ( + "Only per-tensor quantization is supported for static input scales" + ) w13_input_scale = torch.nn.Parameter( torch.ones(num_experts, dtype=torch.float32), requires_grad=False ) diff --git a/python/sglang/srt/layers/quantization/quark/weights.py b/python/sglang/srt/layers/quantization/quark/weights.py index c4926a913..beaef8b89 100644 --- a/python/sglang/srt/layers/quantization/quark/weights.py +++ b/python/sglang/srt/layers/quantization/quark/weights.py @@ -67,9 +67,9 @@ def _load_gptoss_quark_expert_weights(model, weights, quark_expert_pat): moe_ep_size = get_parallel().moe_ep_size intermediate_size = model.config.intermediate_size - assert ( - intermediate_size % mxfp4_block == 0 - ), f"{intermediate_size=} must be divisible by {mxfp4_block=}" + assert intermediate_size % mxfp4_block == 0, ( + f"{intermediate_size=} must be divisible by {mxfp4_block=}" + ) intermediate_size_block = intermediate_size // mxfp4_block per_rank_intermediate_size_block = math.ceil(intermediate_size_block / moe_tp_size) diff --git a/python/sglang/srt/layers/quantization/quark_int4fp8_moe.py b/python/sglang/srt/layers/quantization/quark_int4fp8_moe.py index 1fdb45780..2813260f3 100644 --- a/python/sglang/srt/layers/quantization/quark_int4fp8_moe.py +++ b/python/sglang/srt/layers/quantization/quark_int4fp8_moe.py @@ -387,9 +387,9 @@ class QuarkInt4Fp8MoEMethod(FusedMoEMethodBase): int4_rescale = ( layer.w13_fp8_scale[expert_id][shard_id] / max_w13_scale_fp8 ) - layer.w13_int4_scale[expert_id][ - start : start + shard_size - ] *= int4_rescale + layer.w13_int4_scale[expert_id][start : start + shard_size] *= ( + int4_rescale + ) start += shard_size layer.w13_fp8_scale = torch.nn.Parameter(max_w13_scales, requires_grad=False) @@ -432,9 +432,9 @@ class QuarkInt4Fp8MoEMethod(FusedMoEMethodBase): moe_runner_config = self.moe_runner_config # TODO: add triton kernel and add check get_bool_env_var("CK_MOE") - assert ( - not moe_runner_config.no_combine - ), f"no_combine={moe_runner_config.no_combine} is not supported." + assert not moe_runner_config.no_combine, ( + f"no_combine={moe_runner_config.no_combine} is not supported." + ) quant_info = AiterMoeQuantInfo( w13_weight=layer.w13_weight, diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index 2dbd79041..6c46caff7 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -191,8 +191,7 @@ def initialize_bf16_gemm_config(server_args: ServerArgs) -> None: ) if not get_platform().is_sm100: raise ValueError( - f"--bf16-gemm-backend {backend.value} requires " - "SM100/SM103 (Blackwell)" + f"--bf16-gemm-backend {backend.value} requires SM100/SM103 (Blackwell)" ) from sglang.kernels.ops.gemm.cutedsl_bf16_gemm import ( @@ -1079,11 +1078,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, BaseFusedOp): return StandardCombineInput(hidden_states=output) else: assert backend.is_triton() - assert ( - moe_runner_config.activation == "silu" - ), f"activation = {moe_runner_config.activation} is not supported \ + assert moe_runner_config.activation == "silu", ( + f"activation = {moe_runner_config.activation} is not supported \ for Triton PATH, please drop --moe-runner-backend triton to use \ the sgl-kernel-xpu path, which supports more activations." + ) quant_info = self.get_triton_quant_info(layer) return self.runner.run(dispatch_output, quant_info) diff --git a/python/sglang/srt/layers/quantization/utils.py b/python/sglang/srt/layers/quantization/utils.py index c1e56460e..5a130d6e6 100644 --- a/python/sglang/srt/layers/quantization/utils.py +++ b/python/sglang/srt/layers/quantization/utils.py @@ -453,9 +453,9 @@ def quantize_weights( zero_points: bool = False, ref_zero_points_after_scales: bool = False, ): - assert ( - quant_type.is_integer() - ), "Floating point quantization may work but has not been tested" + assert quant_type.is_integer(), ( + "Floating point quantization may work but has not been tested" + ) assert not zero_points or group_size is not None, ( "to have group zero points, group_size must be provided " "(-1 group_size is channelwise)" @@ -555,12 +555,12 @@ def gptq_quantize_weights( size_k, _ = w.shape assert w.is_floating_point(), "w must be float" - assert ( - quant_type in SUPPORTED_GPTQ_QUANT_TYPES - ), f"Unsupported gptq type = {quant_type}" - assert group_size in SUPPORTED_GROUP_SIZES + [ - size_k - ], f"Unsupported groupsize = {group_size}" + assert quant_type in SUPPORTED_GPTQ_QUANT_TYPES, ( + f"Unsupported gptq type = {quant_type}" + ) + assert group_size in SUPPORTED_GROUP_SIZES + [size_k], ( + f"Unsupported groupsize = {group_size}" + ) w_ref, w_q, w_s, _ = quantize_weights(w, quant_type, group_size) @@ -568,10 +568,10 @@ def gptq_quantize_weights( g_idx = torch.empty(0, dtype=torch.int, device=w.device) rand_perm = torch.empty(0, dtype=torch.int, device=w.device) if act_order: - assert ( - group_size < size_k - ), "For act_order, groupsize = {} must be less than size_k = {}".format( - group_size, size_k + assert group_size < size_k, ( + "For act_order, groupsize = {} must be less than size_k = {}".format( + group_size, size_k + ) ) w_ref, w_q, g_idx, rand_perm = permute_rows(w_q, w_ref, group_size, test_perm) @@ -674,18 +674,14 @@ def prepare_static_weights_for_trtllm_fp4_moe( ) # packed fp4 gemm1_scales_linear_fp4 = gemm1_scales_linear_fp4_bytes.view( torch.float8_e4m3fn - ).reshape( - num_experts, gemm1_rows, hidden_size // 16 - ) # fp8 scaling factors + ).reshape(num_experts, gemm1_rows, hidden_size // 16) # fp8 scaling factors gemm2_weights_fp4 = gemm2_weights.view(torch.float8_e4m3fn).reshape( num_experts, hidden_size, intermediate_size // 2 ) # packed fp4 gemm2_scales_linear_fp4 = gemm2_scales_linear_fp4_bytes.view( torch.float8_e4m3fn - ).reshape( - num_experts, hidden_size, intermediate_size // 16 - ) # fp8 scaling factors + ).reshape(num_experts, hidden_size, intermediate_size // 16) # fp8 scaling factors # Pre-allocate output tensors so per-expert shuffles write directly into # contiguous slices instead of building lists + torch.stack(). This avoids diff --git a/python/sglang/srt/layers/quantization/w8a8_fp8.py b/python/sglang/srt/layers/quantization/w8a8_fp8.py index 30ed7ddb8..428594edc 100644 --- a/python/sglang/srt/layers/quantization/w8a8_fp8.py +++ b/python/sglang/srt/layers/quantization/w8a8_fp8.py @@ -101,7 +101,6 @@ class W8A8Fp8Config(QuantizationConfig): class W8A8Fp8LinearMethod(LinearMethodBase): - def __init__(self, quantization_config: W8A8Fp8Config): self.cutlass_fp8_supported = cutlass_fp8_supported() self.quantization_config = quantization_config diff --git a/python/sglang/srt/layers/quantization/w8a8_int8.py b/python/sglang/srt/layers/quantization/w8a8_int8.py index 8dc8dd2c5..482b08dc0 100644 --- a/python/sglang/srt/layers/quantization/w8a8_int8.py +++ b/python/sglang/srt/layers/quantization/w8a8_int8.py @@ -155,7 +155,6 @@ class W8A8Int8Config(QuantizationConfig): class W8A8Int8LinearMethod(LinearMethodBase): - def __init__(self, quantization_config: W8A8Int8Config): self.quantization_config = quantization_config diff --git a/python/sglang/srt/layers/rotary_embedding/base.py b/python/sglang/srt/layers/rotary_embedding/base.py index ca5c45cd7..f3161b35c 100644 --- a/python/sglang/srt/layers/rotary_embedding/base.py +++ b/python/sglang/srt/layers/rotary_embedding/base.py @@ -242,9 +242,9 @@ class RotaryEmbedding(BaseFusedOp): fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """A PyTorch-native implementation of forward().""" - assert ( - fused_set_kv_buffer_arg is None - ), "fused_set_kv_buffer_arg is not supported for native implementation" + assert fused_set_kv_buffer_arg is None, ( + "fused_set_kv_buffer_arg is not supported for native implementation" + ) if offsets is not None: positions = positions + offsets @@ -284,9 +284,9 @@ class RotaryEmbedding(BaseFusedOp): fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """A PyTorch-npu implementation of forward().""" - assert ( - fused_set_kv_buffer_arg is None - ), "fused_set_kv_buffer_arg is not supported for npu implementation" + assert fused_set_kv_buffer_arg is None, ( + "fused_set_kv_buffer_arg is not supported for npu implementation" + ) if ( query.dtype == torch.bfloat16 and self.cos_sin_cache.dtype == torch.float @@ -343,9 +343,9 @@ class RotaryEmbedding(BaseFusedOp): offsets: Optional[torch.Tensor] = None, fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: - assert ( - fused_set_kv_buffer_arg is None - ), "fused_set_kv_buffer_arg is not supported for cpu implementation" + assert fused_set_kv_buffer_arg is None, ( + "fused_set_kv_buffer_arg is not supported for cpu implementation" + ) positions = torch.add(positions, offsets) if offsets is not None else positions if _is_cpu_amx_available: @@ -386,7 +386,6 @@ class RotaryEmbedding(BaseFusedOp): fused_args=fused_set_kv_buffer_arg, ) else: - if fused_set_kv_buffer_arg is not None and _is_hip: extra_args = fused_set_kv_buffer_arg k_cache = fused_set_kv_buffer_arg["key_cache"] @@ -419,9 +418,9 @@ class RotaryEmbedding(BaseFusedOp): **extra_args, ) else: - assert ( - fused_set_kv_buffer_arg is None - ), "save kv cache is not supported for fallback_rotary_embedding." + assert fused_set_kv_buffer_arg is None, ( + "save kv cache is not supported for fallback_rotary_embedding." + ) self.cos_sin_cache = self.cos_sin_cache.to( query.device, dtype=query.dtype ) @@ -449,9 +448,9 @@ class RotaryEmbedding(BaseFusedOp): offsets: Optional[torch.Tensor] = None, fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: - assert ( - fused_set_kv_buffer_arg is None - ), "fused_set_kv_buffer_arg is not supported for xpu implementation" + assert fused_set_kv_buffer_arg is None, ( + "fused_set_kv_buffer_arg is not supported for xpu implementation" + ) positions = torch.add(positions, offsets) if offsets is not None else positions # Fused_qk_rope only supports aligned head_size diff --git a/python/sglang/srt/layers/rotary_embedding/factory.py b/python/sglang/srt/layers/rotary_embedding/factory.py index 6f2aa7e29..32a9f0ecc 100644 --- a/python/sglang/srt/layers/rotary_embedding/factory.py +++ b/python/sglang/srt/layers/rotary_embedding/factory.py @@ -433,9 +433,9 @@ def get_rope_cpu( assert rope_scaling is not None scaling_type = rope_scaling["rope_type"] - assert ( - scaling_type == "deepseek_yarn" - ), "Only deepseek_yarn is supported for CPU for now" + assert scaling_type == "deepseek_yarn", ( + "Only deepseek_yarn is supported for CPU for now" + ) scaling_factor = _get_rope_param(rope_scaling, "factor", 1.0, scaling_type) original_max_position = _get_rope_param( diff --git a/python/sglang/srt/layers/rotary_embedding/mrope.py b/python/sglang/srt/layers/rotary_embedding/mrope.py index 06fa25d6a..f60c24f4d 100644 --- a/python/sglang/srt/layers/rotary_embedding/mrope.py +++ b/python/sglang/srt/layers/rotary_embedding/mrope.py @@ -110,9 +110,9 @@ class MRotaryEmbedding(RotaryEmbedding): return None section = self.mrope_section num_pairs = self.rotary_dim // 2 - assert ( - len(section) == 3 and sum(section) == num_pairs - ), f"mrope_section {section} must be three axes summing to {num_pairs}" + assert len(section) == 3 and sum(section) == num_pairs, ( + f"mrope_section {section} must be three axes summing to {num_pairs}" + ) if self.mrope_interleaved_glm: axes = [] spent = [0, 0, 0] @@ -179,9 +179,9 @@ class MRotaryEmbedding(RotaryEmbedding): key: torch.Tensor, fused_set_kv_buffer_arg=None, ) -> Tuple[torch.Tensor, torch.Tensor]: - assert ( - fused_set_kv_buffer_arg is None - ), "save kv cache is not supported for MRotaryEmbedding." + assert fused_set_kv_buffer_arg is None, ( + "save kv cache is not supported for MRotaryEmbedding." + ) assert positions.ndim == 1 or positions.ndim == 2 cos_sin = self.cos_sin_cache[positions] @@ -282,9 +282,9 @@ class MRotaryEmbedding(RotaryEmbedding): key: torch.Tensor, fused_set_kv_buffer_arg=None, ) -> Tuple[torch.Tensor, torch.Tensor]: - assert ( - fused_set_kv_buffer_arg is None - ), "fused_set_kv_buffer_arg is not supported for npu implementation" + assert fused_set_kv_buffer_arg is None, ( + "fused_set_kv_buffer_arg is not supported for npu implementation" + ) if query.shape[1] > 4096: return self.forward_native(positions, query, key, fused_set_kv_buffer_arg) rotary_mode = "half" if self.is_neox_style else "interleave" diff --git a/python/sglang/srt/layers/rotary_embedding/rope_variant.py b/python/sglang/srt/layers/rotary_embedding/rope_variant.py index d49c76128..795e1105b 100644 --- a/python/sglang/srt/layers/rotary_embedding/rope_variant.py +++ b/python/sglang/srt/layers/rotary_embedding/rope_variant.py @@ -235,9 +235,9 @@ class FourierRotaryEmbedding(nn.Module): / self.rotary_dim ) ) - assert ( - inv_freq[:-1] > inv_freq[1:] - ).all(), "Expected inv_freq to be in decreasing order" + assert (inv_freq[:-1] > inv_freq[1:]).all(), ( + "Expected inv_freq to be in decreasing order" + ) inv_freq_idx_selected = torch.ones_like(inv_freq, dtype=torch.bool) if self.num_inv_freq is not None: inv_freq_idx_selected[self.num_inv_freq :] = False @@ -302,9 +302,9 @@ class FourierRotaryEmbedding(nn.Module): dtype=query.dtype ) cos, sin = cos_sin.chunk(2, dim=-1) - assert ( - query.dim() == key.dim() == 3 - ), "Expected query key (seq_len, heads, head_dim)" + assert query.dim() == key.dim() == 3, ( + "Expected query key (seq_len, heads, head_dim)" + ) assert cos.dim() <= 3 and sin.dim() <= 3 need_reshape = False if cos.dim() == 3: @@ -535,7 +535,6 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbedding): class Llama3RotaryEmbedding(RotaryEmbedding): - def __init__( self, head_size: int, @@ -581,7 +580,6 @@ class Llama3RotaryEmbedding(RotaryEmbedding): class Llama4VisionRotaryEmbedding(RotaryEmbedding): - def __init__( self, head_size: int, diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py index 8f10c1abf..dd9d6c5ae 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/python/sglang/srt/layers/sampler.py @@ -356,9 +356,9 @@ class Sampler(nn.Module): else: backend = get_exec().kernel.sampling_backend if backend == "flashinfer": - assert ( - sampling_info.sampling_seed is None - ), "Sampling seed is not supported for flashinfer backend" + assert sampling_info.sampling_seed is None, ( + "Sampling seed is not supported for flashinfer backend" + ) if sampling_info.need_min_p_sampling: probs = top_k_renorm_prob(probs, sampling_info.top_ks) probs = top_p_renorm_prob(probs, sampling_info.top_ps) @@ -574,9 +574,9 @@ class Sampler(nn.Module): Used for deterministic sampling with simple cases (no top-k/top-p/min-p). Requires sampling_seed to be set in sampling_info. """ - assert ( - sampling_info.sampling_seed is not None - ), "sampling_seed is required for sampling from logprobs" + assert sampling_info.sampling_seed is not None, ( + "sampling_seed is required for sampling from logprobs" + ) sampled_index = multinomial_with_seed( logprobs, sampling_info.sampling_seed, positions ) @@ -604,9 +604,9 @@ class Sampler(nn.Module): batch_next_token_ids = torch.multinomial(probs, num_samples=1).view(-1) return batch_next_token_ids.to(torch.int32) else: - assert ( - self.use_ascend_backend - ), "Only ascend backend supports sampling from logits" + assert self.use_ascend_backend, ( + "Only ascend backend supports sampling from logits" + ) batch_next_token_ids = top_k_top_p_min_p_sampling_from_logits_ascend( logits, sampling_info.top_ks, @@ -744,9 +744,9 @@ def top_k_top_p_min_p_sampling_from_probs_torch( if need_min_p_sampling: # TODO: probs_sort should be re-normalized for the use of multinomial_with_seed - assert ( - sampling_seed is None - ), "With sampling seed, multinomial_with_seed will provide wrong results" + assert sampling_seed is None, ( + "With sampling seed, multinomial_with_seed will provide wrong results" + ) min_p_thresholds = probs_sort[:, 0] * min_ps probs_sort[probs_sort < min_p_thresholds.view(-1, 1)] = 0.0 diff --git a/python/sglang/srt/layers/utils/cp_utils.py b/python/sglang/srt/layers/utils/cp_utils.py index 2a692eeb4..0171952cf 100644 --- a/python/sglang/srt/layers/utils/cp_utils.py +++ b/python/sglang/srt/layers/utils/cp_utils.py @@ -137,9 +137,9 @@ def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor): if is_dsa_prefill_cp_round_robin_split(): cp_size = get_parallel().attn_cp_size - assert ( - input_.shape[0] % cp_size == 0 - ), f"Expect input shape 0 can divided by cp size, but got input shape {input_.shape}, cp size {cp_size}" + assert input_.shape[0] % cp_size == 0, ( + f"Expect input shape 0 can divided by cp size, but got input shape {input_.shape}, cp size {cp_size}" + ) return dsa_cp_round_robin_split_data(input_) input_list = list( diff --git a/python/sglang/srt/layers/vocab_parallel_embedding.py b/python/sglang/srt/layers/vocab_parallel_embedding.py index 8576d7ef6..74dd32e82 100644 --- a/python/sglang/srt/layers/vocab_parallel_embedding.py +++ b/python/sglang/srt/layers/vocab_parallel_embedding.py @@ -273,9 +273,9 @@ class VocabParallelEmbedding(torch.nn.Module): num_added_embeddings = num_embeddings - self.org_vocab_size self.use_presharded_weights = use_presharded_weights if use_presharded_weights: - assert ( - num_added_embeddings == 0 - ), "Lora is not supported with presharded weights." + assert num_added_embeddings == 0, ( + "Lora is not supported with presharded weights." + ) self.org_vocab_size_padded = pad_vocab_size( self.org_vocab_size, self.padding_size @@ -497,7 +497,9 @@ class VocabParallelEmbedding(torch.nn.Module): assert loaded_weight.shape[output_dim] == ( self.org_vocab_size // (self.tp_size if self.use_presharded_weights else 1) - ), f"{self.org_vocab_size=} {self.use_presharded_weights=} {loaded_weight.shape[output_dim]=}" + ), ( + f"{self.org_vocab_size=} {self.use_presharded_weights=} {loaded_weight.shape[output_dim]=}" + ) # Copy the data. if not self.use_presharded_weights: diff --git a/python/sglang/srt/lora/backend/ascend_backend.py b/python/sglang/srt/lora/backend/ascend_backend.py index 4aa35dbb3..0b96fe2dd 100644 --- a/python/sglang/srt/lora/backend/ascend_backend.py +++ b/python/sglang/srt/lora/backend/ascend_backend.py @@ -24,9 +24,9 @@ class AscendLoRABackend(BaseLoRABackend): def run_lora_a_embedding( self, input_ids, weights, vocab_size, extra_embeddings=None, *args, **kwargs ): - assert ( - extra_embeddings is None - ), "Ascend LoRA embedding backend does not support extra embeddings (added tokens)." + assert extra_embeddings is None, ( + "Ascend LoRA embedding backend does not support extra embeddings (added tokens)." + ) total_seq_len = input_ids.shape[0] if weights.numel() == 0: @@ -268,9 +268,9 @@ class AscendLoRABackend(BaseLoRABackend): bs = forward_batch.batch_size if use_cuda_graph: - assert ( - self.npu_graph_batch_info is not None - ), "NPU Graph batch info is not initialized." + assert self.npu_graph_batch_info is not None, ( + "NPU Graph batch info is not initialized." + ) batch_info = self.npu_graph_batch_info batch_info.bs = forward_batch.batch_size batch_info.num_segments = forward_batch.batch_size diff --git a/python/sglang/srt/lora/backend/base_backend.py b/python/sglang/srt/lora/backend/base_backend.py index 5f70a02f5..6d3cf22f8 100644 --- a/python/sglang/srt/lora/backend/base_backend.py +++ b/python/sglang/srt/lora/backend/base_backend.py @@ -394,9 +394,9 @@ def _compute_moe_lora_info( max_len: int, ) -> tuple[torch.Tensor, torch.Tensor]: if token_lora_mapping is not None: - assert ( - num_tokens <= token_lora_mapping.shape[0] - ), "num_tokens must be less than or equal to the shape of token_lora_mapping" + assert num_tokens <= token_lora_mapping.shape[0], ( + "num_tokens must be less than or equal to the shape of token_lora_mapping" + ) token_lora_mapping = token_lora_mapping[:num_tokens] else: token_lora_mapping = torch.empty( @@ -404,9 +404,9 @@ def _compute_moe_lora_info( ) if adapter_enabled is not None: - assert ( - len(lora_ranks) <= adapter_enabled.shape[0] - ), "lora_ranks must be less than or equal to the shape of adapter_enabled" + assert len(lora_ranks) <= adapter_enabled.shape[0], ( + "lora_ranks must be less than or equal to the shape of adapter_enabled" + ) else: adapter_enabled = torch.empty( len(lora_ranks), dtype=torch.int32, device=lora_ranks.device diff --git a/python/sglang/srt/lora/backend/chunked_backend.py b/python/sglang/srt/lora/backend/chunked_backend.py index 92a77b8d7..180e90cf3 100644 --- a/python/sglang/srt/lora/backend/chunked_backend.py +++ b/python/sglang/srt/lora/backend/chunked_backend.py @@ -53,9 +53,9 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend): *args, **kwargs, ) -> torch.Tensor: - assert ( - extra_embeddings is None - ), "Extra embeddings for lora a is not supported yet in chunked backend" + assert extra_embeddings is None, ( + "Extra embeddings for lora a is not supported yet in chunked backend" + ) return chunked_embedding_lora_a_forward( input_ids=input_ids, weights=weights, diff --git a/python/sglang/srt/lora/backend/torch_backend.py b/python/sglang/srt/lora/backend/torch_backend.py index 9519c17cf..4e3ada965 100644 --- a/python/sglang/srt/lora/backend/torch_backend.py +++ b/python/sglang/srt/lora/backend/torch_backend.py @@ -51,9 +51,9 @@ class TorchNativeLoRABackend(BaseLoRABackend): *args, **kwargs, ) -> torch.Tensor: - assert ( - extra_embeddings is None - ), "Extra embeddings for lora a is not supported yet in chunked backend" + assert extra_embeddings is None, ( + "Extra embeddings for lora a is not supported yet in chunked backend" + ) output_tensor = sgemm_lora_a_embedding_fwd( inputs=input_ids, weights=weights, @@ -248,9 +248,9 @@ class TorchNativeLoRABackend(BaseLoRABackend): num_segments = len(weight_indices_tensor) if use_cuda_graph: - assert ( - self.cuda_graph_batch_info is not None - ), "CUDA Graph batch info is not initialized." + assert self.cuda_graph_batch_info is not None, ( + "CUDA Graph batch info is not initialized." + ) batch_info = self.cuda_graph_batch_info batch_info.bs = forward_batch.batch_size batch_info.num_segments = num_segments diff --git a/python/sglang/srt/lora/backend/triton_backend.py b/python/sglang/srt/lora/backend/triton_backend.py index f8f0e6d12..53a413275 100644 --- a/python/sglang/srt/lora/backend/triton_backend.py +++ b/python/sglang/srt/lora/backend/triton_backend.py @@ -287,9 +287,9 @@ class TritonLoRABackend(BaseLoRABackend): bs = forward_batch.batch_size if use_cuda_graph: - assert ( - self.cuda_graph_batch_info is not None - ), "CUDA Graph batch info is not initialized." + assert self.cuda_graph_batch_info is not None, ( + "CUDA Graph batch info is not initialized." + ) batch_info = self.cuda_graph_batch_info if forward_batch.forward_mode.is_target_verify(): # seg_lens were pre-filled at the captured per-request width diff --git a/python/sglang/srt/lora/layers.py b/python/sglang/srt/lora/layers.py index b4525fe30..173e1000e 100644 --- a/python/sglang/srt/lora/layers.py +++ b/python/sglang/srt/lora/layers.py @@ -113,9 +113,9 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA): if hasattr(base_layer, "tp_size") and base_layer.tp_size > 1: from sglang.srt.layers.communicator import get_attn_tp_context - assert ( - not get_attn_tp_context().allow_input_scattered - ), "VocabParallelEmbeddingWithLoRA with TP > 1 under input_scattered mode (e.g., DeepSeek-v2 MLA with --enable-attn-tp-input-scattered) is not fully supported and may produce incorrect results. Consider disabling input_scattered or removing embed_tokens from LoRA target modules." + assert not get_attn_tp_context().allow_input_scattered, ( + "VocabParallelEmbeddingWithLoRA with TP > 1 under input_scattered mode (e.g., DeepSeek-v2 MLA with --enable-attn-tp-input-scattered) is not fully supported and may produce incorrect results. Consider disabling input_scattered or removing embed_tokens from LoRA target modules." + ) offsets = [0, self.embed_dim] self.output_offset = torch.tensor( offsets, diff --git a/python/sglang/srt/lora/lora_drainer.py b/python/sglang/srt/lora/lora_drainer.py index d60c5787a..6af2f025e 100644 --- a/python/sglang/srt/lora/lora_drainer.py +++ b/python/sglang/srt/lora/lora_drainer.py @@ -143,14 +143,14 @@ class LoRADrainer: min_eligible_adapter = min( eligible_to_drain_adapters, - key=lambda adapter_id: self.adapter_to_stats[ - adapter_id - ].max_remaining_tokens, + key=lambda adapter_id: ( + self.adapter_to_stats[adapter_id].max_remaining_tokens + ), ) - self.adapter_to_stats[min_eligible_adapter].is_draining_for = ( - starving_adapter - ) + self.adapter_to_stats[ + min_eligible_adapter + ].is_draining_for = starving_adapter logger.debug( f"LoRA adapter {min_eligible_adapter} is draining for {starving_adapter}" ) diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py index 950ff1f68..fc94f8b2c 100644 --- a/python/sglang/srt/lora/lora_manager.py +++ b/python/sglang/srt/lora/lora_manager.py @@ -242,12 +242,12 @@ class LoRAManager: Args: lora_ref (LoRARef): The LoRARef object containing the LoRA name, path, and ID. """ - assert ( - lora_ref.lora_name is not None and lora_ref.lora_path is not None - ), "LoRARef must have both lora_name and lora_path set for loading." - assert ( - lora_ref.lora_id not in self.loras - ), f"LoRA adapter with ID {lora_ref.lora_id} is already loaded. This should have been verified before request is sent to the backend." + assert lora_ref.lora_name is not None and lora_ref.lora_path is not None, ( + "LoRARef must have both lora_name and lora_path set for loading." + ) + assert lora_ref.lora_id not in self.loras, ( + f"LoRA adapter with ID {lora_ref.lora_id} is already loaded. This should have been verified before request is sent to the backend." + ) try: # load configs @@ -337,9 +337,9 @@ class LoRAManager: adapter = self.configs.get(lora_ref.lora_id) lora_ref = self.lora_refs.get(lora_ref.lora_id) - assert ( - adapter is not None and lora_ref is not None - ), f"LoRA adapter with ID {lora_ref.lora_id} is not loaded. This should have been verified before request is sent to the backend." + assert adapter is not None and lora_ref is not None, ( + f"LoRA adapter with ID {lora_ref.lora_id} is not loaded. This should have been verified before request is sent to the backend." + ) try: pending_events = getattr(self, "pending_lora_load_events", {}) @@ -379,9 +379,9 @@ class LoRAManager: for lora_id in lora_ids: if lora_id is not None: lora_ref = self.lora_refs.get(lora_id) - assert ( - lora_ref is not None - ), f"LoRA ID {lora_id} not found in lora_refs." + assert lora_ref is not None, ( + f"LoRA ID {lora_id} not found in lora_refs." + ) pinned_loras_in_batch += int(lora_ref.pinned) assert pinned_loras_in_batch <= self.num_pinned_loras, ( @@ -570,7 +570,9 @@ class LoRAManager: assert lora_paths or ( max_lora_rank is not None and target_modules is not None - ), "When no initial --lora-paths is provided, you need to specify both --max-lora-rank and --lora-target-modules for LoRA initialization." + ), ( + "When no initial --lora-paths is provided, you need to specify both --max-lora-rank and --lora-target-modules for LoRA initialization." + ) self.init_lora_adapters(lora_paths) self.init_lora_shapes( @@ -838,12 +840,12 @@ class LoRAManager: """ Load a single LoRA adapter from tensors and config dict. """ - assert ( - lora_ref.lora_name is not None and lora_ref.lora_path is not None - ), "LoRARef must have both lora_name and lora_path set for loading." - assert ( - lora_ref.lora_id not in self.loras - ), f"LoRA adapter with ID {lora_ref.lora_id} is already loaded. This should have been verified before request is sent to the backend." + assert lora_ref.lora_name is not None and lora_ref.lora_path is not None, ( + "LoRARef must have both lora_name and lora_path set for loading." + ) + assert lora_ref.lora_id not in self.loras, ( + f"LoRA adapter with ID {lora_ref.lora_id} is already loaded. This should have been verified before request is sent to the backend." + ) try: new_adapter = LoRAConfig.from_dict( diff --git a/python/sglang/srt/lora/lora_moe_runner_marlin.py b/python/sglang/srt/lora/lora_moe_runner_marlin.py index 318e5c595..405643a0e 100644 --- a/python/sglang/srt/lora/lora_moe_runner_marlin.py +++ b/python/sglang/srt/lora/lora_moe_runner_marlin.py @@ -75,9 +75,9 @@ class MarlinLoraRunnerCore(DispatchMoeRunnerCore): topk_ids = topk_output.topk_ids assert runner_config.activation == "silu", "Only SiLU activation is supported." - assert ( - torch.cuda.get_device_capability(hidden_states.device)[0] >= 9 - ), "MarlinLoraRunnerCore requires CUDA compute capability >= 9" + assert torch.cuda.get_device_capability(hidden_states.device)[0] >= 9, ( + "MarlinLoraRunnerCore requires CUDA compute capability >= 9" + ) routed_scaling_factor = runner_config.routed_scaling_factor M, K = hidden_states.shape diff --git a/python/sglang/srt/lora/lora_registry.py b/python/sglang/srt/lora/lora_registry.py index 4e72efab5..09f0c771a 100644 --- a/python/sglang/srt/lora/lora_registry.py +++ b/python/sglang/srt/lora/lora_registry.py @@ -191,12 +191,12 @@ class LoRARegistry: This method itself is not synchronized, which is safe because it should only be called during LoRA unloading, which itself is guaranteed to be sequential. """ - assert ( - lora_id not in self._registry - ), "wait_for_unload should only be called after the LoRA adapter has been unregistered. " - assert ( - lora_id in self._counters - ), "The LoRA ID should still have a counter if it has been registered before." + assert lora_id not in self._registry, ( + "wait_for_unload should only be called after the LoRA adapter has been unregistered. " + ) + assert lora_id in self._counters, ( + "The LoRA ID should still have a counter if it has been registered before." + ) # Wait until no requests are using this LoRA adapter. await self._counters[lora_id].wait_for_zero() diff --git a/python/sglang/srt/lora/marlin_lora_temp/moe_runner.py b/python/sglang/srt/lora/marlin_lora_temp/moe_runner.py index 3cdac75f4..5e59c03f4 100644 --- a/python/sglang/srt/lora/marlin_lora_temp/moe_runner.py +++ b/python/sglang/srt/lora/marlin_lora_temp/moe_runner.py @@ -261,8 +261,7 @@ def fused_experts_experimental_sgl_marlin_lora( ep_active = local_experts < global_experts if ep_active: assert E == local_experts, ( - f"Marlin has {E} local experts but runner_config declares " - f"{local_experts}" + f"Marlin has {E} local experts but runner_config declares {local_experts}" ) assert ( lora_info.gate_up_lora_b_weights.shape[1] diff --git a/python/sglang/srt/lora/mem_pool.py b/python/sglang/srt/lora/mem_pool.py index 3668317c7..0270ef3e6 100644 --- a/python/sglang/srt/lora/mem_pool.py +++ b/python/sglang/srt/lora/mem_pool.py @@ -865,9 +865,9 @@ class LoRAMemoryPool: # to avoid contamination from the residual weight of the evicted adapters. buffer_view.zero_() else: - assert ( - buffer_view.shape == weight.shape - ), f"LoRA buffer shape {buffer_view.shape} does not match weight shape {weight.shape}." + assert buffer_view.shape == weight.shape, ( + f"LoRA buffer shape {buffer_view.shape} does not match weight shape {weight.shape}." + ) copy_weight_into_buffer(buffer_view, weight) if uid is None: @@ -1457,9 +1457,9 @@ class LoRAMemoryPool: # Non-last PP stages do not own lm_head, so adapters can # legitimately contain lm_head LoRA weights with no local # module to load them into, otherwise we should have been able to load this weight. - assert ( - not get_pp_group().is_last_rank - ), f"Failed to load lm_head LoRA weight: {name}, this is only expected to happen on non-last PP stages." + assert not get_pp_group().is_last_rank, ( + f"Failed to load lm_head LoRA weight: {name}, this is only expected to happen on non-last PP stages." + ) continue else: # Zero out embedding/lm_head buffers for adapters without embedding LoRA diff --git a/python/sglang/srt/lora/torch_ops/graph_lora_ops.py b/python/sglang/srt/lora/torch_ops/graph_lora_ops.py index 17dff6e8b..70860f292 100644 --- a/python/sglang/srt/lora/torch_ops/graph_lora_ops.py +++ b/python/sglang/srt/lora/torch_ops/graph_lora_ops.py @@ -23,7 +23,6 @@ def sgemm_lora_a_embedding_graph_fwd( ) for lora_idx in range(num_loras): - batch_token_mask = weight_indices[:total_seq_len] == lora_idx x_seq = torch.where(batch_token_mask, inputs, 0) @@ -59,7 +58,6 @@ def sgemm_lora_a_graph_fwd( ) for lora_idx in range(num_loras): - batch_token_mask = (weight_indices[:total_seq_len] == lora_idx).unsqueeze(1) x_seq = torch.where(batch_token_mask, inputs, 0) @@ -98,7 +96,6 @@ def sgemm_lora_b_graph_fwd( ) for lora_idx in range(num_loras): - batch_token_mask = (weight_indices[:total_seq_len] == lora_idx).unsqueeze(1) inputs_masked = torch.where(batch_token_mask, inputs, 0) diff --git a/python/sglang/srt/lora/torch_ops/lora_ops.py b/python/sglang/srt/lora/torch_ops/lora_ops.py index 16b25dc35..fccdd538c 100644 --- a/python/sglang/srt/lora/torch_ops/lora_ops.py +++ b/python/sglang/srt/lora/torch_ops/lora_ops.py @@ -29,7 +29,6 @@ def sgemm_lora_a_embedding_fwd( rank = lora_ranks[lora_idx] if rank > 0: - x_seq = inputs[token_offset : token_offset + seq_len] w_seq = weights[lora_idx, :rank] @@ -70,7 +69,6 @@ def sgemm_lora_a_fwd( rank = lora_ranks[lora_idx] if rank > 0: - x_seq = inputs[token_offset : token_offset + seq_len] w_seq = weights[lora_idx, : num_slices * rank] @@ -120,7 +118,6 @@ def sgemm_lora_b_fwd( rank = lora_ranks[lora_idx] if rank > 0: - for slice_idx in range(num_slices): slice_start_input = slice_idx * rank slice_end_input = (slice_idx + 1) * rank diff --git a/python/sglang/srt/lora/trtllm_lora_temp/__init__.py b/python/sglang/srt/lora/trtllm_lora_temp/__init__.py index 47f6274b3..425e33986 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/__init__.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/__init__.py @@ -148,7 +148,16 @@ def install_two_stream_overrides() -> None: :func:`get_original_row_forward`, :func:`get_original_moe_lora_func` so the new versions can fall back when their per-batch gate says single-stream. """ - global _INSTALLED, _ORIGINAL_QKV_FORWARD, _ORIGINAL_ROW_FORWARD, _ORIGINAL_MERGED_FORWARD, _ORIGINAL_COLUMN_FORWARD, _ORIGINAL_REPLICATED_FORWARD, _ORIGINAL_MOE_LORA_FUNC, _ORIGINAL_FP4_MOE_LORA_FUNC, _ORIGINAL_BF16_MOE_LORA_FUNC + global \ + _INSTALLED, \ + _ORIGINAL_QKV_FORWARD, \ + _ORIGINAL_ROW_FORWARD, \ + _ORIGINAL_MERGED_FORWARD, \ + _ORIGINAL_COLUMN_FORWARD, \ + _ORIGINAL_REPLICATED_FORWARD, \ + _ORIGINAL_MOE_LORA_FUNC, \ + _ORIGINAL_FP4_MOE_LORA_FUNC, \ + _ORIGINAL_BF16_MOE_LORA_FUNC if _INSTALLED: return diff --git a/python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py b/python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py index dc2c7a71c..972b1e4ce 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py @@ -328,9 +328,9 @@ def fused_experts_none_to_experimental_sgl_trtllm_bf16_lora( from sglang.srt.layers.moe.utils import RoutingMethodType from sglang.srt.model_executor.runner_utils.capture_mode import get_is_capture_mode - assert ( - runner_config.activation == "silu" and runner_config.is_gated - ), "experimental_sgl_trtllm BF16 LoRA currently supports the gated SwiGLU path only." + assert runner_config.activation == "silu" and runner_config.is_gated, ( + "experimental_sgl_trtllm BF16 LoRA currently supports the gated SwiGLU path only." + ) hidden_states = dispatch_output.hidden_states topk_output = dispatch_output.topk_output @@ -484,9 +484,9 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora( from sglang.srt.layers.moe.topk import TopKOutputChecker from sglang.srt.model_executor.runner_utils.capture_mode import get_is_capture_mode - assert ( - runner_config.activation == "silu" and runner_config.is_gated - ), "experimental_sgl_trtllm NVFP4 LoRA currently supports the gated SwiGLU path only." + assert runner_config.activation == "silu" and runner_config.is_gated, ( + "experimental_sgl_trtllm NVFP4 LoRA currently supports the gated SwiGLU path only." + ) hidden_states = dispatch_output.hidden_states topk_output = dispatch_output.topk_output diff --git a/python/sglang/srt/lora/trtllm_lora_temp/lora_layer.py b/python/sglang/srt/lora/trtllm_lora_temp/lora_layer.py index 1d279a2e6..40ec490ca 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/lora_layer.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/lora_layer.py @@ -140,15 +140,15 @@ def init_experimental_sgl_trtllm_lora(layer, base_layer) -> None: layer._quant_info.w2_weight = _g2.reshape(_g2.shape[0], _g2.shape[1], -1) return - assert getattr( - quant_method, "block_quant", False - ), "experimental_sgl_trtllm LoRA currently requires FP8 block quant." - assert ( - not use_mxfp8 - ), "experimental_sgl_trtllm LoRA currently targets the non-MX FP8 Qwen path." - assert ( - weight_block_size is not None - ), "experimental_sgl_trtllm LoRA needs the FP8 weight block size." + assert getattr(quant_method, "block_quant", False), ( + "experimental_sgl_trtllm LoRA currently requires FP8 block quant." + ) + assert not use_mxfp8, ( + "experimental_sgl_trtllm LoRA currently targets the non-MX FP8 Qwen path." + ) + assert weight_block_size is not None, ( + "experimental_sgl_trtllm LoRA needs the FP8 weight block size." + ) w13_weight_scale = getattr(base_layer, "w13_weight_scale_inv", None) if w13_weight_scale is None: w13_weight_scale = getattr(base_layer, "w13_weight_scale", None) diff --git a/python/sglang/srt/lora/trtllm_lora_temp/moe_overlap.py b/python/sglang/srt/lora/trtllm_lora_temp/moe_overlap.py index 2cb65e5c7..12a2d62a4 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/moe_overlap.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/moe_overlap.py @@ -373,9 +373,9 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora_two_stream( from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput from sglang.srt.layers.moe.topk import TopKOutputChecker - assert ( - runner_config.activation == "silu" and runner_config.is_gated - ), "experimental_sgl_trtllm NVFP4 LoRA currently supports the gated SwiGLU path only." + assert runner_config.activation == "silu" and runner_config.is_gated, ( + "experimental_sgl_trtllm NVFP4 LoRA currently supports the gated SwiGLU path only." + ) topk_output = dispatch_output.topk_output assert TopKOutputChecker.format_is_standard(topk_output) assert runner_config.top_k is not None @@ -623,9 +623,9 @@ def fused_experts_none_to_experimental_sgl_trtllm_bf16_lora_two_stream( from sglang.srt.layers.moe.topk import TopKOutputChecker from sglang.srt.layers.moe.utils import RoutingMethodType - assert ( - runner_config.activation == "silu" and runner_config.is_gated - ), "experimental_sgl_trtllm BF16 LoRA currently supports the gated SwiGLU path only." + assert runner_config.activation == "silu" and runner_config.is_gated, ( + "experimental_sgl_trtllm BF16 LoRA currently supports the gated SwiGLU path only." + ) topk_output = dispatch_output.topk_output assert TopKOutputChecker.format_is_standard(topk_output) assert runner_config.top_k is not None diff --git a/python/sglang/srt/lora/trtllm_lora_temp/sgl_fp8_moe.py b/python/sglang/srt/lora/trtllm_lora_temp/sgl_fp8_moe.py index 58ff08eb0..3bcc0f4a1 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/sgl_fp8_moe.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/sgl_fp8_moe.py @@ -111,9 +111,9 @@ def fused_experts_fp8_sgl( # during torch.compile for piecewise cuda graph. # Use custom op wrapper for torch.compile compatibility. if use_routed_topk: - assert ( - runner_config.top_k is not None - ), "runner_config.top_k is required for flashinfer_trtllm_routed." + assert runner_config.top_k is not None, ( + "runner_config.top_k is required for flashinfer_trtllm_routed." + ) assert TopKOutputChecker.format_is_standard(topk_output) packed_topk_ids = fused_pack_topk( topk_ids=topk_output.topk_ids, diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 556fb6cd2..cceac56e3 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -79,9 +79,7 @@ class LayerDoneCounter: def update_producer(self): self.producer_index = (self.producer_index + 1) % self.num_counters - assert self.events[ - self.producer_index - ].finish_event.query(), ( + assert self.events[self.producer_index].finish_event.query(), ( "Producer finish event should be ready before being reused." ) return self.producer_index @@ -100,7 +98,6 @@ class LayerDoneCounter: class CacheOperation: - counter = 0 def __init__( @@ -283,7 +280,6 @@ class PrefetchOperation(StorageOperation): class HiCacheController: - def __init__( self, token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator, @@ -714,9 +710,9 @@ class HiCacheController: should_split_heads = False if tp_lcm_size: - assert ( - tp_lcm_size % self.tp_size == 0 - ), "tp_lcm_size must be divisible by tp_size." + assert tp_lcm_size % self.tp_size == 0, ( + "tp_lcm_size must be divisible by tp_size." + ) should_split_heads = ( not is_rank_replicated and self.mem_pool_host.layout == "page_head" diff --git a/python/sglang/srt/managers/hisparse_coordinator.py b/python/sglang/srt/managers/hisparse_coordinator.py index 061d7545c..396755654 100644 --- a/python/sglang/srt/managers/hisparse_coordinator.py +++ b/python/sglang/srt/managers/hisparse_coordinator.py @@ -781,9 +781,9 @@ class HiSparseCoordinator: Returns: Device KV cache indices for the selected tokens. Shape: (num_reqs, top_k) """ - assert ( - not self.is_dsv4_hisparse - ), "naive_load_topk is not implemented for dsv4 hisparse" + assert not self.is_dsv4_hisparse, ( + "naive_load_topk is not implemented for dsv4 hisparse" + ) num_reqs = req_pool_indices.size(0) top_k_indices = torch.full( (num_reqs, self.top_k), -1, dtype=torch.int32, device=self.device @@ -798,9 +798,9 @@ class HiSparseCoordinator: req_idx = int(req_pool_indices[i].item()) selected_tokens = top_k_tokens[i, :top_n].to(dtype=torch.int64) - assert torch.all( - selected_tokens >= 0 - ), f"Req {req_idx}: selected tokens contain negative positions" + assert torch.all(selected_tokens >= 0), ( + f"Req {req_idx}: selected tokens contain negative positions" + ) assert torch.all(selected_tokens < seq_len), ( f"Req {req_idx}: selected tokens {selected_tokens.tolist()} " f"out of range for seq_len={seq_len}" diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py index e086a7c25..ff88f27e8 100644 --- a/python/sglang/srt/managers/multi_tokenizer_mixin.py +++ b/python/sglang/srt/managers/multi_tokenizer_mixin.py @@ -598,9 +598,9 @@ class MultiDetokenizerRouter: # Single request: route by its own http_worker_ipc. if isinstance(recv_obj, BaseReq): - assert ( - recv_obj.http_worker_ipc is not None - ), f"Single req {recv_obj.rid=} missing http_worker_ipc" + assert recv_obj.http_worker_ipc is not None, ( + f"Single req {recv_obj.rid=} missing http_worker_ipc" + ) self._send(self._pick(recv_obj.http_worker_ipc), recv_obj) continue diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index 9ac8f3d02..6996b8e77 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/python/sglang/srt/managers/overlap_utils.py @@ -125,7 +125,6 @@ CONFIDENCE_RELAY_RING_DEPTH: int = CONFIDENCE_RELAY_RING_LAG + 1 class ResolvedConfidence(msgspec.Struct): - confidence: torch.Tensor generation: torch.Tensor @@ -169,7 +168,6 @@ class RelayPayload: class ConfidenceRelay(msgspec.Struct): - device: torch.device req_pool_size: int pool: Any diff --git a/python/sglang/srt/managers/prefill_delayer.py b/python/sglang/srt/managers/prefill_delayer.py index a0bc19303..c8408e843 100644 --- a/python/sglang/srt/managers/prefill_delayer.py +++ b/python/sglang/srt/managers/prefill_delayer.py @@ -118,9 +118,9 @@ class PrefillDelayer: or envs.SGLANG_NCCL_ALL_GATHER_IN_OVERLAP_SCHEDULER_SYNC_BATCH.get() ) if use_nccl: - assert ( - device_group is not None - ), "device_group is required when using NCCL for PrefillDelayer all-gather" + assert device_group is not None, ( + "device_group is required when using NCCL for PrefillDelayer all-gather" + ) self._gather_group = device_group self._gather_device = device else: @@ -141,9 +141,9 @@ class PrefillDelayer: self._curr_state: Optional[_State] = None self.skip_first_delayer = True - assert ( - not get_schedule().disable_overlap_schedule - ), "To use PrefillDelayer, disable_overlap_schedule must be False." + assert not get_schedule().disable_overlap_schedule, ( + "To use PrefillDelayer, disable_overlap_schedule must be False." + ) def _negotiate_should_allow_prefill( self, diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index f586bc277..d8263b63b 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -2453,9 +2453,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): else: self.encoder_out_cache_loc = torch.cat(encoder_out_cache_loc) - assert ( - len(self.out_cache_loc) == self.extend_num_tokens - ), f"Expected {len(self.out_cache_loc)}, got {self.extend_num_tokens}" + assert len(self.out_cache_loc) == self.extend_num_tokens, ( + f"Expected {len(self.out_cache_loc)}, got {self.extend_num_tokens}" + ) if self.extend_input_logprob_token_ids is not None: new_token_ids_parts = [] diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index ae9fa0dc9..5aa27c38e 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -572,9 +572,7 @@ class PrefillAdder: # fail-loud `RuntimeError`. `None` outside the unified Mamba pool. self.rem_mamba_slots = None if self._mamba_slot_cost: - self.rem_mamba_slots = ( - self.token_to_kv_pool_allocator.mamba_allocator.schedulable_available_size() - ) + self.rem_mamba_slots = self.token_to_kv_pool_allocator.mamba_allocator.schedulable_available_size() if self.is_hybrid_ssm_cache: self.rem_mamba_slots += self.tree_cache.mamba_evictable_size() @@ -1316,9 +1314,9 @@ class PrefillAdder: if self.rem_dllm_tokens <= 0: return AddReqResult.OTHER - assert ( - truncation_align_size is None - ), "truncation_align_size is not supported for dllm prefill" + assert truncation_align_size is None, ( + "truncation_align_size is not supported for dllm prefill" + ) if ( tile_stop := self._check_prefill_tile_budget(input_tokens) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 17d83c74f..4d2755acf 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2321,19 +2321,31 @@ class Scheduler( spec_algorithm=self.spec_algorithm, get_running_batch=lambda: self.running_batch, get_waiting_queue=lambda: self.waiting_queue, - waiting_queue_prefix_matched=lambda: self.policy.waiting_queue_prefix_matched( - self.waiting_queue + waiting_queue_prefix_matched=lambda: ( + self.policy.waiting_queue_prefix_matched(self.waiting_queue) + ), + get_recent_cache_hit_rate=lambda: ( + self.metrics_reporter.recent_cache_hit_rate ), - get_recent_cache_hit_rate=lambda: self.metrics_reporter.recent_cache_hit_rate, get_stats=lambda: self.metrics_reporter.stats, get_chunked_req=lambda: self.chunked_req, - get_disagg_prefill_bootstrap_queue=lambda: self.disagg_prefill_bootstrap_queue, - get_disagg_prefill_inflight_queue=lambda: self.disagg_prefill_inflight_queue, + get_disagg_prefill_bootstrap_queue=lambda: ( + self.disagg_prefill_bootstrap_queue + ), + get_disagg_prefill_inflight_queue=lambda: ( + self.disagg_prefill_inflight_queue + ), get_disagg_decode_prealloc_queue=lambda: self.disagg_decode_prealloc_queue, get_disagg_decode_transfer_queue=lambda: self.disagg_decode_transfer_queue, - get_spec_total_num_accept_tokens=lambda: self.metrics_reporter.spec_total_num_accept_tokens, - get_spec_total_num_forward_ct=lambda: self.metrics_reporter.spec_total_num_forward_ct, - get_total_prefill_uncached_tokens=lambda: self.total_prefill_uncached_tokens, + get_spec_total_num_accept_tokens=lambda: ( + self.metrics_reporter.spec_total_num_accept_tokens + ), + get_spec_total_num_forward_ct=lambda: ( + self.metrics_reporter.spec_total_num_forward_ct + ), + get_total_prefill_uncached_tokens=lambda: ( + self.total_prefill_uncached_tokens + ), get_total_prefill_busy_us=lambda: self.total_prefill_busy_us, get_decode_moment_totals=lambda: self.decode_moment_totals, ) @@ -3688,9 +3700,8 @@ class Scheduler( running_batch.batch_is_full = True if running_batch.batch_is_full: - if ( - not self.enable_priority_preemption - or not adder.preempt_to_schedule(req) + if not self.enable_priority_preemption or not adder.preempt_to_schedule( + req ): break diff --git a/python/sglang/srt/managers/scheduler_components/new_token_ratio_tracker.py b/python/sglang/srt/managers/scheduler_components/new_token_ratio_tracker.py index 8a771a5b0..fbc1dee21 100644 --- a/python/sglang/srt/managers/scheduler_components/new_token_ratio_tracker.py +++ b/python/sglang/srt/managers/scheduler_components/new_token_ratio_tracker.py @@ -44,8 +44,6 @@ class NewTokenRatioTracker: new_estimate_ratio = ( total_decoded_tokens + envs.SGLANG_RETRACT_DECODE_STEPS.get() * len(reqs) - ) / ( - total_max_new_tokens + 1 - ) # avoid zero division + ) / (total_max_new_tokens + 1) # avoid zero division new_estimate_ratio = min(1.0, new_estimate_ratio) return new_estimate_ratio diff --git a/python/sglang/srt/managers/scheduler_components/weight_updater.py b/python/sglang/srt/managers/scheduler_components/weight_updater.py index 362de6889..925ee721b 100644 --- a/python/sglang/srt/managers/scheduler_components/weight_updater.py +++ b/python/sglang/srt/managers/scheduler_components/weight_updater.py @@ -209,9 +209,9 @@ class SchedulerWeightUpdaterManager: ) def release_memory_occupation(self, recv_req: ReleaseMemoryOccupationReqInput): - assert ( - self.is_fully_idle() - ), "release_memory_occupation should be called only when server is idle." + assert self.is_fully_idle(), ( + "release_memory_occupation should be called only when server is idle." + ) tags = recv_req.tags @@ -338,9 +338,9 @@ class SchedulerWeightUpdaterManager: if self.draft_worker is not None: draft_url = params.get("draft_url", None) - assert ( - draft_url is not None - ), "draft_url must be provided when draft model is enabled" + assert draft_url is not None, ( + "draft_url must be provided when draft model is enabled" + ) self.draft_worker.model_runner.weight_exporter.save_remote_model(draft_url) def save_sharded_model(self, params): diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py index f02298037..857444388 100644 --- a/python/sglang/srt/managers/tokenizer_control_mixin.py +++ b/python/sglang/srt/managers/tokenizer_control_mixin.py @@ -427,9 +427,9 @@ class TokenizerControlMixin: request: Optional[fastapi.Request] = None, ) -> Tuple[bool, str]: self.auto_create_handle_loop() - assert ( - get_parallel().dp_size == 1 or get_parallel().enable_dp_attention - ), "dp_size must be 1 or dp attention must be enabled for update weights from distributed" + assert get_parallel().dp_size == 1 or get_parallel().enable_dp_attention, ( + "dp_size must be 1 or dp attention must be enabled for update weights from distributed" + ) results = await self.init_weights_update_group_communicator(obj) return FanOutCommunicator.merge_results(results) @@ -440,9 +440,9 @@ class TokenizerControlMixin: request: Optional[fastapi.Request] = None, ) -> Tuple[bool, str]: self.auto_create_handle_loop() - assert ( - get_parallel().dp_size == 1 or get_parallel().enable_dp_attention - ), "dp_size must be 1 or dp attention must be enabled for destroy parameter update group" + assert get_parallel().dp_size == 1 or get_parallel().enable_dp_attention, ( + "dp_size must be 1 or dp attention must be enabled for destroy parameter update group" + ) results = await self.destroy_weights_update_group_communicator(obj) return FanOutCommunicator.merge_results(results) @@ -453,9 +453,9 @@ class TokenizerControlMixin: request: Optional[fastapi.Request] = None, ) -> Tuple[bool, str]: self.auto_create_handle_loop() - assert ( - get_parallel().dp_size == 1 or get_parallel().enable_dp_attention - ), "dp_size must be 1 or dp attention must be enabled for update weights from distributed" + assert get_parallel().dp_size == 1 or get_parallel().enable_dp_attention, ( + "dp_size must be 1 or dp attention must be enabled for update weights from distributed" + ) if obj.abort_all_requests: self.abort_request(abort_all=True) @@ -486,9 +486,9 @@ class TokenizerControlMixin: ) -> Tuple[bool, str]: self.auto_create_handle_loop() # TODO: support DP - assert ( - get_parallel().dp_size == 1 - ), "dp_size must be 1 for init_weights_send_group_for_remote_instance" + assert get_parallel().dp_size == 1, ( + "dp_size must be 1 for init_weights_send_group_for_remote_instance" + ) result = ( await self.init_weights_send_group_for_remote_instance_communicator(obj) )[0] @@ -501,9 +501,9 @@ class TokenizerControlMixin: ) -> Tuple[bool, str]: self.auto_create_handle_loop() # TODO: support DP - assert ( - get_parallel().dp_size == 1 - ), "dp_size must be 1 for send_weights_to_remote_instance" + assert get_parallel().dp_size == 1, ( + "dp_size must be 1 for send_weights_to_remote_instance" + ) result = (await self.send_weights_to_remote_instance_communicator(obj))[0] return result.success, result.message @@ -513,9 +513,9 @@ class TokenizerControlMixin: request: Optional[fastapi.Request] = None, ) -> Tuple[bool, str]: self.auto_create_handle_loop() - assert ( - get_parallel().dp_size == 1 or get_parallel().enable_dp_attention - ), "dp_size must be 1 or dp attention must be enabled for update weights from tensor" + assert get_parallel().dp_size == 1 or get_parallel().enable_dp_attention, ( + "dp_size must be 1 or dp attention must be enabled for update weights from tensor" + ) if obj.abort_all_requests: self.abort_request(abort_all=True) @@ -551,9 +551,9 @@ class TokenizerControlMixin: self.auto_create_handle_loop() try: # For now, we only support single data parallel instance - assert ( - get_parallel().dp_size == 1 or get_parallel().enable_dp_attention - ), "dp_size must be 1 or dp attention must be enabled for update weights from IPC" + assert get_parallel().dp_size == 1 or get_parallel().enable_dp_attention, ( + "dp_size must be 1 or dp attention must be enabled for update weights from IPC" + ) logger.info("Starting IPC weight update") async with self.is_pause_cond: @@ -583,9 +583,9 @@ class TokenizerControlMixin: self: TokenizerManager, obj: UnloadLoRAAdapterReqInput, ) -> UnloadLoRAAdapterReqOutput: - assert ( - self.lora_update_lock.locked() - ), "self.lora_update_lock must be locked in order for self._unload_lora_adapter_locked() to be called" + assert self.lora_update_lock.locked(), ( + "self.lora_update_lock must be locked in order for self._unload_lora_adapter_locked() to be called" + ) # Unregister the LoRA adapter from the registry to stop new requests for this adapter # from being started. @@ -614,9 +614,9 @@ class TokenizerControlMixin: "LoRA is not enabled. Please set `--enable-lora` to enable LoRA." ) - assert ( - get_parallel().dp_size == 1 or get_parallel().enable_dp_attention - ), "dp_size must be 1 or dp attention must be enabled for dynamic lora loading" + assert get_parallel().dp_size == 1 or get_parallel().enable_dp_attention, ( + "dp_size must be 1 or dp attention must be enabled for dynamic lora loading" + ) logger.info( "Start load Lora adapter. Lora name=%s, path=%s", obj.lora_name, @@ -692,9 +692,9 @@ class TokenizerControlMixin: "LoRA is not enabled. Please set `--enable-lora` to enable LoRA." ) - assert ( - get_parallel().dp_size == 1 or get_parallel().enable_dp_attention - ), "dp_size must be 1 or dp attention must be enabled for dynamic lora loading" + assert get_parallel().dp_size == 1 or get_parallel().enable_dp_attention, ( + "dp_size must be 1 or dp attention must be enabled for dynamic lora loading" + ) logger.info( "Start load Lora adapter from tensors. Lora name=%s", obj.lora_name, @@ -768,13 +768,13 @@ class TokenizerControlMixin: "LoRA is not enabled. Please set `--enable-lora` to enable LoRA." ) - assert ( - obj.lora_name is not None - ), "lora_name must be provided to unload LoRA adapter" + assert obj.lora_name is not None, ( + "lora_name must be provided to unload LoRA adapter" + ) - assert ( - get_parallel().dp_size == 1 or get_parallel().enable_dp_attention - ), "dp_size must be 1 or dp attention must be enabled for dynamic lora loading" + assert get_parallel().dp_size == 1 or get_parallel().enable_dp_attention, ( + "dp_size must be 1 or dp attention must be enabled for dynamic lora loading" + ) logger.info( "Start unload Lora adapter. Lora name=%s", obj.lora_name, @@ -847,9 +847,9 @@ class TokenizerControlMixin: async def get_internal_state(self: TokenizerManager) -> List[Dict[Any, Any]]: self.auto_create_handle_loop() req = GetInternalStateReq() - responses: List[GetInternalStateReqOutput] = ( - await self.get_internal_state_communicator(req) - ) + responses: List[ + GetInternalStateReqOutput + ] = await self.get_internal_state_communicator(req) # Many DP ranks return [res.internal_state for res in responses] @@ -857,9 +857,9 @@ class TokenizerControlMixin: self: TokenizerManager, obj: SetInternalStateReq ) -> List[bool]: self.auto_create_handle_loop() - responses: List[SetInternalStateReqOutput] = ( - await self.set_internal_state_communicator(obj) - ) + responses: List[ + SetInternalStateReqOutput + ] = await self.set_internal_state_communicator(obj) return [res.updated for res in responses] async def dumper_control( diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 2d61e1417..c5a5e24cc 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -1768,9 +1768,9 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): # Record response sent time right before we log finished results and metrics. if not state.time_stats.response_sent_to_client_time: state.time_stats.set_response_sent_to_client_time() - out["meta_info"][ - "response_sent_to_client_ts" - ] = state.time_stats.get_response_sent_to_client_realtime() + out["meta_info"]["response_sent_to_client_ts"] = ( + state.time_stats.get_response_sent_to_client_realtime() + ) self.request_logger.log_finished_request( obj, out, @@ -1798,9 +1798,9 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): # Record response sent time right before we send response. if not state.time_stats.response_sent_to_client_time: state.time_stats.set_response_sent_to_client_time() - out["meta_info"][ - "response_sent_to_client_ts" - ] = state.time_stats.get_response_sent_to_client_realtime() + out["meta_info"]["response_sent_to_client_ts"] = ( + state.time_stats.get_response_sent_to_client_realtime() + ) yield out else: if ( @@ -2027,9 +2027,11 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): self.model_update_lock.writer_lock if not is_paused else nullcontext() ) async with lock_context: - success, message, num_paused_requests = ( - await self._wait_for_model_update_from_disk(obj) - ) + ( + success, + message, + num_paused_requests, + ) = await self._wait_for_model_update_from_disk(obj) if success and obj.flush_cache and self.mm_processor is not None: self.mm_processor.clear_preprocess_cache() @@ -2597,8 +2599,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): # shared batch-output loop; degrade to nested instead. state.input_top_logprobs_flat_fields = None logger.error( - "Falling back to nested input top logprobs for " - "rid=%s: %s", + "Falling back to nested input top logprobs for rid=%s: %s", meta_info.get("id"), e, ) @@ -3084,7 +3085,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): filename = os.path.join( self.crash_dump_folder, hostname, - f'crash_dump_{datetime.now().strftime("%Y-%m-%d_%H-%M-%S")}.pkl', + f"crash_dump_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.pkl", ) os.makedirs(os.path.dirname(filename), exist_ok=True) @@ -3301,9 +3302,9 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): scale_phase=self.elastic_scale_phase, ) self.auto_create_handle_loop() - responses: List[ScaleElasticEPReqOutput] = ( - await self.scale_elastic_ep_communicator(obj) - ) + responses: List[ + ScaleElasticEPReqOutput + ] = await self.scale_elastic_ep_communicator(obj) for res in responses: if not res.success: self.elastic_scale_phase = res.scale_phase diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index 328dd09c7..816b7d2d3 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -614,9 +614,9 @@ class TpModelWorker(BaseTpWorker): else: # FIXME(lsyin): unify the interface of forward_batch assert forward_batch is not None - assert ( - capture_hidden_mode is None - ), "capture_hidden_mode override requires a ScheduleBatch input" + assert capture_hidden_mode is None, ( + "capture_hidden_mode override requires a ScheduleBatch input" + ) # Deprecated kwarg: pre-planners mark the batch themselves now. forward_batch.apply_deprecated_skip_attn_backend_init(skip_attn_backend_init) diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index aece0c5e2..ee08f42bf 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py @@ -386,7 +386,7 @@ def msgpack_decode_explained(data: bytes) -> Any: if m is not None: idx = int(m.group(1)) if 1 <= idx <= len(fields): - msg = f"{msg[:m.start()]}$.{fields[idx - 1]}{msg[m.end():]}" + msg = f"{msg[: m.start()]}$.{fields[idx - 1]}{msg[m.end() :]}" raise MsgpackDecodeError(rid, msg) from e diff --git a/python/sglang/srt/mem_cache/allocator/hisparse.py b/python/sglang/srt/mem_cache/allocator/hisparse.py index 4a10891f8..4be5df84e 100644 --- a/python/sglang/srt/mem_cache/allocator/hisparse.py +++ b/python/sglang/srt/mem_cache/allocator/hisparse.py @@ -158,9 +158,9 @@ class HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): extra_indices = self.hisparse_attn_allocator.alloc( need_size - len(hisparse_indices) ) - assert ( - extra_indices is not None - ), "Hisparse allocation failed in alloc_device_buffer" + assert extra_indices is not None, ( + "Hisparse allocation failed in alloc_device_buffer" + ) buffer_indices = torch.cat([hisparse_indices, extra_indices]) return buffer_indices @@ -216,9 +216,9 @@ class HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): len(logical_indices), num_new_pages=num_new_pages, ) - assert ( - hisparse_indices is not None - ), "Hisparse allocation failed in alloc_extend" + assert hisparse_indices is not None, ( + "Hisparse allocation failed in alloc_extend" + ) self.full_to_hisparse_device_index_mapping[logical_indices] = hisparse_indices return logical_indices @@ -268,7 +268,6 @@ class HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): - def __init__( self, logical_attn_allocator: BaseTokenToKVPoolAllocator, @@ -469,9 +468,9 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): extra_indices = self.hisparse_attn_allocator.alloc( need_size - len(hisparse_indices) ) - assert ( - extra_indices is not None - ), "Hisparse allocation failed in alloc_device_buffer" + assert extra_indices is not None, ( + "Hisparse allocation failed in alloc_device_buffer" + ) buffer_indices = torch.cat([hisparse_indices, extra_indices]) return buffer_indices @@ -538,9 +537,9 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): hisparse_last_loc, len(compressed_logical_indices), ) - assert ( - hisparse_indices is not None - ), "Hisparse allocation failed in alloc_extend" + assert hisparse_indices is not None, ( + "Hisparse allocation failed in alloc_extend" + ) self.full_to_hisparse_device_index_mapping[compressed_logical_indices] = ( hisparse_indices.to(torch.int64) diff --git a/python/sglang/srt/mem_cache/allocator/paged.py b/python/sglang/srt/mem_cache/allocator/paged.py index 8ad1d2329..b50272d26 100755 --- a/python/sglang/srt/mem_cache/allocator/paged.py +++ b/python/sglang/srt/mem_cache/allocator/paged.py @@ -162,9 +162,9 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): def alloc(self, need_size: int): # page-aligned allocation, returning contiguous indices of pages if self.debug_mode: - assert ( - need_size % self.page_size == 0 - ), "The allocation size should be page-aligned" + assert need_size % self.page_size == 0, ( + "The allocation size should be page-aligned" + ) num_pages = need_size // self.page_size if num_pages > len(self.free_pages): diff --git a/python/sglang/srt/mem_cache/chunk_cache.py b/python/sglang/srt/mem_cache/chunk_cache.py index c524a974e..ab3bc4afb 100644 --- a/python/sglang/srt/mem_cache/chunk_cache.py +++ b/python/sglang/srt/mem_cache/chunk_cache.py @@ -127,9 +127,9 @@ class SWAChunkCache(ChunkCache): self.chunked_prefill_size = params.chunked_prefill_size def supports_swa(self) -> bool: - assert ( - self.sliding_window_size is not None - ), "sliding_window_size must be set for SWAChunkCache" + assert self.sliding_window_size is not None, ( + "sliding_window_size must be set for SWAChunkCache" + ) return True def evict(self, params: EvictParams) -> EvictResult: diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index b1497c6f5..43cd1b15f 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -66,9 +66,9 @@ def free_swa_out_of_window_slots( return # For swa radix cache, we need to evict the tokens that are not in the tree cache and also not in the sliding window - assert ( - req.kv.cache_protected_len % page_size == 0 - ), "cache_protected_len must be page aligned" + assert req.kv.cache_protected_len % page_size == 0, ( + "cache_protected_len must be page aligned" + ) req.kv.swa_evicted_seqlen = max( req.kv.swa_evicted_seqlen, req.kv.swa_dead_lo(page_size) ) @@ -255,9 +255,9 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr assert (not req.kv.holds_kv) == req.kv.is_kv_released # MambaRadixCache may alloc mamba state before alloc KV cache if not req.kv.holds_kv: - assert ( - tree_cache.supports_mamba() - ), "Only MambaRadixCache allow freeing before alloc" + assert tree_cache.supports_mamba(), ( + "Only MambaRadixCache allow freeing before alloc" + ) # TODO (csy, hanming): clean up this early allocation logic if req.kv.holds_mamba: tree_cache.req_to_token_pool.mamba_allocator.free( @@ -286,9 +286,9 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr if isinstance(tree_cache.req_to_token_pool, HybridReqToTokenPool) and ( not tree_cache.supports_mamba() ): - assert ( - req.kv.holds_mamba - ), "mamba state is freed while the tree cache does not manage mamba states" + assert req.kv.holds_mamba, ( + "mamba state is freed while the tree cache does not manage mamba states" + ) tree_cache.req_to_token_pool.free_mamba_cache(req) # The DSV4-NPU ReqToTokenPool subclass's free() additionally releases the # c4/c128 state pages; other ReqToTokenPool subclasses are a no-op here. @@ -306,9 +306,9 @@ def _release_overallocated_kv_indices( # strip_thinking_cache intentionally reports output tokens as overallocated # so they fall into the free path below (#22373). if spec_algo is None and not get_serving().strip_thinking_cache: - assert ( - start_p == end_p - ), f"Unexpected overallocated KV cache, {req.kv.kv_committed_len=}, {req.kv.kv_allocated_len=}" + assert start_p == end_p, ( + f"Unexpected overallocated KV cache, {req.kv.kv_committed_len=}, {req.kv.kv_allocated_len=}" + ) if page_size > 1: start_p = ceil_align(start_p, page_size) diff --git a/python/sglang/srt/mem_cache/deepseek_v4_compress_state.py b/python/sglang/srt/mem_cache/deepseek_v4_compress_state.py index 4c5dd567b..915f0b0a7 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_compress_state.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_compress_state.py @@ -73,9 +73,9 @@ class KVAndScore: assert len(tensors) > 0, "At least one tensor is required for concatenation." item_size = tensors[0]._item_size for v in tensors: - assert ( - v._item_size == item_size - ), "All tensors must have the same item size." + assert v._item_size == item_size, ( + "All tensors must have the same item size." + ) return KVAndScore(torch.cat([v.kv_score for v in tensors], dim=dim)) diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index ec2adad9f..8287abbf8 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -182,7 +182,6 @@ class DeepSeekV4SingleKVPool(KVCache): class HiSparseC4DevicePool(DeepSeekV4SingleKVPool): - def __init__( self, size: int, @@ -509,7 +508,6 @@ class DeepSeekV4UnifiedKVPool: class DeepSeekV4TokenToKVPool(BaseSWAKVPool): - def __init__( self, max_num_reqs: int, @@ -1028,9 +1026,9 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): def get_attention_compress_states(self, layer_id: int) -> CompressStatePool: self.wait_layer_transfer(layer_id) compress_state_pool = self.compress_state_pools[layer_id] - assert ( - compress_state_pool is not None - ), "Only c4/c128 layers have attention states." + assert compress_state_pool is not None, ( + "Only c4/c128 layers have attention states." + ) return compress_state_pool def get_online_c128_mtp_state_slot_offset(self) -> int: @@ -1099,9 +1097,9 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): def get_indexer_compress_states(self, layer_id: int) -> CompressStatePool: self.wait_layer_transfer(layer_id) indexer_compress_state_pool = self.indexer_compress_state_pools[layer_id] - assert ( - indexer_compress_state_pool is not None - ), "Only c4 layers have indexer states." + assert indexer_compress_state_pool is not None, ( + "Only c4 layers have indexer states." + ) return indexer_compress_state_pool def _swa_local_layer_id(self, layer_id: int) -> int: diff --git a/python/sglang/srt/mem_cache/dsa_cache_layer_split.py b/python/sglang/srt/mem_cache/dsa_cache_layer_split.py index ba1cd1e52..a4861f55a 100644 --- a/python/sglang/srt/mem_cache/dsa_cache_layer_split.py +++ b/python/sglang/srt/mem_cache/dsa_cache_layer_split.py @@ -201,9 +201,9 @@ class LayerSplitDSATokenToKVPool(DSATokenToKVPool): layer_shard_size: int, **kwargs, ): - assert ( - layer_shard_rank is not None and layer_shard_size > 1 - ), "LayerSplitDSATokenToKVPool requires layer_shard_size > 1" + assert layer_shard_rank is not None and layer_shard_size > 1, ( + "LayerSplitDSATokenToKVPool requires layer_shard_size > 1" + ) self.layer_shard_rank = layer_shard_rank self.layer_shard_size = layer_shard_size self.layer_shard_enabled = True diff --git a/python/sglang/srt/mem_cache/embedding_cache_controller.py b/python/sglang/srt/mem_cache/embedding_cache_controller.py index be912813b..cf5d1fee1 100644 --- a/python/sglang/srt/mem_cache/embedding_cache_controller.py +++ b/python/sglang/srt/mem_cache/embedding_cache_controller.py @@ -583,7 +583,7 @@ class EmbeddingCacheController: self._lru_touch(mm_hash) else: logger.debug( - f"Req {req_id}: {mm_hash} is FILLING; " f"treating as miss." + f"Req {req_id}: {mm_hash} is FILLING; treating as miss." ) continue diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index e8d403f07..d0402364d 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -371,7 +371,6 @@ class MetadataCache: class HiCacheFile(HiCacheStorage): - def __init__( self, storage_config: HiCacheStorageConfig, file_path: str = "/tmp/hicache" ): diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 561599d35..396dd018f 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -80,7 +80,6 @@ logger = logging.getLogger(__name__) class HiRadixCache(RadixCache): - def __init__(self, params: CacheInitParams, server_args: ServerArgs): self._enable_metrics_flag = params.enable_metrics @@ -105,7 +104,6 @@ class HiRadixCache(RadixCache): # Filled by attach_hybrid_minimax_sparse_pool_to_hiradix_cache. self.token_to_kv_pool_host = None elif isinstance(self.kv_cache, MLATokenToKVPool): - _parallel = get_parallel() self.token_to_kv_pool_host = MLATokenToKVPoolHost( self.kv_cache, @@ -342,7 +340,6 @@ class HiRadixCache(RadixCache): labels.update(extra_metric_labels) existing_collector = getattr(self, "storage_metrics_collector", None) if existing_collector is None: - storage_cls = resolve_collector_class( STAT_LOGGER_ROLE_STORAGE, StorageMetricsCollector, @@ -1184,9 +1181,9 @@ class HiRadixCache(RadixCache): self._update_leaf_status(node) self._update_host_leaf_status(node) if node.parent is None: - assert ( - node is self.root_node - ), f"This request holds the node from another tree" + assert node is self.root_node, ( + f"This request holds the node from another tree" + ) node = node.parent return DecLockRefResult(delta=delta) @@ -1397,9 +1394,9 @@ class HiRadixCache(RadixCache): last_hit_node = node nodes_to_load = [] while node.evicted: - assert ( - node.backuped - ), "No backup available on evicted nodes, should not happen" + assert node.backuped, ( + "No backup available on evicted nodes, should not happen" + ) nodes_to_load.insert(0, node) node = node.parent else: diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py index 18acefeb8..e0a806670 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/python/sglang/srt/mem_cache/kv_cache_builder.py @@ -51,7 +51,6 @@ from sglang.srt.runtime_context import ( from sglang.srt.utils import is_hip if TYPE_CHECKING: - from torch.distributed import ProcessGroup from sglang.srt.configs.model_config import ModelConfig diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 81f84e68e..3cb24e143 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -314,9 +314,9 @@ class KVCacheConfigurator: ) if not self.spec_algorithm.is_none() and self.is_draft_worker: - assert ( - self.memory_pool_config is not None - ), "Draft worker requires memory_pool_config" + assert self.memory_pool_config is not None, ( + "Draft worker requires memory_pool_config" + ) config = self.memory_pool_config else: config = self._resolve_memory_pool_config(pre_model_load_memory) @@ -713,9 +713,9 @@ class KVCacheConfigurator: config = self.mambaish_config assert config is not None and self.is_hybrid_swa assert self.page_size >= 1, f"page_size must be >= 1, got {self.page_size}" - assert ( - not self.use_mla_backend - ), "unified tri-pool does not support an MLA full side yet" + assert not self.use_mla_backend, ( + "unified tri-pool does not support an MLA full side yet" + ) # Mirror the non-shared path's extra_max_context_len computation. extra_max_context_len = 4 if get_spec().speculative_num_draft_tokens is not None: @@ -827,9 +827,9 @@ class KVCacheConfigurator: # Both sub-pools are page-aware; the SWA composite runs alloc_extend_kernel # once in virtual space and binds the new pages on both sub-allocators. assert self.page_size >= 1, f"page_size must be >= 1, got {self.page_size}" - assert ( - not self.use_mla_backend - ), "unified memory pool does not support MLA-SWA hybrid yet" + assert not self.use_mla_backend, ( + "unified memory pool does not support MLA-SWA hybrid yet" + ) # Mirror the non-shared path's extra_max_context_len computation. extra_max_context_len = 4 if get_spec().speculative_num_draft_tokens is not None: @@ -1252,9 +1252,9 @@ class KVCacheConfigurator: if quant_method is not None and is_float4_e2m1fn_x2( self.kv_cache_dtype ): - assert ( - not enable_page_major - ), "page-major KV layout is not supported with fp4 KV cache" + assert not enable_page_major, ( + "page-major KV layout is not supported with fp4 KV cache" + ) token_to_kv_pool = self._build_mha_kv_pool( max_total_num_tokens=sizes.max_total_num_tokens, mha_pool_class=mha_pool_class, @@ -2080,9 +2080,9 @@ class KVCacheConfigurator: else: additional_ratio = MAMBA_CACHE_V2_ADDITIONAL_RATIO_OVERLAP else: - assert ( - not mamba_extra_buffer_lazy_enabled() - ), "Lazy extra buffer requires overlap schedule (--disable-overlap-schedule is incompatible)" + assert not mamba_extra_buffer_lazy_enabled(), ( + "Lazy extra buffer requires overlap schedule (--disable-overlap-schedule is incompatible)" + ) additional_ratio = MAMBA_CACHE_V2_ADDITIONAL_RATIO_NO_OVERLAP elif skip_decode_lock: # no_buffer under skip: add the base drop back so effective stays 3, @@ -2433,9 +2433,9 @@ def calculate_mla_kv_cache_dim( # kv_lora_rank + scale storage (kv_lora_rank // quant_block_size * 4 bytes) + rope dimension storage # Note: rope dimension is stored in original dtype (bf16), not quantized to fp8 if kv_cache_dtype == torch.float8_e4m3fn: - assert ( - kv_lora_rank % quant_block_size == 0 - ), f"kv_lora_rank {kv_lora_rank} must be multiple of quant_block_size {quant_block_size}" + assert kv_lora_rank % quant_block_size == 0, ( + f"kv_lora_rank {kv_lora_rank} must be multiple of quant_block_size {quant_block_size}" + ) return ( kv_lora_rank diff --git a/python/sglang/srt/mem_cache/kv_index_translator.py b/python/sglang/srt/mem_cache/kv_index_translator.py index 7701c95dd..780456eb7 100644 --- a/python/sglang/srt/mem_cache/kv_index_translator.py +++ b/python/sglang/srt/mem_cache/kv_index_translator.py @@ -227,8 +227,7 @@ class KVIndexTranslator: if sliding_window: assert self._swa_v2p_table is not None, ( - "fill_packed_read_stream: sliding_window on a pool with no swa " - "sub-pool" + "fill_packed_read_stream: sliding_window on a pool with no swa sub-pool" ) build_kv_read_table_packed( req_to_token=self.req_to_token, diff --git a/python/sglang/srt/mem_cache/kv_vmm_backing.py b/python/sglang/srt/mem_cache/kv_vmm_backing.py index fd69ea333..88f626ece 100644 --- a/python/sglang/srt/mem_cache/kv_vmm_backing.py +++ b/python/sglang/srt/mem_cache/kv_vmm_backing.py @@ -212,9 +212,9 @@ class KvVmmBufferOwner: self.ensure_prefix(self.page_size) for t in self.tensors: - assert ( - t.is_cuda and t.device.index == self.device_id - ), f"post-capture KV buffer landed on {t.device}, expected cuda:{self.device_id}" + assert t.is_cuda and t.device.index == self.device_id, ( + f"post-capture KV buffer landed on {t.device}, expected cuda:{self.device_id}" + ) # -- backing -------------------------------------------------------------- diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index 3a9d46a85..0847748d0 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -71,7 +71,6 @@ _MAMBA_DEBUG_ASSERTS = os.environ.get("SGLANG_MAMBA_DEBUG_ASSERTS", "0") == "1" class TreeNode: - counter = 0 last_access_time_counter_float = float64(1.0) @@ -238,9 +237,9 @@ class LRUList: Move a (existing) node to most recently used position """ assert node.id in self.cache, f"Resetting node {node.id=} not in lru list" - assert ( - not self.mamba or node.mamba_value is not None - ), f"Resetting mamba tombstone node in mamba lru list: {node.id=}" + assert not self.mamba or node.mamba_value is not None, ( + f"Resetting mamba tombstone node in mamba lru list: {node.id=}" + ) if self.mamba: node.mamba_last_access_time = get_last_access_time() self._remove_node(node) @@ -254,9 +253,9 @@ class LRUList: prev_node = self.head while node != root_node: if not self.mamba or node.mamba_value is not None: - assert ( - node.id in self.cache - ), f"Resetting node {node.id=} not in lru list when resetting node and parents mru" + assert node.id in self.cache, ( + f"Resetting node {node.id=} not in lru list when resetting node and parents mru" + ) self._remove_node(node) self._add_node_after(prev_node, node) prev_node = node @@ -266,12 +265,12 @@ class LRUList: """ Insert a (new) node as most recently used """ - assert ( - not self.mamba or node.mamba_value is not None - ), f"Inserting mamba tombstone node in mamba lru list: {node.id=}" - assert ( - node.id not in self.cache - ), f"Inserting node {node.id=} already in lru list, existing node: {self.cache[node.id].id=}" + assert not self.mamba or node.mamba_value is not None, ( + f"Inserting mamba tombstone node in mamba lru list: {node.id=}" + ) + assert node.id not in self.cache, ( + f"Inserting node {node.id=} already in lru list, existing node: {self.cache[node.id].id=}" + ) if self.mamba: node.mamba_last_access_time = get_last_access_time() self.cache[node.id] = node @@ -282,9 +281,9 @@ class LRUList: Remove node from lru list """ assert node.id in self.cache, f"Removing node {node.id=} not in lru list" - assert ( - not self.mamba or node.mamba_value is not None - ), f"Removing mamba tombstone node from mamba lru list: {node.id=}" + assert not self.mamba or node.mamba_value is not None, ( + f"Removing mamba tombstone node from mamba lru list: {node.id=}" + ) del self.cache[node.id] self._remove_node(node) @@ -307,9 +306,9 @@ class LRUList: Get the previous (i.e. more recently used) node that is not locked """ if check_id: - assert ( - node.id in self.cache - ), f"Getting prev of node {node.id=} not in lru list" + assert node.id in self.cache, ( + f"Getting prev of node {node.id=} not in lru list" + ) x = getattr(node, self.prv) # x = node.prev while getattr(x, self.lock_ref) > 0: x = getattr(x, self.prv) # x = x.prev @@ -323,9 +322,9 @@ class LRUList: Get the previous (i.e. more recently used) leaf node that is not locked """ if check_id: - assert ( - node.id in self.cache - ), f"Getting prev of node {node.id=} not in lru list" + assert node.id in self.cache, ( + f"Getting prev of node {node.id=} not in lru list" + ) x = getattr(node, self.prv) # x = node.prev while getattr(x, self.lock_ref) > 0 or len(x.children) > 0: x = getattr(x, self.prv) # x = x.prev @@ -396,28 +395,28 @@ class LRUList: # lists have independent recency, so they use different stamps) nodes.sort(key=lambda n: getattr(n, self.time_attr)) # the root node is not in the lru list - assert len(nodes) == ( - total_lru + (0 if self.mamba else 1) - ), f"len(nodes): {len(nodes)}, total_lru: {total_lru}" + assert len(nodes) == (total_lru + (0 if self.mamba else 1)), ( + f"len(nodes): {len(nodes)}, total_lru: {total_lru}" + ) x_lru = self._get_lru() for x in nodes: if x == tree_cache.root_node: # root node is not in the lru list continue - assert ( - x_lru is not None and x_lru.id in self.cache - ), f"Incorrect LRU list, x_lru is None or not in cache: {x_lru=}, {x.id=}" + assert x_lru is not None and x_lru.id in self.cache, ( + f"Incorrect LRU list, x_lru is None or not in cache: {x_lru=}, {x.id=}" + ) - assert ( - x == x_lru - ), f"Incorrect LRU list, {self.mamba=}, x: {x.id=} != x_lru: {x_lru.id=}, {getattr(x, self.time_attr)=}, {getattr(x_lru, self.time_attr)=}" - assert ( - x_lru.full_lock_ref == 0 - ), f"x_lru should not be locked when idle, {x_lru.full_lock_ref=}, {x_lru.id=}" - assert ( - x_lru.mamba_lock_ref == 0 - ), f"x_lru should not be locked when idle, {x_lru.mamba_lock_ref=}, {x_lru.id=}" + assert x == x_lru, ( + f"Incorrect LRU list, {self.mamba=}, x: {x.id=} != x_lru: {x_lru.id=}, {getattr(x, self.time_attr)=}, {getattr(x_lru, self.time_attr)=}" + ) + assert x_lru.full_lock_ref == 0, ( + f"x_lru should not be locked when idle, {x_lru.full_lock_ref=}, {x_lru.id=}" + ) + assert x_lru.mamba_lock_ref == 0, ( + f"x_lru should not be locked when idle, {x_lru.mamba_lock_ref=}, {x_lru.id=}" + ) x_lru = getattr(x, self.prv) if self.mamba: @@ -427,9 +426,9 @@ class LRUList: evictable_size = tree_cache.full_evictable_size() lru_list_evictable_size = self.sanity_check_evictable_size() - assert ( - evictable_size == lru_list_evictable_size - ), f"{self.mamba=}, total nodes: {total_nodes}, total lru: {total_lru}, evictable size: {evictable_size} != lru list evictable size: {lru_list_evictable_size}" + assert evictable_size == lru_list_evictable_size, ( + f"{self.mamba=}, total nodes: {total_nodes}, total lru: {total_lru}, evictable size: {evictable_size} != lru list evictable size: {lru_list_evictable_size}" + ) except Exception as e: if get_parallel().tp_rank == 0: msg = f"Mamba Radix tree sanity check failed, ping @yizhang2077: {e}" @@ -464,9 +463,9 @@ class MambaRadixCache(BasePrefixCache): ) if not self.enable_mamba_extra_buffer: - assert ( - self.page_size == 1 - ), f"Page size must be 1 for MambaRadixCache v1, got {self.page_size}" + assert self.page_size == 1, ( + f"Page size must be 1 for MambaRadixCache v1, got {self.page_size}" + ) if self.token_to_kv_pool_allocator: self.device = self.token_to_kv_pool_allocator.device @@ -592,9 +591,9 @@ class MambaRadixCache(BasePrefixCache): page_aligned_len = len(kv_indices) page_aligned_kv_indices = kv_indices.to(dtype=torch.int64, copy=True) - assert ( - cache_len == page_aligned_len - ), f"It is required {cache_len=}, {page_aligned_len=}, {kv_len_to_handle=}, {len(req.origin_input_ids)=}, {len(req.output_ids)=} ping @yizhang2077 if you see this" + assert cache_len == page_aligned_len, ( + f"It is required {cache_len=}, {page_aligned_len=}, {kv_len_to_handle=}, {len(req.origin_input_ids)=}, {len(req.output_ids)=} ping @yizhang2077 if you see this" + ) # Radix Cache takes one ref in memory pool # insert the token_ids and kv_indices into the radix tree @@ -706,9 +705,9 @@ class MambaRadixCache(BasePrefixCache): page_aligned_len = len(kv_indices) page_aligned_kv_indices = kv_indices.to(dtype=torch.int64, copy=True) - assert page_aligned_len == len( - kv_indices - ), f"page_aligned_len != len(kv_indices), {page_aligned_len=}, {len(kv_indices)=}, {cache_len=}, {self.page_size=}, {self.mamba_cache_chunk_size=}" + assert page_aligned_len == len(kv_indices), ( + f"page_aligned_len != len(kv_indices), {page_aligned_len=}, {len(kv_indices)=}, {cache_len=}, {self.page_size=}, {self.mamba_cache_chunk_size=}" + ) page_aligned_token_ids = token_ids[:page_aligned_len] @@ -778,12 +777,12 @@ class MambaRadixCache(BasePrefixCache): if not mamba_exist: assert torch.equal(new_last_node.mamba_value, mamba_value_donated) - assert ( - req.kv.cache_protected_len <= len(new_indices) + self.page_size - 1 - ), f"{req.kv.cache_protected_len=}, {len(new_indices)=}, {len(page_aligned_token_ids)=}, {mamba_exist=}" - assert new_prefix_len <= len( - new_indices - ), f"{new_prefix_len=}, {len(new_indices)=}" + assert req.kv.cache_protected_len <= len(new_indices) + self.page_size - 1, ( + f"{req.kv.cache_protected_len=}, {len(new_indices)=}, {len(page_aligned_token_ids)=}, {mamba_exist=}" + ) + assert new_prefix_len <= len(new_indices), ( + f"{new_prefix_len=}, {len(new_indices)=}" + ) self.req_to_token_pool.write( (req.kv.req_pool_idx, slice(req.kv.cache_protected_len, len(new_indices))), @@ -813,9 +812,9 @@ class MambaRadixCache(BasePrefixCache): def _evict_leaf_node( self, x: TreeNode, is_evict_mamba: bool ) -> Tuple[int, int, TreeNode, TreeNode]: - assert ( - x.full_lock_ref == 0 and x.mamba_lock_ref == 0 - ), f"evict leaf node invalid with {x.id=} {x.full_lock_ref=} {x.mamba_lock_ref=}" + assert x.full_lock_ref == 0 and x.mamba_lock_ref == 0, ( + f"evict leaf node invalid with {x.id=} {x.full_lock_ref=} {x.mamba_lock_ref=}" + ) assert x.mamba_value is not None, f"leaf node mamba value is not None, {x.id=}" # 1. a leaf node, free full tokens and mamba @@ -868,9 +867,9 @@ class MambaRadixCache(BasePrefixCache): # evict lru leaf nodes until mamba_num_tokens is reached while mamba_num_evicted < mamba_num and (self.mamba_lru_list.in_list(x)): assert x.mamba_value is not None, f"node has no mamba value, {x.id=}" - assert ( - len(x.mamba_value) == 1 - ), f"node has abnormal mamba length, {x.id=}, {len(x.mamba_value)=}" + assert len(x.mamba_value) == 1, ( + f"node has abnormal mamba length, {x.id=}, {len(x.mamba_value)=}" + ) assert x != self.root_node, f"root node is not evictable, {x.id=}" assert x.mamba_lock_ref == 0, f"node is in use by mamba kv indices, {x.id=}" @@ -903,9 +902,9 @@ class MambaRadixCache(BasePrefixCache): x = self.full_lru_list.get_leaf_lru_no_lock() while full_num_evicted < full_num_tokens and self.full_lru_list.in_list(x): - assert ( - x != self.root_node - ), f"root node should not exist in full lru list, {x.id=}" + assert x != self.root_node, ( + f"root node should not exist in full lru list, {x.id=}" + ) full_num_evicted_delta, _, x, x_next = self._evict_leaf_node(x, False) full_num_evicted += full_num_evicted_delta @@ -936,9 +935,9 @@ class MambaRadixCache(BasePrefixCache): while node != self.root_node: # lock full from node to root - assert ( - node.full_lock_ref >= 0 - ), f"inc_lock_ref on node with {node.full_lock_ref=}, {node.id=}" + assert node.full_lock_ref >= 0, ( + f"inc_lock_ref on node with {node.full_lock_ref=}, {node.id=}" + ) if node.full_lock_ref == 0: self.full_evictable_size_ -= len(node.value) self.full_protected_size_ += len(node.value) @@ -958,18 +957,18 @@ class MambaRadixCache(BasePrefixCache): return DecLockRefResult() if node.mamba_value is not None: - assert ( - node.mamba_lock_ref > 0 - ), f"dec_lock_ref on node with {node.mamba_lock_ref=}, {node.id=}" + assert node.mamba_lock_ref > 0, ( + f"dec_lock_ref on node with {node.mamba_lock_ref=}, {node.id=}" + ) if node.mamba_lock_ref == 1: self.mamba_evictable_size_ += len(node.mamba_value) self.mamba_protected_size_ -= len(node.mamba_value) node.mamba_lock_ref -= 1 while node != self.root_node: - assert ( - node.full_lock_ref > 0 - ), f"dec_lock_ref on node with {node.full_lock_ref=}, {node.id=}" + assert node.full_lock_ref > 0, ( + f"dec_lock_ref on node with {node.full_lock_ref=}, {node.id=}" + ) if node.full_lock_ref == 1: self.full_evictable_size_ += len(node.value) self.full_protected_size_ -= len(node.value) @@ -1326,9 +1325,9 @@ class MambaRadixCache(BasePrefixCache): # if locked, means node is in use, skip if node.parent.full_lock_ref > 0: break - assert ( - node.parent.mamba_lock_ref == 0 - ), f"tombstone mamba_lock_ref should always be 0, {node.parent.full_lock_ref=}, {node.parent.mamba_lock_ref=}, {node.parent.id=}" + assert node.parent.mamba_lock_ref == 0, ( + f"tombstone mamba_lock_ref should always be 0, {node.parent.full_lock_ref=}, {node.parent.mamba_lock_ref=}, {node.parent.id=}" + ) # delete tombstone node evicts full tokens self.kv_events.record_remove(node.parent) self.token_to_kv_pool_allocator.free_segment(node.parent.value, start_pos=0) @@ -1340,9 +1339,9 @@ class MambaRadixCache(BasePrefixCache): return node, full_num_evicted def _delete_leaf(self, node: TreeNode) -> None: - assert ( - node.mamba_value is not None - ), f"Invariant violated: leaf node is a tombstone, {node.id=}" + assert node.mamba_value is not None, ( + f"Invariant violated: leaf node is a tombstone, {node.id=}" + ) assert len(node.children) == 0, f"leaf node has children, {node.id=}" key = node.key.child_key(self.page_size) v = node.parent.children.pop(key, None) @@ -1357,9 +1356,9 @@ class MambaRadixCache(BasePrefixCache): node.mamba_value = None def _delete_tombstone_leaf(self, node: TreeNode) -> None: - assert ( - node.mamba_value is None - ), f"Deleting a unexpected non-tombstone leaf node, {node.id=}" + assert node.mamba_value is None, ( + f"Deleting a unexpected non-tombstone leaf node, {node.id=}" + ) assert len(node.children) == 0, f"leaf node has children, {node.id=}" key = node.key.child_key(self.page_size) v = node.parent.children.pop(key, None) @@ -1406,9 +1405,9 @@ class MambaRadixCache(BasePrefixCache): for key, child in current_node.children.items(): stack.append((child, current_indent + 2)) - assert key == child.key.child_key( - self.page_size - ), f"{key=}, {child.key.child_key(self.page_size)=}" + assert key == child.key.child_key(self.page_size), ( + f"{key=}, {child.key.child_key(self.page_size)=}" + ) def _total_size_helper(self) -> Tuple[int, int]: total_size = 0 diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 4a5ab085a..14ae9cdc5 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -294,9 +294,9 @@ class ReqToTokenPool: # Indices of reqs that already have a req_pool_idx and will reuse # their existing slot (e.g. chunked prefill continuing across chunks). reusing = [i for i, r in enumerate(reqs) if r.kv.holds_kv] - assert all( - reqs[i].kv.kv_allocated_len > 0 for i in reusing - ), "a reused row must carry allocated KV" + assert all(reqs[i].kv.kv_allocated_len > 0 for i in reusing), ( + "a reused row must carry allocated KV" + ) select_index = self.alloc_rows(len(reqs) - len(reusing)) if select_index is None: @@ -555,9 +555,9 @@ class MambaPool: # mamba layers/slots share one contiguous byte buffer; conv and # temporal are strided views into it (see mem_cache/layout/ # page_major.py). Only the standard CUDA Triton path is supported. - assert not _is_npu and not ( - _is_cpu and _cpu_has_amx_support - ), "envelope_layout mamba is only supported on the CUDA path" + assert not _is_npu and not (_is_cpu and _cpu_has_amx_support), ( + "envelope_layout mamba is only supported on the CUDA path" + ) max_slots = size + 1 entry_bytes = mamba_entry_bytes( layer_num=num_mamba_layers, @@ -1360,9 +1360,9 @@ class HybridReqToTokenPool(ReqToTokenPool): pass else: mid = self.mamba_allocator.alloc(1) - assert ( - mid is not None - ), f"Not enough space for mamba cache, try to increase --mamba-full-memory-ratio or --max-mamba-cache-size. {mid=}, {self.mamba_pool.size=}, {self.mamba_allocator.available_size()=}, {len(reqs)=}" + assert mid is not None, ( + f"Not enough space for mamba cache, try to increase --mamba-full-memory-ratio or --max-mamba-cache-size. {mid=}, {self.mamba_pool.size=}, {self.mamba_allocator.available_size()=}, {len(reqs)=}" + ) req.kv.mamba_pool_idx = mid[0] req.kv.mamba_needs_clear = True # GDN ReplaySSM: a freshly (re)assigned slot starts an empty @@ -1384,13 +1384,13 @@ class HybridReqToTokenPool(ReqToTokenPool): mamba_ping_pong_track_buffers.append( req.kv.mamba_ping_pong_track_buffer ) - assert len(select_index) == len( - mamba_indices - ), "Not enough space for mamba cache, try to increase --mamba-full-memory-ratio or --max-mamba-cache-size." + assert len(select_index) == len(mamba_indices), ( + "Not enough space for mamba cache, try to increase --mamba-full-memory-ratio or --max-mamba-cache-size." + ) if self.enable_mamba_extra_buffer: - assert len(select_index) == len( - mamba_ping_pong_track_buffers - ), "Not enough space for mamba ping pong idx, try to increase --mamba-full-memory-ratio." + assert len(select_index) == len(mamba_ping_pong_track_buffers), ( + "Not enough space for mamba ping pong idx, try to increase --mamba-full-memory-ratio." + ) mamba_index_tensor = torch.stack(mamba_indices).to(dtype=torch.int32) self.req_index_to_mamba_index_mapping[select_index] = mamba_index_tensor if self.enable_mamba_extra_buffer: @@ -1534,7 +1534,9 @@ class HybridReqToTokenPool(ReqToTokenPool): assert mamba_ping_pong_track_buffer_to_keep in [ 0, 1, - ], f"mamba_ping_pong_track_buffer_to_keep must be 0 or 1, {mamba_ping_pong_track_buffer_to_keep=}" + ], ( + f"mamba_ping_pong_track_buffer_to_keep must be 0 or 1, {mamba_ping_pong_track_buffer_to_keep=}" + ) # Avoid Python-list advanced indexing on a device tensor. # The ping-pong buffer size is either 2 (normal) or 1 (spec decode). if self.mamba_ping_pong_track_buffer_size == 2: @@ -1854,7 +1856,9 @@ class MHATokenToKVPool(KVCache): self.v_head_dim = ( swa_v_head_dim if swa_v_head_dim is not None - else v_head_dim if v_head_dim is not None else head_dim + else v_head_dim + if v_head_dim is not None + else head_dim ) # Layout: NHD (default) | HND (SGLANG_USE_HND_KVCACHE) | vectorized_5d (ROCm AITER). @@ -2886,9 +2890,9 @@ class MHATokenToKVPool(KVCache): if N == 0: return - assert ( - self._kv_copy_config is not None - ), "KV copy not initialized. Set enable_kv_cache_copy=True in __init__" + assert self._kv_copy_config is not None, ( + "KV copy not initialized. Set enable_kv_cache_copy=True in __init__" + ) cfg = self._kv_copy_config cap = int(cfg.get("num_locs_upper", 256)) @@ -3675,9 +3679,9 @@ class HybridLinearKVPool(KVCache): TokenToKVPoolClass = current_platform.get_mha_kv_pool_cls() quant_method_kwarg = {} elif _is_npu: - assert not is_float4_e2m1fn_x2( - dtype - ), "FP4 is not supported on NPU yet." + assert not is_float4_e2m1fn_x2(dtype), ( + "FP4 is not supported on NPU yet." + ) from sglang.srt.hardware_backend.npu.memory_pool_npu import ( NPUMHATokenToKVPool, ) @@ -4476,13 +4480,13 @@ class DSATokenToKVPool(MLATokenToKVPool): if _is_hip: if aiter_can_use_preshuffle_paged_mqa(): - assert ( - self.page_size % 16 == 0 - ), f"HIP preshuffle requires page_size to be a multiple of 16, got {self.page_size}" + assert self.page_size % 16 == 0, ( + f"HIP preshuffle requires page_size to be a multiple of 16, got {self.page_size}" + ) else: - assert ( - self.page_size == 1 - ), f"HIP legacy DSA path requires page_size == 1, got {self.page_size}" + assert self.page_size == 1, ( + f"HIP legacy DSA path requires page_size == 1, got {self.page_size}" + ) else: assert self.page_size == 64 self.index_key_cache = self._create_index_key_cache() diff --git a/python/sglang/srt/mem_cache/multi_ended_allocator.py b/python/sglang/srt/mem_cache/multi_ended_allocator.py index 6d0dfdc54..8371e425f 100644 --- a/python/sglang/srt/mem_cache/multi_ended_allocator.py +++ b/python/sglang/srt/mem_cache/multi_ended_allocator.py @@ -359,9 +359,9 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): self.high_peer: Optional[MultiEndedAllocator] = None # Inverse history of relocations (spec rollback), at PAGE granularity. - self._inverse_history: List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = ( - [] - ) + self._inverse_history: List[ + Tuple[torch.Tensor, torch.Tensor, torch.Tensor] + ] = [] # --- Lazy compaction state (all unused when lazy_compaction=False) --- # `_free_phys_pages`: GPU free list of physical PAGE ids, sorted at `_flush`. @@ -1222,9 +1222,9 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): v2p stays -1 and translation yields negative ids → CUDA OOB. """ with record_function("MultiEndedAlloc.alloc_extend"): - assert ( - self.is_id_owner - ), f"alloc_extend on a non-id-owner allocator ({self.sub_pool_name!r})" + assert self.is_id_owner, ( + f"alloc_extend on a non-id-owner allocator ({self.sub_pool_name!r})" + ) if num_new_pages is None: num_new_pages = get_num_new_pages( seq_lens=seq_lens_cpu, @@ -1291,9 +1291,9 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): virtual page on THIS sub-allocator (else v2p stays -1 → CUDA OOB). """ with record_function("MultiEndedAlloc.alloc_decode"): - assert ( - self.is_id_owner - ), f"alloc_decode on a non-id-owner allocator ({self.sub_pool_name!r})" + assert self.is_id_owner, ( + f"alloc_decode on a non-id-owner allocator ({self.sub_pool_name!r})" + ) bs = len(seq_lens) # CPU-only count BEFORE the kernel, to snapshot the exact slice the # kernel will consume. diff --git a/python/sglang/srt/mem_cache/pool_host/base.py b/python/sglang/srt/mem_cache/pool_host/base.py index 1c258a771..fa06cbf5b 100644 --- a/python/sglang/srt/mem_cache/pool_host/base.py +++ b/python/sglang/srt/mem_cache/pool_host/base.py @@ -386,9 +386,9 @@ class HostKVCache(abc.ABC): @synchronized def alloc(self, need_size: int) -> Optional[torch.Tensor]: - assert ( - need_size % self.logical_page_size == 0 - ), "The requested size should be a multiple of the page size." + assert need_size % self.logical_page_size == 0, ( + "The requested size should be a multiple of the page size." + ) if need_size > self.available_size(): return None diff --git a/python/sglang/srt/mem_cache/pool_host/common.py b/python/sglang/srt/mem_cache/pool_host/common.py index ee27e8b30..df7240298 100644 --- a/python/sglang/srt/mem_cache/pool_host/common.py +++ b/python/sglang/srt/mem_cache/pool_host/common.py @@ -23,9 +23,9 @@ class HostTensorAllocator: self.dims = None def allocate(self, dims: tuple, dtype: torch.dtype, device: str) -> torch.Tensor: - assert ( - device == "cpu" - ), f"HostTensorAllocator only supports CPU allocations; got device={device!r}" + assert device == "cpu", ( + f"HostTensorAllocator only supports CPU allocations; got device={device!r}" + ) self.dtype = dtype self.dims = dims return alloc_mmap(dims, dtype) @@ -46,9 +46,9 @@ class ShmHostTensorAllocator(HostTensorAllocator): return self.mms[0] if self.mms else None def allocate(self, dims: tuple, dtype: torch.dtype, device: str) -> torch.Tensor: - assert ( - device == "cpu" - ), f"ShmHostTensorAllocator only supports CPU allocations; got device={device!r}" + assert device == "cpu", ( + f"ShmHostTensorAllocator only supports CPU allocations; got device={device!r}" + ) self.dtype = dtype self.dims = dims from sglang.srt.mem_cache.storage.mmap import alloc_shm @@ -189,8 +189,7 @@ def _cuda_host_unregister_ranges( if rc != 0: failed_ranges.append((ptr, size)) logger.warning( - "cudaHostUnregister failed during %s (rc=%d, %s) " - "for ptr=%#x size=%d", + "cudaHostUnregister failed during %s (rc=%d, %s) for ptr=%#x size=%d", operation, rc, cudart.cudaGetErrorString(rc), diff --git a/python/sglang/srt/mem_cache/pool_host/mamba.py b/python/sglang/srt/mem_cache/pool_host/mamba.py index 2c982a204..669d99662 100644 --- a/python/sglang/srt/mem_cache/pool_host/mamba.py +++ b/python/sglang/srt/mem_cache/pool_host/mamba.py @@ -252,9 +252,9 @@ class MambaPoolHost(HostKVCache): @synchronized def alloc(self, need_size: int) -> Optional[torch.Tensor]: - assert ( - need_size % self.page_size == 0 - ), "The requested size should be a multiple of the page size." + assert need_size % self.page_size == 0, ( + "The requested size should be a multiple of the page size." + ) if need_size > self.available_size(): return None diff --git a/python/sglang/srt/mem_cache/pure_swa_radix_cache.py b/python/sglang/srt/mem_cache/pure_swa_radix_cache.py index 071d155c1..6dbed5b86 100644 --- a/python/sglang/srt/mem_cache/pure_swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/pure_swa_radix_cache.py @@ -34,9 +34,9 @@ class PureSWARadixCache(RadixCache): self.sliding_window_size = params.sliding_window_size def supports_swa(self) -> bool: - assert ( - self.sliding_window_size is not None - ), "sliding_window_size must be set for PureSWARadixCache" + assert self.sliding_window_size is not None, ( + "sliding_window_size must be set for PureSWARadixCache" + ) return True def swa_evictable_size(self): diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 8cebee66a..e6d5ba594 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -236,7 +236,6 @@ class RadixKey: class TreeNode: - counter = 0 def __init__(self, id: Optional[int] = None, priority: int = 0): @@ -553,9 +552,9 @@ class RadixCache(BasePrefixCache): match_result.device_indices, match_result.last_device_node, ) - assert len(new_indices) == len( - radix_key - ), f"{len(new_indices)=}, {len(radix_key)=}" + assert len(new_indices) == len(radix_key), ( + f"{len(new_indices)=}, {len(radix_key)=}" + ) self.req_to_token_pool.write( (req.kv.req_pool_idx, slice(req.kv.cache_protected_len, len(new_indices))), @@ -650,9 +649,9 @@ class RadixCache(BasePrefixCache): node.lock_ref -= 1 self._update_leaf_status(node) if node.parent is None: - assert ( - node is self.root_node - ), "This request holds the node from another tree" + assert node is self.root_node, ( + "This request holds the node from another tree" + ) node = node.parent return DecLockRefResult(delta=delta) @@ -804,9 +803,9 @@ class RadixCache(BasePrefixCache): for key, child in current_node.children.items(): stack.append((child, current_indent + 2)) - assert key == child.key.child_key( - self.page_size - ), f"{key=}, {child.key.child_key(self.page_size)=}" + assert key == child.key.child_key(self.page_size), ( + f"{key=}, {child.key.child_key(self.page_size)=}" + ) def _delete_leaf(self, node): key = node.key.child_key(self.page_size) diff --git a/python/sglang/srt/mem_cache/radix_cache_cpp.py b/python/sglang/srt/mem_cache/radix_cache_cpp.py index 52bdd9013..d1dc22ba5 100644 --- a/python/sglang/srt/mem_cache/radix_cache_cpp.py +++ b/python/sglang/srt/mem_cache/radix_cache_cpp.py @@ -52,9 +52,9 @@ class RadixCacheCpp(BasePrefixCache): self.disable = params.disable self.enable_write_cancel = enable_write_cancel - assert ( - params.enable_kv_cache_events is False - ), "HiRadixCache does not support kv cache events yet" + assert params.enable_kv_cache_events is False, ( + "HiRadixCache does not support kv cache events yet" + ) # record the nodes with ongoing write through self.ongoing_write_through: Set[IOHandle] = set() diff --git a/python/sglang/srt/mem_cache/rust_tree_core/adapter.py b/python/sglang/srt/mem_cache/rust_tree_core/adapter.py index 55f20e39b..a553d00ad 100644 --- a/python/sglang/srt/mem_cache/rust_tree_core/adapter.py +++ b/python/sglang/srt/mem_cache/rust_tree_core/adapter.py @@ -74,9 +74,9 @@ def _radix_key_buffer(key: RadixKey) -> array: """The key's token ids honoring `limit`; view-independent since the binding derives its own atoms.""" token_ids = key.raw_token_ids() - assert ( - isinstance(token_ids, array) and token_ids.typecode == "q" - ), f"tree keys must carry array('q') token ids, got {type(token_ids).__name__}" + assert isinstance(token_ids, array) and token_ids.typecode == "q", ( + f"tree keys must carry array('q') token ids, got {type(token_ids).__name__}" + ) return token_ids @@ -478,9 +478,9 @@ class RustUnifiedTreeCore(UnifiedTreeCoreInterface): self, node_id: NodeId, is_write_back: bool ) -> EvictDeviceLeafResult: # The binding reads is_write_back from the core's construction config. - assert ( - is_write_back == self.is_write_back - ), "is_write_back must match the core's construction config" + assert is_write_back == self.is_write_back, ( + "is_write_back must match the core's construction config" + ) binding_result = self._binding.evict_device_leaf(node_id) backup = binding_result.backup_kv result = EvictDeviceLeafResult( diff --git a/python/sglang/srt/mem_cache/storage/aibrix_kvcache/aibrix_kvcache_storage.py b/python/sglang/srt/mem_cache/storage/aibrix_kvcache/aibrix_kvcache_storage.py index 95f4d08a4..d4197a10a 100644 --- a/python/sglang/srt/mem_cache/storage/aibrix_kvcache/aibrix_kvcache_storage.py +++ b/python/sglang/srt/mem_cache/storage/aibrix_kvcache/aibrix_kvcache_storage.py @@ -83,9 +83,9 @@ class AibrixKVCacheStorage(HiCacheStorage): kv_blocks = handle.to_tensors() assert len(kv_blocks) == len(target_locations) for i in range(len(kv_blocks)): - assert ( - target_locations[i].nbytes == kv_blocks[i].nbytes - ), f"{target_locations[i].nbytes}, {kv_blocks[i].nbytes}" + assert target_locations[i].nbytes == kv_blocks[i].nbytes, ( + f"{target_locations[i].nbytes}, {kv_blocks[i].nbytes}" + ) target_locations[i].copy_(kv_blocks[i].flatten()) handle.release() return target_locations @@ -120,9 +120,9 @@ class AibrixKVCacheStorage(HiCacheStorage): logger.warning("aibrix_kvcache set allocate not enough") return False for i in range(len(tensors)): - assert ( - tensors[i].nbytes == values[i].nbytes - ), f"{tensors[i].nbytes}, {values[i].nbytes}" + assert tensors[i].nbytes == values[i].nbytes, ( + f"{tensors[i].nbytes}, {values[i].nbytes}" + ) tensors[i].reshape(values[i].shape).copy_(values[i]).reshape( tensors[i].shape ) diff --git a/python/sglang/srt/mem_cache/storage/eic/test_unit.py b/python/sglang/srt/mem_cache/storage/eic/test_unit.py index 03d348ad8..10d1f5c75 100644 --- a/python/sglang/srt/mem_cache/storage/eic/test_unit.py +++ b/python/sglang/srt/mem_cache/storage/eic/test_unit.py @@ -67,9 +67,9 @@ def test_set(eic_client): set_opt = eic.SetOption() set_opt.ttl_second = 3 status_code, set_outcome = eic_client.mset(data_keys, data_vals, set_opt) - assert ( - status_code == eic.StatusCode.SUCCESS - ), f"Set failed with status code: {status_code}" + assert status_code == eic.StatusCode.SUCCESS, ( + f"Set failed with status code: {status_code}" + ) def test_get(eic_client): @@ -87,9 +87,9 @@ def test_get(eic_client): ) get_opt = eic.GetOption() status_code, data_vals, get_outcome = eic_client.mget(data_keys, get_opt, data_vals) - assert ( - status_code == eic.StatusCode.SUCCESS - ), f"Get failed with status code: {status_code}" + assert status_code == eic.StatusCode.SUCCESS, ( + f"Get failed with status code: {status_code}" + ) def test_exists(eic_client): @@ -99,9 +99,9 @@ def test_exists(eic_client): data_keys.append(key) exists_opt = eic.ExistOption() status_code, exists_outcome = eic_client.mexist(data_keys, exists_opt) - assert ( - status_code == eic.StatusCode.SUCCESS - ), f"Exists failed with status code: {status_code}" + assert status_code == eic.StatusCode.SUCCESS, ( + f"Exists failed with status code: {status_code}" + ) def main(): diff --git a/python/sglang/srt/mem_cache/storage/flexkv/flexkv_comm.py b/python/sglang/srt/mem_cache/storage/flexkv/flexkv_comm.py index 00cabf6ae..609ba8f66 100644 --- a/python/sglang/srt/mem_cache/storage/flexkv/flexkv_comm.py +++ b/python/sglang/srt/mem_cache/storage/flexkv/flexkv_comm.py @@ -627,9 +627,9 @@ class FlexKVLayerDoneCounter: def update_producer(self) -> int: self.producer_index = (self.producer_index + 1) % self.num_counters - assert self.events[ - self.producer_index - ]._finished, "Producer event should be finished before reuse" + assert self.events[self.producer_index]._finished, ( + "Producer event should be finished before reuse" + ) return self.producer_index def set_consumer(self, task_id: int) -> None: diff --git a/python/sglang/srt/mem_cache/storage/flexkv/flexkv_connector.py b/python/sglang/srt/mem_cache/storage/flexkv/flexkv_connector.py index a1fb14462..7c4211179 100644 --- a/python/sglang/srt/mem_cache/storage/flexkv/flexkv_connector.py +++ b/python/sglang/srt/mem_cache/storage/flexkv/flexkv_connector.py @@ -397,8 +397,7 @@ class FlexKVConnector: payload = self._sync_ctx.scatter_pp(None) if payload.get("cmd") != CMD_LAYERWISE: raise RuntimeError( - f"Tag mismatch: expected CMD_LAYERWISE, got " - f"{payload.get('cmd')}" + f"Tag mismatch: expected CMD_LAYERWISE, got {payload.get('cmd')}" ) producer_id = int(payload["counter_id"]) self.layer_done_counter.register_task_with_explicit_counter_id( @@ -518,7 +517,7 @@ class FlexKVConnector: payload = self._sync_ctx.scatter_pp(None) if payload.get("cmd") != CMD_PUT_META: raise RuntimeError( - f"Tag mismatch: expected CMD_PUT_META, got " f"{payload.get('cmd')}" + f"Tag mismatch: expected CMD_PUT_META, got {payload.get('cmd')}" ) fkv_task_id = int(payload["fkv_task_id"]) mask_list = payload.get("unmatched_mask", []) @@ -758,7 +757,7 @@ class FlexKVConnector: raise if attempt % 30 == 0: logger.info( - "[FlexKV] GPU register retry %s attempt=%d/%d " "error=%s", + "[FlexKV] GPU register retry %s attempt=%d/%d error=%s", self._label, attempt + 1, max_retries, @@ -772,9 +771,9 @@ class FlexKVConnector: indexer_buffers: Optional[List[torch.Tensor]] = None, ) -> None: assert len(kv_caches) > 0 - assert ( - kv_caches[0].ndim == 3 - ), f"Expected 3D KV cache tensor, got shape={kv_caches[0].shape}" + assert kv_caches[0].ndim == 3, ( + f"Expected 3D KV cache tensor, got shape={kv_caches[0].shape}" + ) is_mla = self.model_config.use_mla num_blocks, num_kv_heads, head_size = kv_caches[0].shape @@ -896,8 +895,7 @@ class FlexKVConnector: ) from exc if not ack or ack[0] != 1: raise RuntimeError( - f"FlexKV layerwise worker NACK'd eventfd transfer " - f"(ack={ack!r})" + f"FlexKV layerwise worker NACK'd eventfd transfer (ack={ack!r})" ) logger.info( "[FlexKV] Eventfd handshake complete %s counters=%d layers=%d", diff --git a/python/sglang/srt/mem_cache/storage/hf3fs/storage_hf3fs.py b/python/sglang/srt/mem_cache/storage/hf3fs/storage_hf3fs.py index 84ab181eb..7370f403e 100644 --- a/python/sglang/srt/mem_cache/storage/hf3fs/storage_hf3fs.py +++ b/python/sglang/srt/mem_cache/storage/hf3fs/storage_hf3fs.py @@ -228,7 +228,7 @@ class HiCacheHF3FS(HiCacheStorage): logger.info( f"[Rank {self.rank}] HiCacheHF3FS Client Initializing: " f"file_path={self.file_path}, " - f"file_size={self.file_size / (2 ** 30):.2f} GB, " + f"file_size={self.file_size / (2**30):.2f} GB, " f"num_pages={self.num_pages}, " f"is_mla_model={self.is_mla_model}" ) diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py index 6266059cf..ec6c606a5 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py @@ -330,7 +330,6 @@ class MooncakeBaseStore: class MooncakeStore(HiCacheStorage, MooncakeBaseStore): - @staticmethod def _standalone_required_bytes(mem_pool: Any) -> int: """Compute total bytes of host buffers that must be visible to the real client. diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/test_mooncake_store.py b/python/sglang/srt/mem_cache/storage/mooncake_store/test_mooncake_store.py index 9929ea228..6df8d3a79 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/test_mooncake_store.py +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/test_mooncake_store.py @@ -117,9 +117,9 @@ def test_single_operation(): assert result is True, f"❌get operation failed for key: {key}" # Compare the data using proper tensor indices - assert torch.allclose( - set_slice, get_slice, atol=1e-6 - ), f"❌get operation failed for key: {key}" + assert torch.allclose(set_slice, get_slice, atol=1e-6), ( + f"❌get operation failed for key: {key}" + ) logger.info(f"✅ Single operation passed") @@ -154,15 +154,16 @@ def test_batch_operation(config: HiCacheStorageConfig): assert all(result), "batch set operation failed" # Test batch exists operation - assert ( - store.batch_exists(keys) == kv_num - ), "keys should exist after batch set operation" + assert store.batch_exists(keys) == kv_num, ( + "keys should exist after batch set operation" + ) # Test batch get operation get_slices = [ buffer[ - (kv_num * entries_per_page + i) - * value_elements : (kv_num * entries_per_page + i + 1) + (kv_num * entries_per_page + i) * value_elements : ( + kv_num * entries_per_page + i + 1 + ) * value_elements ] for i in range(kv_num * entries_per_page) @@ -171,9 +172,9 @@ def test_batch_operation(config: HiCacheStorageConfig): result = store.batch_get_v1(keys, get_indices) assert all(result), "❌batch get operation failed" for i in range(kv_num * entries_per_page): - assert torch.allclose( - set_slices[i], get_slices[i], atol=1e-6 - ), f"❌batch get operation failed for key: {keys[i // entries_per_page]}" + assert torch.allclose(set_slices[i], get_slices[i], atol=1e-6), ( + f"❌batch get operation failed for key: {keys[i // entries_per_page]}" + ) logger.info(f"✅ Batch operation passed") diff --git a/python/sglang/srt/mem_cache/storage/simm/hicache_simm.py b/python/sglang/srt/mem_cache/storage/simm/hicache_simm.py index 66d274ab4..dfb9341ca 100644 --- a/python/sglang/srt/mem_cache/storage/simm/hicache_simm.py +++ b/python/sglang/srt/mem_cache/storage/simm/hicache_simm.py @@ -129,7 +129,6 @@ def get_numa_nic_mapping() -> Dict[int, List[str]]: class HiCacheSiMM(HiCacheStorage): - def __init__( self, storage_config: HiCacheStorageConfig = None, mem_pool: HostKVCache = None ): @@ -239,7 +238,7 @@ class HiCacheSiMM(HiCacheStorage): if not all(got_block.as_ref()[: len(warmup_key)] == warmup_tensor): logger.warning(f"SiMM client warmup key {warmup_key} data wrong") logger.info( - f"finish SiMM client warm up, cost {(time.perf_counter_ns() - start_time)/1000:.2f} us" + f"finish SiMM client warm up, cost {(time.perf_counter_ns() - start_time) / 1000:.2f} us" ) def register_mem_pool_host(self, mem_pool_host: HostKVCache): @@ -334,7 +333,7 @@ class HiCacheSiMM(HiCacheStorage): if self.config.enable_profile: logger.info( f"SiMM batch_get_v1 {len(keys)} keys, total size: {total_size / 1024**2} MiB, \ - using {(t2 - t1)/1000} us, Throughput: {total_size / 1024**3 / ((t2 - t1) / 1000**3):.2f} GiB/s" + using {(t2 - t1) / 1000} us, Throughput: {total_size / 1024**3 / ((t2 - t1) / 1000**3):.2f} GiB/s" ) return self._batch_postprocess(get_results, is_set_operate=False) @@ -355,7 +354,7 @@ class HiCacheSiMM(HiCacheStorage): t2 = time.perf_counter_ns() if self.config.enable_profile: logger.info( - f"SiMM batch exists {len(keys)} keys, using {(t2 - t1)/1000} us" + f"SiMM batch exists {len(keys)} keys, using {(t2 - t1) / 1000} us" ) set_keys = [] @@ -385,7 +384,7 @@ class HiCacheSiMM(HiCacheStorage): if self.config.enable_profile: logger.info( f"SiMM batch_put_v1 {len(keys)} keys, total size: {total_size / 1024**2} MiB, \ - using {(t3 - t2)/1000} us, Throughput: {total_size / 1024**3 / ((t3 - t2) / 1000**3):.2f} GiB/s" + using {(t3 - t2) / 1000} us, Throughput: {total_size / 1024**3 / ((t3 - t2) / 1000**3):.2f} GiB/s" ) return self._batch_postprocess(set_results, is_set_operate=True) @@ -511,7 +510,7 @@ class HiCacheSiMM(HiCacheStorage): t2 = time.perf_counter_ns() if self.config.enable_profile: logger.info( - f"SiMM batch exists {len(keys)} keys, using {(t2 - t1)/1000} us" + f"SiMM batch exists {len(keys)} keys, using {(t2 - t1) / 1000} us" ) for i in range(len(query_keys)): if not exist_result[i]: diff --git a/python/sglang/srt/mem_cache/storage/simm/test_simm.py b/python/sglang/srt/mem_cache/storage/simm/test_simm.py index 74e4ec65b..fa3660e26 100644 --- a/python/sglang/srt/mem_cache/storage/simm/test_simm.py +++ b/python/sglang/srt/mem_cache/storage/simm/test_simm.py @@ -82,9 +82,9 @@ def test_single_operation(): assert result is True, f"❌get operation failed for key: {key}" # Compare the data using proper tensor indices - assert torch.allclose( - set_slice, get_slice, atol=1e-6 - ), f"❌get operation failed for key: {key}" + assert torch.allclose(set_slice, get_slice, atol=1e-6), ( + f"❌get operation failed for key: {key}" + ) logger.info(f"✅ Single operation passed") @@ -116,15 +116,14 @@ def test_batch_operation(config: HiCacheStorageConfig): assert all(result), f"❌batch set operation failed" # Test batch exists operation - assert store.batch_exists( - exist_keys - ), f"❌keys should exist after batch set operation" + assert store.batch_exists(exist_keys), ( + f"❌keys should exist after batch set operation" + ) # Test batch get operation get_slices = [ buffer[ - (len(set_keys) + i) - * value_elements : (len(set_keys) + i + 1) + (len(set_keys) + i) * value_elements : (len(set_keys) + i + 1) * value_elements ] for i in range(len(get_keys)) @@ -133,9 +132,9 @@ def test_batch_operation(config: HiCacheStorageConfig): result = store.batch_get_v1(get_keys, get_indices) assert all(result), f"❌batch get operation failed" for i in range(len(get_keys)): - assert torch.allclose( - set_slices[i], get_slices[i], atol=1e-6 - ), f"❌batch get operation failed for key: {get_keys[i]}" + assert torch.allclose(set_slices[i], get_slices[i], atol=1e-6), ( + f"❌batch get operation failed for key: {get_keys[i]}" + ) logger.info(f"✅ Batch operation passed") diff --git a/python/sglang/srt/mem_cache/storage/umbp/umbp_store.py b/python/sglang/srt/mem_cache/storage/umbp/umbp_store.py index f64f096f3..c6affd41b 100644 --- a/python/sglang/srt/mem_cache/storage/umbp/umbp_store.py +++ b/python/sglang/srt/mem_cache/storage/umbp/umbp_store.py @@ -361,8 +361,7 @@ class UMBPStore(HiCacheStorage): ssd_backend = str(extra["ssd_backend"]).strip().lower() if ssd_backend not in ("file", "spdk", "spdk_proxy"): raise ValueError( - "extra_config['ssd_backend'] must be one of: " - "file, spdk, spdk_proxy" + "extra_config['ssd_backend'] must be one of: file, spdk, spdk_proxy" ) cfg.ssd.ssd_backend = ssd_backend if "spdk_nvme_pci_addr" in extra: diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index 5b6862d48..f37260171 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -56,7 +56,6 @@ logger = logging.getLogger(__name__) class TreeNode: - counter = 0 swa_uuid_counter = 1 last_access_time_counter_float = float64(1.0) @@ -177,9 +176,9 @@ class LRUList: Move a (existing) node to most recently used position """ assert node.id in self.cache, f"Resetting node {node.id=} not in lru list" - assert ( - not self.is_swa_list or not node.swa_tombstone - ), f"Resetting swa tombstone node in swa lru list: {node.id=}" + assert not self.is_swa_list or not node.swa_tombstone, ( + f"Resetting swa tombstone node in swa lru list: {node.id=}" + ) self._remove_node(node) self._add_node(node) @@ -192,9 +191,9 @@ class LRUList: while node != root_node: # for swa lru list, only reset non-tombstone nodes if not self.is_swa_list or not node.swa_tombstone: - assert ( - node.id in self.cache - ), f"Resetting node {node.id=} not in lru list when resetting node and parents mru" + assert node.id in self.cache, ( + f"Resetting node {node.id=} not in lru list when resetting node and parents mru" + ) self._remove_node(node) self._add_node_after(prev_node, node) prev_node = node @@ -204,12 +203,12 @@ class LRUList: """ Insert a (new) node as most recently used """ - assert ( - not self.is_swa_list or not node.swa_tombstone - ), f"Inserting swa tombstone node in swa lru list: {node.id=}" - assert ( - node.id not in self.cache - ), f"Inserting node {node.id=} already in lru list, existing node: {self.cache[node.id].id=}" + assert not self.is_swa_list or not node.swa_tombstone, ( + f"Inserting swa tombstone node in swa lru list: {node.id=}" + ) + assert node.id not in self.cache, ( + f"Inserting node {node.id=} already in lru list, existing node: {self.cache[node.id].id=}" + ) self.cache[node.id] = node self._add_node(node) @@ -218,9 +217,9 @@ class LRUList: Remove node from lru list """ assert node.id in self.cache, f"Removing node {node.id=} not in lru list" - assert ( - not self.is_swa_list or not node.swa_tombstone - ), f"Removing swa tombstone node from swa lru list: {node.id=}" + assert not self.is_swa_list or not node.swa_tombstone, ( + f"Removing swa tombstone node from swa lru list: {node.id=}" + ) del self.cache[node.id] self._remove_node(node) @@ -243,9 +242,9 @@ class LRUList: Get the previous (i.e. more recently used) node that is not locked """ if check_id: - assert ( - node.id in self.cache - ), f"Getting prev of node {node.id=} not in lru list" + assert node.id in self.cache, ( + f"Getting prev of node {node.id=} not in lru list" + ) x = getattr(node, self.prv) # x = node.prev while getattr(x, self.lock_ref) > 0: x = getattr(x, self.prv) # x = x.prev @@ -259,9 +258,9 @@ class LRUList: Get the previous (i.e. more recently used) leaf node that is not locked """ if check_id: - assert ( - node.id in self.cache - ), f"Getting prev of node {node.id=} not in lru list" + assert node.id in self.cache, ( + f"Getting prev of node {node.id=} not in lru list" + ) x = getattr(node, self.prv) # x = node.prev while getattr(x, self.lock_ref) > 0 or len(x.children) > 0: x = getattr(x, self.prv) # x = x.prev @@ -306,9 +305,9 @@ class LRUList: # heapify based on last_access_time heapq.heapify(nodes) # the root node is not in the lru list - assert ( - len(nodes) == len(self.cache) + 1 - ), f"len(nodes): {len(nodes)} != len(self.cache) + 1: {len(self.cache) + 1}" + assert len(nodes) == len(self.cache) + 1, ( + f"len(nodes): {len(nodes)} != len(self.cache) + 1: {len(self.cache) + 1}" + ) x_lru = self._get_lru() while len(nodes): @@ -316,15 +315,15 @@ class LRUList: if x == tree_cache.root_node: # root node is not in the lru list continue - assert ( - x == x_lru - ), f"Incorrect LRU list, {self.is_swa_list=}, x: {x.id=} != x_lru: {x_lru.id=}" - assert ( - x_lru.full_lock_ref == 0 - ), f"x_lru should not be locked when idle, {x_lru.full_lock_ref=}, {x_lru.swa_uuid=}, {x_lru.id=}" - assert ( - x_lru.swa_lock_ref == 0 - ), f"x_lru should not be locked when idle, {x_lru.swa_lock_ref=}, {x_lru.swa_uuid=}, {x_lru.id=}" + assert x == x_lru, ( + f"Incorrect LRU list, {self.is_swa_list=}, x: {x.id=} != x_lru: {x_lru.id=}" + ) + assert x_lru.full_lock_ref == 0, ( + f"x_lru should not be locked when idle, {x_lru.full_lock_ref=}, {x_lru.swa_uuid=}, {x_lru.id=}" + ) + assert x_lru.swa_lock_ref == 0, ( + f"x_lru should not be locked when idle, {x_lru.swa_lock_ref=}, {x_lru.swa_uuid=}, {x_lru.id=}" + ) x_lru = getattr(x, self.prv) if self.is_swa_list: @@ -334,9 +333,9 @@ class LRUList: evictable_size = tree_cache.full_evictable_size() lru_list_evictable_size = self.sanity_check_evictable_size() - assert ( - evictable_size == lru_list_evictable_size - ), f"{self.is_swa_list=}, total nodes: {total_nodes}, total lru plus 1: {total_lru_plus_1}, evictable size: {evictable_size} != lru list evictable size: {lru_list_evictable_size}" + assert evictable_size == lru_list_evictable_size, ( + f"{self.is_swa_list=}, total nodes: {total_nodes}, total lru plus 1: {total_lru_plus_1}, evictable size: {evictable_size} != lru list evictable size: {lru_list_evictable_size}" + ) except Exception as e: msg = f"SWA Radix tree sanity check failed, ping @hanming-lu: {e}" logger.error(msg) @@ -369,9 +368,9 @@ class SWARadixCache(BasePrefixCache): ##### Public API ##### def supports_swa(self) -> bool: - assert ( - self.sliding_window_size is not None - ), "sliding_window_size must be set for SWARadixCache" + assert self.sliding_window_size is not None, ( + "sliding_window_size must be set for SWARadixCache" + ) return True def swa_reprefill_tail_tokens(self) -> int: @@ -618,9 +617,9 @@ class SWARadixCache(BasePrefixCache): x = self.full_lru_list.get_leaf_lru_no_lock() while full_num_evicted < full_num_tokens and self.full_lru_list.in_list(x): - assert ( - x != self.root_node - ), f"root node should not exist in full lru list, {x.id=}" + assert x != self.root_node, ( + f"root node should not exist in full lru list, {x.id=}" + ) assert x.full_lock_ref == 0, f"node is in use, {x.id=}" # 1. free node kv indices, evict full and swa tokens @@ -683,9 +682,9 @@ class SWARadixCache(BasePrefixCache): self.swa_evictable_size_ -= len(x.value) x.swa_tombstone = True else: - assert ( - x.full_lock_ref == 0 - ), f"leaf node with full lock must also have swa lock, {x.id=}" + assert x.full_lock_ref == 0, ( + f"leaf node with full lock must also have swa lock, {x.id=}" + ) # 1. a leaf node, free full and swa tokens self.kv_events.record_remove(x) node_full_evicted, node_swa_evicted = self._free_node_value(x) @@ -724,9 +723,9 @@ class SWARadixCache(BasePrefixCache): swa_uuid_for_lock = None while node != self.root_node: # lock full from node to root - assert ( - node.full_lock_ref >= 0 - ), f"inc_lock_ref on node with {node.full_lock_ref=}, {node.id=}" + assert node.full_lock_ref >= 0, ( + f"inc_lock_ref on node with {node.full_lock_ref=}, {node.id=}" + ) if node.full_lock_ref == 0: self.full_evictable_size_ -= len(node.value) self.full_protected_size_ += len(node.value) @@ -736,9 +735,9 @@ class SWARadixCache(BasePrefixCache): # When we reach the sliding window size, we will set the swa_uuid_for_lock. # caller needs to pass the swa_uuid_for_lock to dec_lock_ref if swa_lock_size < self.sliding_window_size: - assert ( - not node.swa_tombstone - ), f"inc_lock_swa on swa_tombstone node, {node.id=}" + assert not node.swa_tombstone, ( + f"inc_lock_swa on swa_tombstone node, {node.id=}" + ) if node.swa_lock_ref == 0: self.swa_evictable_size_ -= len(node.value) self.swa_protected_size_ += len(node.value) @@ -773,21 +772,21 @@ class SWARadixCache(BasePrefixCache): dec_lock_swa = not skip_swa while node != self.root_node: - assert ( - node.full_lock_ref > 0 - ), f"dec_lock_ref on node with {node.full_lock_ref=}, {node.id=}" + assert node.full_lock_ref > 0, ( + f"dec_lock_ref on node with {node.full_lock_ref=}, {node.id=}" + ) if node.full_lock_ref == 1: self.full_evictable_size_ += len(node.value) self.full_protected_size_ -= len(node.value) node.full_lock_ref -= 1 if dec_lock_swa: - assert ( - not node.swa_tombstone - ), f"dec_lock_ref on swa_tombstone node, {node.id=}" - assert ( - node.swa_lock_ref > 0 - ), f"dec_lock_ref on node with {node.swa_lock_ref=}, {node.id=}" + assert not node.swa_tombstone, ( + f"dec_lock_ref on swa_tombstone node, {node.id=}" + ) + assert node.swa_lock_ref > 0, ( + f"dec_lock_ref on node with {node.swa_lock_ref=}, {node.id=}" + ) if node.swa_lock_ref == 1: self.swa_evictable_size_ += len(node.value) @@ -832,12 +831,12 @@ class SWARadixCache(BasePrefixCache): return while node != self.root_node: - assert ( - not node.swa_tombstone - ), f"dec_swa_lock_only on swa_tombstone node, {node.id=}" - assert ( - node.swa_lock_ref > 0 - ), f"dec_swa_lock_only on node with {node.swa_lock_ref=}, {node.id=}" + assert not node.swa_tombstone, ( + f"dec_swa_lock_only on swa_tombstone node, {node.id=}" + ) + assert node.swa_lock_ref > 0, ( + f"dec_swa_lock_only on node with {node.swa_lock_ref=}, {node.id=}" + ) if node.swa_lock_ref == 1: self.swa_protected_size_ -= len(node.value) @@ -1186,12 +1185,12 @@ class SWARadixCache(BasePrefixCache): # For page_size > 1 and chunked prefill case, update_kv_after_len may be not page-aligned due to a trailing partial page # (kept in the request but not inserted into the radix tree) appended to prefix_indices. if node.swa_tombstone: - assert ( - node.swa_lock_ref == 0 - ), f"tombstone swa_lock_ref should always be 0, {node.full_lock_ref=}, {node.swa_lock_ref=}, {node.id=}" - assert ( - swa_evicted_seqlen % self.page_size == 0 - ), f"swa_evicted_seqlen must be page aligned, {swa_evicted_seqlen=}, {self.page_size=}" + assert node.swa_lock_ref == 0, ( + f"tombstone swa_lock_ref should always be 0, {node.full_lock_ref=}, {node.swa_lock_ref=}, {node.id=}" + ) + assert swa_evicted_seqlen % self.page_size == 0, ( + f"swa_evicted_seqlen must be page aligned, {swa_evicted_seqlen=}, {self.page_size=}" + ) if swa_evicted_seqlen <= total_prefix_length: # Branch 1: all swa tokens of value[:prefix_len] are not evicted, so we can insert it to the tree directly. if node.full_lock_ref > 0: @@ -1359,9 +1358,9 @@ class SWARadixCache(BasePrefixCache): # if locked, means node is in use, skip if node.parent.full_lock_ref > 0: break - assert ( - node.parent.swa_lock_ref == 0 - ), f"tombstone swa_lock_ref should always be 0, {node.parent.full_lock_ref=}, {node.parent.swa_lock_ref=}, {node.parent.id=}" + assert node.parent.swa_lock_ref == 0, ( + f"tombstone swa_lock_ref should always be 0, {node.parent.full_lock_ref=}, {node.parent.swa_lock_ref=}, {node.parent.id=}" + ) # delete tombstone node evicts full tokens self.kv_events.record_remove(node.parent) node_full_evicted, _ = self._free_node_value(node.parent) @@ -1389,9 +1388,9 @@ class SWARadixCache(BasePrefixCache): self.swa_evictable_size_ -= len(node.key) def _delete_tombstone_leaf(self, node: TreeNode) -> None: - assert ( - node.swa_tombstone - ), f"Deleting a unexpected non-tombstone leaf node, {node.id=}" + assert node.swa_tombstone, ( + f"Deleting a unexpected non-tombstone leaf node, {node.id=}" + ) assert len(node.children) == 0, f"leaf node has children, {node.id=}" key = node.key.child_key(self.page_size) v = node.parent.children.pop(key, None) @@ -1438,9 +1437,9 @@ class SWARadixCache(BasePrefixCache): for key, child in current_node.children.items(): stack.append((child, current_indent + 2)) - assert key == child.key.child_key( - self.page_size - ), f"{key=}, {child.key.child_key(self.page_size)=}" + assert key == child.key.child_key(self.page_size), ( + f"{key=}, {child.key.child_key(self.page_size)=}" + ) def _total_size_helper(self) -> Tuple[int, int]: total_size = 0 diff --git a/python/sglang/srt/mem_cache/unified_cache/components/full_component.py b/python/sglang/srt/mem_cache/unified_cache/components/full_component.py index c9102411d..95fecb2ab 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/full_component.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/full_component.py @@ -106,8 +106,8 @@ class FullComponent(TreeComponent): self, match_device_only: bool = False ) -> Callable[[UnifiedTreeNode], bool]: if match_device_only: - return ( - lambda node: node.component_data[self.component_type].value is not None + return lambda node: ( + node.component_data[self.component_type].value is not None ) # HiCache: evicted + backuped nodes are valid match boundaries. @@ -290,9 +290,9 @@ class FullComponent(TreeComponent): delta = 0 while cur is not root: cd = cur.component_data[ct] - assert ( - cd.value is not None - ), f"FULL invariant broken: evicted ancestor {cur.id} above device-on segment" + assert cd.value is not None, ( + f"FULL invariant broken: evicted ancestor {cur.id} above device-on segment" + ) if cd.lock_ref == 0: key_len = len(cd.value) self.tree_core.component_evictable_size_[ct] -= key_len diff --git a/python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py b/python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py index 0e6042af1..f60d9a713 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py @@ -62,13 +62,13 @@ class MambaComponent(TreeComponent): def __init__(self, cache: UnifiedRadixCache, params: CacheInitParams): from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool - assert isinstance( - params.req_to_token_pool, HybridReqToTokenPool - ), f"MambaComponent requires HybridReqToTokenPool, got {type(params.req_to_token_pool)}" + assert isinstance(params.req_to_token_pool, HybridReqToTokenPool), ( + f"MambaComponent requires HybridReqToTokenPool, got {type(params.req_to_token_pool)}" + ) if not params.enable_mamba_extra_buffer: - assert ( - params.page_size == 1 - ), f"MambaComponent requires page_size=1 when mamba_extra_buffer is disabled, got {params.page_size}" + assert params.page_size == 1, ( + f"MambaComponent requires page_size=1 when mamba_extra_buffer is disabled, got {params.page_size}" + ) super().__init__(cache, params) self.mamba_cache_chunk_size = mamba_cache_chunk_size() # params.page_size is the tree page the allocator actually uses, already diff --git a/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py index 6e9dccfc7..fd03d1b35 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py @@ -70,7 +70,9 @@ class SWAComponent(TreeComponent): assert isinstance( params.token_to_kv_pool_allocator, SWATokenToKVPoolAllocator - ), f"SWAComponent requires SWATokenToKVPoolAllocator, got {type(params.token_to_kv_pool_allocator)}" + ), ( + f"SWAComponent requires SWATokenToKVPoolAllocator, got {type(params.token_to_kv_pool_allocator)}" + ) super().__init__(cache, params) self._session_leaf_covered_len: dict[str, dict[UnifiedTreeNode, int]] = {} self.sliding_window_size = params.sliding_window_size @@ -239,9 +241,9 @@ class SWAComponent(TreeComponent): # were just allocated by the in-flight request, so every page must be # live; a violation means we would hand the node the sink and serve # zeros, which is worth a hard failure rather than silent corruption. - assert bool( - (physical > 0).all() - ), f"incoming swa pages must all be live, got {physical.tolist()}" + assert bool((physical > 0).all()), ( + f"incoming swa pages must all be live, got {physical.tolist()}" + ) swa.bind(kept_pages, physical) swa.virtual_to_physical.index_fill_(0, incoming_pages, -1) swa.clear_inverse_history() @@ -351,12 +353,12 @@ class SWAComponent(TreeComponent): full_cd = node.component_data[BASE_COMPONENT_TYPE] swa_evicted_seqlen = params.swa_evicted_seqlen - assert ( - node.component_data[self.component_type].lock_ref == 0 - ), f"tombstone {self.component_type} lock_ref should be 0, node {node.id}" - assert ( - swa_evicted_seqlen % self.tree_core.page_size == 0 - ), f"{self.component_type}: swa_evicted_seqlen must be page-aligned, {swa_evicted_seqlen=}" + assert node.component_data[self.component_type].lock_ref == 0, ( + f"tombstone {self.component_type} lock_ref should be 0, node {node.id}" + ) + assert swa_evicted_seqlen % self.tree_core.page_size == 0, ( + f"{self.component_type}: swa_evicted_seqlen must be page-aligned, {swa_evicted_seqlen=}" + ) if swa_evicted_seqlen <= total_prefix_len: # Branch 1: entire value_slice is within SWA window — recover @@ -427,13 +429,13 @@ class SWAComponent(TreeComponent): ct = self.component_type if node.component_data[ct].value is not None: return - assert ( - node.component_data[ct].lock_ref == 0 - ), f"tombstone {ct} lock_ref should be 0 on unevict, node {node.id}" + assert node.component_data[ct].lock_ref == 0, ( + f"tombstone {ct} lock_ref should be 0 on unevict, node {node.id}" + ) swa_evicted_seqlen = params.swa_evicted_seqlen - assert ( - swa_evicted_seqlen % self.tree_core.page_size == 0 - ), f"{ct}: swa_evicted_seqlen must be page-aligned, {swa_evicted_seqlen=}" + assert swa_evicted_seqlen % self.tree_core.page_size == 0, ( + f"{ct}: swa_evicted_seqlen must be page-aligned, {swa_evicted_seqlen=}" + ) if swa_evicted_seqlen <= total_prefix_len: pass # entire node is within the SWA window @@ -534,9 +536,9 @@ class SWAComponent(TreeComponent): new_parent.component_data[self.component_type].lock_ref = child.component_data[ self.component_type ].lock_ref - new_parent.component_data[self.component_type].session_ref = ( - child.component_data[self.component_type].session_ref - ) + new_parent.component_data[ + self.component_type + ].session_ref = child.component_data[self.component_type].session_ref assert new_parent.component_data[self.component_type].session_ids is None child_swa_value = child.component_data[self.component_type].value @@ -554,9 +556,9 @@ class SWAComponent(TreeComponent): child_swa_host_value = child.component_data[self.component_type].host_value if child_swa_host_value is not None: split_len = len(new_parent.key) - new_parent.component_data[self.component_type].host_value = ( - child_swa_host_value[:split_len].clone() - ) + new_parent.component_data[ + self.component_type + ].host_value = child_swa_host_value[:split_len].clone() child.component_data[self.component_type].host_value = child_swa_host_value[ split_len: ].clone() diff --git a/python/sglang/srt/mem_cache/unified_cache/components/tree_component.py b/python/sglang/srt/mem_cache/unified_cache/components/tree_component.py index 55b6ca8fc..da628f407 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/tree_component.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/tree_component.py @@ -82,7 +82,6 @@ class PreparePrefetchResult: class CacheTransferPhase(str, Enum): - BACKUP_HOST = "backup_host" # D→H LOAD_BACK = "load_back" # H→D BACKUP_STORAGE = "backup_storage" # H→Storage @@ -508,9 +507,9 @@ class TreeComponent(ABC): def evict_device_start(self, request_cnt: int) -> None: """Begin this component's device-eviction walk (build its cursor/heap).""" - assert ( - not self.is_evict_device_ongoing - ), f"{self.component_type} device eviction already in progress" + assert not self.is_evict_device_ongoing, ( + f"{self.component_type} device eviction already in progress" + ) self._evict_device_start(request_cnt) self.is_evict_device_ongoing = True @@ -525,16 +524,16 @@ class TreeComponent(ABC): Implementations must return after one allocator-relevant internal mutation so the caller can drain pending frees before continuing. """ - assert ( - self.is_evict_device_ongoing - ), f"{self.component_type} device eviction not started" + assert self.is_evict_device_ongoing, ( + f"{self.component_type} device eviction not started" + ) return self._evict_device_next_node(tracker, device_frees, host_frees) def evict_device_end(self) -> None: """Clear this component's device-eviction walk state.""" - assert ( - self.is_evict_device_ongoing - ), f"{self.component_type} device eviction not started" + assert self.is_evict_device_ongoing, ( + f"{self.component_type} device eviction not started" + ) self._evict_device_end() self.is_evict_device_ongoing = False diff --git a/python/sglang/srt/mem_cache/unified_cache/storage_attachment.py b/python/sglang/srt/mem_cache/unified_cache/storage_attachment.py index 2700c9b92..8a0d2e8bf 100644 --- a/python/sglang/srt/mem_cache/unified_cache/storage_attachment.py +++ b/python/sglang/srt/mem_cache/unified_cache/storage_attachment.py @@ -280,7 +280,6 @@ class StorageAttachment: existing_collector = cache.storage_metrics_collector if existing_collector is None: - storage_cls = resolve_collector_class( STAT_LOGGER_ROLE_STORAGE, StorageMetricsCollector, diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py index 70e4ea8b6..e25f7bc95 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py @@ -722,7 +722,9 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): action, ) - def _match_prefix_helper(self, key: RadixKey) -> tuple[ + def _match_prefix_helper( + self, key: RadixKey + ) -> tuple[ list[torch.Tensor], UnifiedTreeNode, UnifiedTreeNode, diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py index a497f34c7..213e7f3fa 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py @@ -35,9 +35,9 @@ class BaseEvictionResult(msgspec.Struct): def __del__(self) -> None: # Drop tripwire: every returned value must be drained before disposal. - assert ( - not self.device_frees and not self.host_frees - ), "BaseEvictionResult dropped with undrained values" + assert not self.device_frees and not self.host_frees, ( + "BaseEvictionResult dropped with undrained values" + ) class EvictDeviceNextNodeResult(BaseEvictionResult): diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index abe547450..5952534c0 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -135,9 +135,9 @@ class MHASubPoolSpec(SubPoolSpec): assert self.head_dim > 0, f"head_dim must be positive; got {self.head_dim}" if self.v_head_dim is None: object.__setattr__(self, "v_head_dim", self.head_dim) - assert ( - self.v_head_dim > 0 - ), f"v_head_dim must be positive; got {self.v_head_dim}" + assert self.v_head_dim > 0, ( + f"v_head_dim must be positive; got {self.v_head_dim}" + ) def k_row_bytes(self) -> int: return self.head_num * self.head_dim * self.store_dtype.itemsize @@ -190,12 +190,12 @@ class MLASubPoolSpec(SubPoolSpec): def __post_init__(self): super().__post_init__() - assert ( - self.kv_lora_rank > 0 - ), f"kv_lora_rank must be positive; got {self.kv_lora_rank}" - assert ( - self.qk_rope_head_dim > 0 - ), f"qk_rope_head_dim must be positive; got {self.qk_rope_head_dim}" + assert self.kv_lora_rank > 0, ( + f"kv_lora_rank must be positive; got {self.kv_lora_rank}" + ) + assert self.qk_rope_head_dim > 0, ( + f"qk_rope_head_dim must be positive; got {self.qk_rope_head_dim}" + ) @property def kv_cache_dim(self) -> int: @@ -308,13 +308,13 @@ class UnifiedKVPool: page_size: int = 1, ): assert page_size >= 1, f"page_size must be >= 1; got {page_size}" - assert ( - len(sub_pool_specs) >= 2 - ), f"UnifiedKVPool needs >= 2 sub-pools; got {len(sub_pool_specs)}" + assert len(sub_pool_specs) >= 2, ( + f"UnifiedKVPool needs >= 2 sub-pools; got {len(sub_pool_specs)}" + ) names = [s.name for s in sub_pool_specs] - assert len(set(names)) == len( - names - ), f"sub-pool names must be unique; got {names}" + assert len(set(names)) == len(names), ( + f"sub-pool names must be unique; got {names}" + ) # Per-spec direction validity already ran in each spec's __post_init__. up_specs = [s for s in sub_pool_specs if s.grow_direction == "up"] down_specs = [s for s in sub_pool_specs if s.grow_direction == "down"] @@ -440,23 +440,23 @@ class UnifiedKVPool: def mha_spec(self, name: str) -> MHASubPoolSpec: s = self._specs_by_name[name] - assert isinstance( - s, MHASubPoolSpec - ), f"sub-pool {name!r} is {type(s).__name__}, expected MHASubPoolSpec" + assert isinstance(s, MHASubPoolSpec), ( + f"sub-pool {name!r} is {type(s).__name__}, expected MHASubPoolSpec" + ) return s def mla_spec(self, name: str) -> MLASubPoolSpec: s = self._specs_by_name[name] - assert isinstance( - s, MLASubPoolSpec - ), f"sub-pool {name!r} is {type(s).__name__}, expected MLASubPoolSpec" + assert isinstance(s, MLASubPoolSpec), ( + f"sub-pool {name!r} is {type(s).__name__}, expected MLASubPoolSpec" + ) return s def mamba_spec(self, name: str) -> MambaSubPoolSpec: s = self._specs_by_name[name] - assert isinstance( - s, MambaSubPoolSpec - ), f"sub-pool {name!r} is {type(s).__name__}, expected MambaSubPoolSpec" + assert isinstance(s, MambaSubPoolSpec), ( + f"sub-pool {name!r} is {type(s).__name__}, expected MambaSubPoolSpec" + ) return s def max_slots(self, name: str) -> int: @@ -812,12 +812,12 @@ class UnifiedMambaPool(MambaPool): self.conv_shard_groups = None self.conv_slice_axis = spec.conv_slice_axis - assert ( - conv_views[0].shape[0] == self.num_mamba_layers - ), f"conv_views layers={conv_views[0].shape[0]} vs expected {self.num_mamba_layers}" - assert ( - conv_views[0].shape[1] == self._max_size + 1 - ), f"conv_views slots={conv_views[0].shape[1]} vs expected {self._max_size + 1}" + assert conv_views[0].shape[0] == self.num_mamba_layers, ( + f"conv_views layers={conv_views[0].shape[0]} vs expected {self.num_mamba_layers}" + ) + assert conv_views[0].shape[1] == self._max_size + 1, ( + f"conv_views slots={conv_views[0].shape[1]} vs expected {self._max_size + 1}" + ) # Per-draft-token intermediate buffers have a different outer size # (spec_state_size+1), so they're NOT in the shared buffer; allocate locally. @@ -1683,12 +1683,12 @@ def init_unified_swa_pools( # Both sub-allocators are page-aware: one virtual ID space at PAGE granularity, # two physical sub-pools compacting pages independently. assert page_size >= 1, f"page_size must be >= 1, got {page_size}" - assert ( - len(full_attention_layer_ids) > 0 - ), "SWA-hybrid with zero full-attention layers is degenerate" - assert ( - len(swa_attention_layer_ids) > 0 - ), "SWA-hybrid with zero SWA-attention layers is degenerate" + assert len(full_attention_layer_ids) > 0, ( + "SWA-hybrid with zero full-attention layers is degenerate" + ) + assert len(swa_attention_layer_ids) > 0, ( + "SWA-hybrid with zero SWA-attention layers is degenerate" + ) store_dtype = _store_dtype_for(kv_cache_dtype) # full-attn at the high-byte end (grow-down), swa at the low-byte end (grow-up). @@ -1865,12 +1865,12 @@ def init_unified_mamba_swa_pools( ) assert page_size >= 1, f"page_size must be >= 1, got {page_size}" - assert ( - len(full_attention_layer_ids) > 0 - ), "tri-pool with zero full-attention layers is degenerate" - assert ( - len(swa_attention_layer_ids) > 0 - ), "tri-pool with zero SWA-attention layers is degenerate" + assert len(full_attention_layer_ids) > 0, ( + "tri-pool with zero full-attention layers is degenerate" + ) + assert len(swa_attention_layer_ids) > 0, ( + "tri-pool with zero SWA-attention layers is degenerate" + ) assert len(mamba_layer_ids) > 0, "tri-pool with zero state layers is degenerate" store_dtype = _store_dtype_for(kv_cache_dtype) diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 1e24a987f..77e556914 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -462,8 +462,8 @@ class UnifiedRadixCache(BasePrefixCache): # content would drop-newest and punch storage holes. write_backlog_cap=2 * self.token_to_kv_pool_allocator.size_full, ) - self.cache_controller.host_write_staged_tokens_fn = ( - lambda: self.buffer_pipeline.write_staged_tokens_ + self.cache_controller.host_write_staged_tokens_fn = lambda: ( + self.buffer_pipeline.write_staged_tokens_ ) # State initialization @@ -998,12 +998,12 @@ class UnifiedRadixCache(BasePrefixCache): new_indices = match_result.device_indices new_last_node = match_result.last_device_node new_prefix_len = result.prefix_len - assert ( - req.kv.cache_protected_len <= len(new_indices) + self.page_size - 1 - ), f"{req.kv.cache_protected_len=}, {len(new_indices)=}, {page_aligned_len=}" - assert new_prefix_len <= len( - new_indices - ), f"{new_prefix_len=}, {len(new_indices)=}" + assert req.kv.cache_protected_len <= len(new_indices) + self.page_size - 1, ( + f"{req.kv.cache_protected_len=}, {len(new_indices)=}, {page_aligned_len=}" + ) + assert new_prefix_len <= len(new_indices), ( + f"{new_prefix_len=}, {len(new_indices)=}" + ) self.req_to_token_pool.write( (req.kv.req_pool_idx, slice(req.kv.cache_protected_len, len(new_indices))), new_indices[req.kv.cache_protected_len :], @@ -1197,9 +1197,9 @@ class UnifiedRadixCache(BasePrefixCache): req.kv.req_pool_idx, window_start:num_tokens ].to(torch.int64) swa_indices = kv_cache.translate_loc_from_full_to_swa(window_indices) - assert bool( - (swa_indices > 0).all() - ), f"unmapped SWA window positions for request {req.rid}" + assert bool((swa_indices > 0).all()), ( + f"unmapped SWA window positions for request {req.rid}" + ) component_transfers[ComponentType.SWA] = [ PoolTransfer( name=PoolName.SWA, @@ -2611,9 +2611,9 @@ class UnifiedRadixCache(BasePrefixCache): self._all_reduce(ready_counts, torch.distributed.ReduceOp.MIN) count_values = list(map(int, ready_counts.tolist())) - assert ( - count_values[-2] == -count_values[-1] - ), "write_back duplicate-reclaim victims diverged across TP ranks" + assert count_values[-2] == -count_values[-1], ( + "write_back duplicate-reclaim victims diverged across TP ranks" + ) return ( count_values[0], count_values[1], @@ -2697,9 +2697,9 @@ class UnifiedRadixCache(BasePrefixCache): ) self._all_reduce(sync_tensor, torch.distributed.ReduceOp.MIN) finish_count = int(sync_tensor[0].item()) - assert ( - sync_tensor[1].item() == -sync_tensor[2].item() - ), "write_back duplicate-reclaim victims diverged across TP ranks" + assert sync_tensor[1].item() == -sync_tensor[2].item(), ( + "write_back duplicate-reclaim victims diverged across TP ranks" + ) while finish_count > 0: ack = cc.ack_load_queue.pop(0) diff --git a/python/sglang/srt/model_executor/cpu_graph_runner.py b/python/sglang/srt/model_executor/cpu_graph_runner.py index d2a0a72d7..6b1e34cbe 100644 --- a/python/sglang/srt/model_executor/cpu_graph_runner.py +++ b/python/sglang/srt/model_executor/cpu_graph_runner.py @@ -137,9 +137,9 @@ def get_batch_sizes_to_capture(model_runner: ModelRunner): # Users can customize the batch sizes supported by cpu_graph, such as: # --cuda-graph-bs-decode 1 2 4 8 16 capture_bs = get_exec().graph.cuda_graph_config.decode.bs - assert ( - max(capture_bs) <= get_exec().graph.torch_compile_max_bs - ), f"{capture_bs=}, {get_exec().graph.torch_compile_max_bs=}" + assert max(capture_bs) <= get_exec().graph.torch_compile_max_bs, ( + f"{capture_bs=}, {get_exec().graph.torch_compile_max_bs=}" + ) capture_bs = [bs for bs in capture_bs if bs <= model_runner.req_to_token_pool.size] capture_bs = list(sorted(set(capture_bs))) assert len(capture_bs) > 0 and capture_bs[0] > 0, f"{capture_bs=}" @@ -619,21 +619,21 @@ class CPUGraphRunner: self.captured_req_width = 1 assert not get_lora().enable_lora, "CPUGraphRunner does not support LoRA yet." - assert ( - not self.enable_two_batch_overlap - ), "CPUGraphRunner does not support two batch overlap yet." - assert ( - not self.require_mlp_tp_gather - ), "CPUGraphRunner does not support MLP TP gather yet." - assert ( - not self.require_mlp_sync - ), "CPUGraphRunner does not support MLP sync yet." - assert ( - not self.require_gathered_buffer - ), "CPUGraphRunner does not support gathered buffer yet." - assert ( - model_runner.spec_algorithm.is_none() - ), "CPUGraphRunner does not support speculative inference yet." + assert not self.enable_two_batch_overlap, ( + "CPUGraphRunner does not support two batch overlap yet." + ) + assert not self.require_mlp_tp_gather, ( + "CPUGraphRunner does not support MLP TP gather yet." + ) + assert not self.require_mlp_sync, ( + "CPUGraphRunner does not support MLP sync yet." + ) + assert not self.require_gathered_buffer, ( + "CPUGraphRunner does not support gathered buffer yet." + ) + assert model_runner.spec_algorithm.is_none(), ( + "CPUGraphRunner does not support speculative inference yet." + ) assert self.dp_size == 1, "CPUGraphRunner does not support DP yet." assert self.pp_size == 1, "CPUGraphRunner does not support PP yet." @@ -952,9 +952,9 @@ class CPUGraphRunner: forward_batch: ForwardBatch, pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> Union[LogitsProcessorOutput, PPProxyTensors]: - assert ( - pp_proxy_tensors is None - ), "PPProxyTensors is not supported in CPUGraphRunner yet." + assert pp_proxy_tensors is None, ( + "PPProxyTensors is not supported in CPUGraphRunner yet." + ) replay_context = ( model_capture_mode if self.is_encoder_decoder else empty_context diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 3926aa58c..f98d5855c 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -1420,9 +1420,9 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): # branch handles decode rows padded to a 1-token extend. if hybrid_ssm or self.seq_lens.shape[0] == 0: dev = self.seq_lens.device - assert ( - self.seq_lens.shape[0] == 0 - ), "extend-idle conversion expects an empty rank" + assert self.seq_lens.shape[0] == 0, ( + "extend-idle conversion expects an empty rank" + ) self.extend_num_tokens = num_tokens self.extend_seq_lens = torch.tensor( [num_tokens], dtype=torch.int32, device=dev diff --git a/python/sglang/srt/model_executor/hook_manager.py b/python/sglang/srt/model_executor/hook_manager.py index 5f6d78a51..54f007d3c 100644 --- a/python/sglang/srt/model_executor/hook_manager.py +++ b/python/sglang/srt/model_executor/hook_manager.py @@ -46,14 +46,13 @@ def register_forward_hooks(model: nn.Module, hook_specs: List[dict[str, Any]]) - if not matched: logger.warning( - f"No modules matched hook spec '{spec_name}' " - f"patterns={target_patterns}" + f"No modules matched hook spec '{spec_name}' patterns={target_patterns}" ) continue for module_name, module in matched: _ = module.register_forward_hook(hook) - logger.info(f"Registered forward hook '{spec_name}' " f"on {module_name}") + logger.info(f"Registered forward hook '{spec_name}' on {module_name}") def resolve_callable(path: Optional[str]) -> Optional[Callable]: diff --git a/python/sglang/srt/model_executor/input_buffers.py b/python/sglang/srt/model_executor/input_buffers.py index 6d1a2d31c..cbbc89de4 100644 --- a/python/sglang/srt/model_executor/input_buffers.py +++ b/python/sglang/srt/model_executor/input_buffers.py @@ -58,7 +58,6 @@ INDEX_SEMANTIC_BUFFERS = frozenset( @dataclass class ForwardInputBuffers: - def reset_index_buffers(self) -> None: """Zero the index-semantic buffers this set declares.""" for f in fields(self): @@ -87,14 +86,14 @@ class ForwardInputBuffers: if isinstance(buffer, dict): for sub_name, sub_buffer in buffer.items(): - assert isinstance( - sub_buffer, torch.Tensor - ), f"Field {name}.{sub_name} is expected to be a torch.Tensor, but got {type(sub_buffer)}." + assert isinstance(sub_buffer, torch.Tensor), ( + f"Field {name}.{sub_name} is expected to be a torch.Tensor, but got {type(sub_buffer)}." + ) buffer[sub_name] = share_input_buffer( f"{name}.{sub_name}", sub_buffer ) else: - assert isinstance( - buffer, torch.Tensor - ), f"Field {name} is expected to be a torch.Tensor, a dict of torch.Tensor, or a dataclass of torch.Tensor, but got {type(buffer)}." + assert isinstance(buffer, torch.Tensor), ( + f"Field {name} is expected to be a torch.Tensor, a dict of torch.Tensor, or a dataclass of torch.Tensor, but got {type(buffer)}." + ) setattr(self, name, share_input_buffer(name, buffer)) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index d6d74ebb8..ace5fd98a 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -310,8 +310,7 @@ class ModelRunner: def sampling_observer(self, observer: Optional[SamplingObserver]) -> None: if observer is not None and not self.supports_sampling_observer(): raise ValueError( - "sampling observers are not supported by the configured " - "sampling path" + "sampling observers are not supported by the configured sampling path" ) self._sampling_observer = observer @@ -481,9 +480,9 @@ class ModelRunner: ) if self.ps.pp_size > 1: - assert ( - self.support_pp - ), "Pipeline Parallel is not compatible with this model." + assert self.support_pp, ( + "Pipeline Parallel is not compatible with this model." + ) # For weight updates self.init_weight_updater() diff --git a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py b/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py index 5cb1b72b1..56d51d7d1 100644 --- a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py +++ b/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py @@ -58,23 +58,23 @@ logger = logging.getLogger(__name__) def _align_pipeline_layers(layers: list, layer_model) -> list: has_start_layer = hasattr(layer_model, "start_layer") has_end_layer = hasattr(layer_model, "end_layer") - assert ( - has_start_layer == has_end_layer - ), "pipeline layer ranges must define start_layer and end_layer together" + assert has_start_layer == has_end_layer, ( + "pipeline layer ranges must define start_layer and end_layer together" + ) start_layer = layer_model.start_layer if has_start_layer else 0 end_layer = layer_model.end_layer if has_end_layer else len(layer_model.layers) - assert isinstance(start_layer, int) and isinstance( - end_layer, int - ), "pipeline layer ranges must define integer start_layer and end_layer" + assert isinstance(start_layer, int) and isinstance(end_layer, int), ( + "pipeline layer ranges must define integer start_layer and end_layer" + ) assert 0 <= start_layer <= end_layer <= len(layer_model.layers), ( f"invalid pipeline layer range [{start_layer}, {end_layer}) for " f"{len(layer_model.layers)} layers" ) if len(layers) == len(layer_model.layers): return layers - assert ( - len(layers) <= end_layer - start_layer - ), f"found {len(layers)} layers in PP range [{start_layer}, {end_layer})" + assert len(layers) <= end_layer - start_layer, ( + f"found {len(layers)} layers in PP range [{start_layer}, {end_layer})" + ) return ( [None] * start_layer + layers + [None] * (len(layer_model.layers) - end_layer) ) diff --git a/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py b/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py index 0ea3c517c..e3c22d6f9 100644 --- a/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py +++ b/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py @@ -93,8 +93,10 @@ def maybe_trigger_remote_instance_nccl_send_group( ``--speculative-draft-draft-load-format`` needs its own send group, and the target's format cannot answer for it.""" if ( - load_format or get_model().load_format - ) == LoadFormat.REMOTE_INSTANCE and get_model().remote_instance_weight_loader_backend == RemoteInstanceWeightLoaderBackend.NCCL: + (load_format or get_model().load_format) == LoadFormat.REMOTE_INSTANCE + and get_model().remote_instance_weight_loader_backend + == RemoteInstanceWeightLoaderBackend.NCCL + ): if tp_rank == 0: instance_ip = NetworkAddress.resolve_host(socket.gethostname()) t = threading.Thread( diff --git a/python/sglang/srt/model_executor/model_runner_components/ngram_embedding_manager.py b/python/sglang/srt/model_executor/model_runner_components/ngram_embedding_manager.py index d4054cce1..3f51035fd 100644 --- a/python/sglang/srt/model_executor/model_runner_components/ngram_embedding_manager.py +++ b/python/sglang/srt/model_executor/model_runner_components/ngram_embedding_manager.py @@ -50,9 +50,9 @@ class NgramEmbeddingManager: device=device, ) chunked_prefill_size = get_schedule().chunked_prefill_size - assert ( - chunked_prefill_size is not None and chunked_prefill_size > 0 - ), "Ngram embedding requires chunked prefill to be enabled (chunked_prefill_size > 0)" + assert chunked_prefill_size is not None and chunked_prefill_size > 0, ( + "Ngram embedding requires chunked prefill to be enabled (chunked_prefill_size > 0)" + ) for module in model.modules(): if isinstance(module, NgramEmbedding): module.init_buffers( diff --git a/python/sglang/srt/model_executor/model_runner_components/startup_weight_load.py b/python/sglang/srt/model_executor/model_runner_components/startup_weight_load.py index 77c844d41..856281b73 100644 --- a/python/sglang/srt/model_executor/model_runner_components/startup_weight_load.py +++ b/python/sglang/srt/model_executor/model_runner_components/startup_weight_load.py @@ -520,8 +520,7 @@ class StartupWeightLoadManager: if changed_names: preview = ", ".join(changed_names[:8]) raise RuntimeError( - "Startup weight commit changed graph-visible tensor storage: " - f"{preview}" + f"Startup weight commit changed graph-visible tensor storage: {preview}" ) unchanged_names = manifest.unchanged_parameter_names( CAPTURE_SAFE_WEIGHT_SENTINEL diff --git a/python/sglang/srt/model_executor/model_runner_components/weight_exporter.py b/python/sglang/srt/model_executor/model_runner_components/weight_exporter.py index 5c155db89..d9f181e56 100644 --- a/python/sglang/srt/model_executor/model_runner_components/weight_exporter.py +++ b/python/sglang/srt/model_executor/model_runner_components/weight_exporter.py @@ -32,15 +32,15 @@ class WeightExporter: group_name, backend="nccl", ): - assert ( - torch.distributed.is_initialized() - ), "Default torch process group must be initialized" + assert torch.distributed.is_initialized(), ( + "Default torch process group must be initialized" + ) assert group_name != "", "Group name cannot be empty" ports_list = ports.split(",") - assert ( - len(ports_list) == self.tp_size - ), f"Expected {self.tp_size} ports, but got {len(ports_list)} ports." + assert len(ports_list) == self.tp_size, ( + f"Expected {self.tp_size} ports, but got {len(ports_list)} ports." + ) group_port = ports_list[self.tp_rank] group_name = f"{group_name}_{group_port}_{self.tp_rank}" @@ -78,15 +78,15 @@ class WeightExporter: ports, group_name, ): - assert ( - torch.distributed.is_initialized() - ), "Default torch process group must be initialized" + assert torch.distributed.is_initialized(), ( + "Default torch process group must be initialized" + ) assert group_name != "", "Group name cannot be empty" ports_list = ports.split(",") - assert ( - len(ports_list) == self.tp_size - ), f"Expected {self.tp_size} ports, but got {len(ports_list)} ports." + assert len(ports_list) == self.tp_size, ( + f"Expected {self.tp_size} ports, but got {len(ports_list)} ports." + ) group_port = ports_list[self.tp_rank] group_name = f"{group_name}_{group_port}_{self.tp_rank}" diff --git a/python/sglang/srt/model_executor/model_runner_components/weight_updater.py b/python/sglang/srt/model_executor/model_runner_components/weight_updater.py index ecbfa7a71..1668c966b 100644 --- a/python/sglang/srt/model_executor/model_runner_components/weight_updater.py +++ b/python/sglang/srt/model_executor/model_runner_components/weight_updater.py @@ -83,9 +83,9 @@ class WeightUpdater: weights/parameters online, and broadcasts them to the inference engine through the `_model_update_group` process group. """ - assert ( - torch.distributed.is_initialized() - ), "Default torch process group must be initialized" + assert torch.distributed.is_initialized(), ( + "Default torch process group must be initialized" + ) assert group_name != "", "Group name cannot be empty" rank = rank_offset + self.tp_rank diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 8bfd033b6..305225d9d 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -462,9 +462,9 @@ class HybridSWAPoolConfigurator(MemoryPoolConfigurator): self._full_layers_num = len(model_config.full_attention_layer_ids) self._swa_layers_num = len(model_config.swa_attention_layer_ids) - assert ( - self._swa_layers_num > 0 - ), "Hybrid SWA model must have at least one SWA layer" + assert self._swa_layers_num > 0, ( + "Hybrid SWA model must have at least one SWA layer" + ) self._swa_full_tokens_ratio = get_schedule().swa_full_tokens_ratio self._sliding_window_size = kvc.sliding_window_size @@ -1014,9 +1014,9 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): def calculate_pool_sizes( self, available_bytes: int, page_size: int ) -> MemoryPoolConfig: - assert ( - page_size % 128 == 0 - ), "page_size must be multiple of 128 for compressed attention" + assert page_size % 128 == 0, ( + "page_size must be multiple of 128 for compressed attention" + ) if self.requested_max_running_requests_per_worker is not None: c128_state_fixed_bytes = self._get_c128_state_fixed_bytes( @@ -1044,9 +1044,9 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): def calculate_pool_sizes_from_max_tokens( self, max_total_num_tokens: int, page_size: int ) -> MemoryPoolConfig: - assert ( - page_size % 128 == 0 - ), "page_size must be multiple of 128 for compressed attention" + assert page_size % 128 == 0, ( + "page_size must be multiple of 128 for compressed attention" + ) sizes = self._compute_dsv4_sizes(max_total_num_tokens, page_size) return self._to_config(sizes) diff --git a/python/sglang/srt/model_executor/runner/base_runner.py b/python/sglang/srt/model_executor/runner/base_runner.py index b27766b51..d45ebfb82 100644 --- a/python/sglang/srt/model_executor/runner/base_runner.py +++ b/python/sglang/srt/model_executor/runner/base_runner.py @@ -254,9 +254,9 @@ class BaseRunner(ABC): if should_run_flashinfer_autotune(self.model_runner): buffers, batch_size = self._autotune_buffers() - assert ( - buffers is not None - ), "_autotune_buffers() must return a reusable buffer set for autotune" + assert buffers is not None, ( + "_autotune_buffers() must return a reusable buffer set for autotune" + ) self._flashinfer_autotune(buffers=buffers, batch_size=batch_size) maybe_flashinfer_autotune_extend(self, decode_num_tokens=batch_size) @@ -427,9 +427,9 @@ class BaseRunner(ABC): ) if mr.spec_algorithm.is_speculative() and not _is_pd_prefill_target: if mr.is_draft_worker: - assert ( - mr.spec_algorithm.supports_target_verify_for_draft() - ), "This should not happen" + assert mr.spec_algorithm.supports_target_verify_for_draft(), ( + "This should not happen" + ) capture_forward_mode = ForwardMode.TARGET_VERIFY num_tokens_per_req = mr.decode_num_tokens_per_req() if extend_num_tokens_per_req is not None: diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index ae9df55fc..82677b901 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -326,9 +326,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): if model_runner.spec_algorithm.is_speculative(): if self.model_runner.is_draft_worker: # Draft workers can use TARGET_VERIFY mode. - if ( - not self.model_runner.spec_algorithm.supports_target_verify_for_draft() - ): + if not self.model_runner.spec_algorithm.supports_target_verify_for_draft(): raise RuntimeError("This should not happen") self.capture_forward_mode = ForwardMode.TARGET_VERIFY elif self.is_dllm: @@ -493,7 +491,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): self.capture() except RuntimeError as e: raise Exception( - f"Capture cuda graph failed: {e}\n" f"{CUDA_GRAPH_CAPTURE_FAILED_MSG}" + f"Capture cuda graph failed: {e}\n{CUDA_GRAPH_CAPTURE_FAILED_MSG}" ) def _record_in_graph_metadata_prep_done(self): @@ -1167,9 +1165,9 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): # Sanity-check: --debug-cuda-graph requires breakable backend. if get_exec().graph.debug_cuda_graph: - assert isinstance( - self.backend, BreakableCudaGraphBackend - ), "Breakable CUDA graph is required for --debug-cuda-graph" + assert isinstance(self.backend, BreakableCudaGraphBackend), ( + "Breakable CUDA graph is required for --debug-cuda-graph" + ) forward_batch, attn_backend, pp_proxy_tensors = self.capture_prepare( bs, stream_idx=stream_idx, num_tokens=num_tokens @@ -1529,7 +1527,6 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): if self.model_runner.is_draft_worker: raise RuntimeError("This should not happen.") else: - capture_mode = ( CaptureHiddenMode.NULL if self.model_runner.spec_algorithm.is_standalone() diff --git a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py index 5c56e55c1..51fa74ecd 100644 --- a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py +++ b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py @@ -277,11 +277,15 @@ def flashinfer_autotune_context(model_runner: ModelRunner, *, run_lm_head: bool) from sglang.srt.layers.logits_processor import autotune_dummy_run_mode skip_ops = get_flashinfer_autotune_skip_ops(mr) - with _autotune_process_group(sync_group), autotune( - True, - cache=str(autotune_cache), - skip_ops=skip_ops, - ), autotune_dummy_run_mode(run_lm_head=run_lm_head): + with ( + _autotune_process_group(sync_group), + autotune( + True, + cache=str(autotune_cache), + skip_ops=skip_ops, + ), + autotune_dummy_run_mode(run_lm_head=run_lm_head), + ): yield torch.cuda.current_stream().wait_stream(mr.forward_stream) logger.info("FlashInfer autotune completed.") diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index c6b06b535..5e2a7f7aa 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -543,9 +543,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): # the contract only when the backend is Breakable; FullCG and # TC_PIECEWISE use the eager init_forward_metadata path. if isinstance(self.backend, BreakableCudaGraphBackend): - self.use_captured_attn_metadata = ( - model_runner.attn_backend.use_captured_forward_metadata_for_breakable_cuda_graph - ) + self.use_captured_attn_metadata = model_runner.attn_backend.use_captured_forward_metadata_for_breakable_cuda_graph else: self.use_captured_attn_metadata = False self.attn_metadata_buffers: Optional[Dict[int, object]] = ( @@ -625,9 +623,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): def _prefill_logits_buffer_rows(self, forward_batch: ForwardBatch) -> int: if not forward_batch.return_logprob: return forward_batch.batch_size - assert ( - self._uses_eager_prefill_tail() - ), "Prefill return_logprob requires an eager logits tail." + assert self._uses_eager_prefill_tail(), ( + "Prefill return_logprob requires an eager logits tail." + ) global_num_tokens = forward_batch.global_num_tokens_for_logprob_cpu if global_num_tokens is not None: diff --git a/python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py b/python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py index d4da89192..368db2ebc 100644 --- a/python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py +++ b/python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py @@ -131,11 +131,14 @@ class BreakableCudaGraphBackend(DedupedCudaGraphMixin, BaseCudaGraphBackend): size = shape_key.size if self._shared_output_buffer is None: self._shared_output_buffer = self._alloc_full_buffer(warmup_out, size) - with graph_pool_capture_scope(), BreakableCUDAGraphCapture( - cuda_graph=graph, - pool=self._pool, - stream=self._capture_stream, - barrier_fn=self._tp_group.barrier, + with ( + graph_pool_capture_scope(), + BreakableCUDAGraphCapture( + cuda_graph=graph, + pool=self._pool, + stream=self._capture_stream, + barrier_fn=self._tp_group.barrier, + ), ): self._precarve.mint() out = captured_fn() diff --git a/python/sglang/srt/model_executor/runner_backend/cuda_graph_dedup_mixin.py b/python/sglang/srt/model_executor/runner_backend/cuda_graph_dedup_mixin.py index 819405a7c..5feaccd83 100644 --- a/python/sglang/srt/model_executor/runner_backend/cuda_graph_dedup_mixin.py +++ b/python/sglang/srt/model_executor/runner_backend/cuda_graph_dedup_mixin.py @@ -259,9 +259,9 @@ class DedupedCudaGraphRegistry: def replay(self, graph: DedupedCudaGraph, stream: int) -> None: assert cuda_rt is not None group = graph.group - assert ( - group is not None - ), "captured CUDA graph does not belong to this dedup state" + assert group is not None, ( + "captured CUDA graph does not belong to this dedup state" + ) raw_graph = graph.raw_graph graph_exec = group.graph_exec diff --git a/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py b/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py index 59f4a2e6b..6bda4fbc2 100644 --- a/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py +++ b/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py @@ -317,9 +317,9 @@ class BreakableCUDAGraphCapture: capture_error_mode: str = "global", barrier_fn: Callable[[], None] | None = None, ): - assert isinstance( - cuda_graph, BreakableCUDAGraph - ), "cuda_graph must be a BreakableCUDAGraph" + assert isinstance(cuda_graph, BreakableCUDAGraph), ( + "cuda_graph must be a BreakableCUDAGraph" + ) self.cuda_graph = cuda_graph self._pool = pool if pool is not None else (0, 0) self._stream = stream diff --git a/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/cuda_utils.py b/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/cuda_utils.py index df86e523e..4fb5e6885 100644 --- a/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/cuda_utils.py +++ b/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/cuda_utils.py @@ -33,8 +33,7 @@ def _cudaGetErrorString(error): def checkCudaErrors(result): if rt is None: raise RuntimeError( - "cuda.bindings is not available. " - "Install it with: pip install cuda-python" + "cuda.bindings is not available. Install it with: pip install cuda-python" ) if result[0] != rt.cudaError_t.cudaSuccess: raise RuntimeError( diff --git a/python/sglang/srt/model_loader/kimi_k3_gguf.py b/python/sglang/srt/model_loader/kimi_k3_gguf.py index 876c73c78..5e75a48bd 100644 --- a/python/sglang/srt/model_loader/kimi_k3_gguf.py +++ b/python/sglang/srt/model_loader/kimi_k3_gguf.py @@ -196,8 +196,9 @@ def kimi_k3_nonexpert_weights_iterator( f"duplicate Kimi-K3 target parameter: {type_name}" ) emitted.add(type_name) - yield type_name, torch.tensor( - int(tensor.tensor_type), dtype=torch.uint8 + yield ( + type_name, + torch.tensor(int(tensor.tensor_type), dtype=torch.uint8), ) emitted.add(runtime_name) yield runtime_name, value diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 3211f0c41..59f5ea01a 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -2894,7 +2894,6 @@ class BitsAndBytesModelLoader(BaseModelLoader): for weight_name, weight_tensor in self._hf_weight_iter( hf_weights_files, use_safetensors ): - if self._is_4bit_weight_name(weight_name): continue @@ -2918,7 +2917,6 @@ class BitsAndBytesModelLoader(BaseModelLoader): for weight_name, weight_tensor in self._hf_weight_iter( hf_weights_files, use_safetensors ): - if any( target_module in weight_name for target_module in self.target_modules ) and weight_name.endswith(".weight"): @@ -2928,7 +2926,6 @@ class BitsAndBytesModelLoader(BaseModelLoader): module in weight_name for module in self.column_parallel_weights_modules ): - total_size = weight_tensor.size(-1) start_index = total_size // tp_size * tp_rank end_index = total_size // tp_size * (tp_rank + 1) @@ -2989,7 +2986,7 @@ class BitsAndBytesModelLoader(BaseModelLoader): self.model_type = type(model).__name__ logger.info( - "Loading weights with BitsAndBytes quantization. " " May take a while ..." + "Loading weights with BitsAndBytes quantization. May take a while ..." ) quant_config = getattr(model_config.hf_config, "quantization_config", None) @@ -3001,8 +2998,7 @@ class BitsAndBytesModelLoader(BaseModelLoader): pre_quant = True else: raise ValueError( - f"BitsAndBytes loader does not support {quant_method} " - "quantization" + f"BitsAndBytes loader does not support {quant_method} quantization" ) # The quant_states in pre_quantized models cannot work with a split @@ -3533,7 +3529,7 @@ class RemoteModelLoader(BaseModelLoader): param_data = param_data.narrow(dim, 0, size) if tensor.shape != param_shape: logger.warning( - "loading tensor of shape %s into " "parameter '%s' of shape %s", + "loading tensor of shape %s into parameter '%s' of shape %s", tensor.shape, key, param_shape, diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index fb96ff210..fd659e7c9 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -1045,8 +1045,7 @@ def _prefetch_all_checkpoints( _prefetch_all() succeeded_event.set() logger.info( - "Rank %d: prefetching checkpoint files into page cache " - "finished in %.2fs", + "Rank %d: prefetching checkpoint files into page cache finished in %.2fs", local_rank, time.perf_counter() - start, ) @@ -1582,7 +1581,6 @@ def runai_safetensors_weights_iterator( device = device if is_distributed and is_cuda_alike() else "cpu" with SafetensorsStreamer() as streamer: - streamer.stream_files( hf_weights_files, device=device, @@ -1799,9 +1797,9 @@ class KVCacheQuantSchema(BaseModel): f"{len(layer_maps)}." ) for i in range(tp_size): - assert ( - i in self.scaling_factor - ), f"KV cache scales map for TP rank {i} not found." + assert i in self.scaling_factor, ( + f"KV cache scales map for TP rank {i} not found." + ) return self @model_validator(mode="after") @@ -1949,9 +1947,9 @@ def pad_loaded_weight(loaded_weight, output_dim, output_sizes): int(output_size / total_output_size * raw_output_size) for output_size in output_sizes ] - assert ( - sum(weight_split_size) == raw_output_size - ), f"Padding the loaded weight failed due to sizes are not divisible cleanly from {output_sizes} to {raw_output_size}" + assert sum(weight_split_size) == raw_output_size, ( + f"Padding the loaded weight failed due to sizes are not divisible cleanly from {output_sizes} to {raw_output_size}" + ) split_weight = loaded_weight.split_with_sizes(weight_split_size, dim=output_dim) for i, output_size in enumerate(output_sizes): diff --git a/python/sglang/srt/models/afmoe.py b/python/sglang/srt/models/afmoe.py index d9573e817..ed5521f77 100644 --- a/python/sglang/srt/models/afmoe.py +++ b/python/sglang/srt/models/afmoe.py @@ -78,7 +78,6 @@ def get_attention_sliding_window_size(config: PretrainedConfig) -> Optional[int] class AfmoeMLP(nn.Module): - def __init__( self, hidden_size: int, @@ -118,7 +117,6 @@ class AfmoeMLP(nn.Module): class AfmoeMoE(nn.Module): - @staticmethod def _custom_routing_function( hidden_states: torch.Tensor, @@ -300,7 +298,6 @@ class AfmoeMoE(nn.Module): class AfmoeAttention(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -421,7 +418,6 @@ class AfmoeAttention(nn.Module): class AfmoeDecoderLayer(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -499,7 +495,6 @@ class AfmoeDecoderLayer(nn.Module): class AfmoeModel(nn.Module): - fall_back_to_pt_during_load = False def __init__( @@ -555,7 +550,6 @@ class AfmoeModel(nn.Module): class AfmoeForCausalLM(nn.Module): - def __init__( self, config: PretrainedConfig, diff --git a/python/sglang/srt/models/apertus.py b/python/sglang/srt/models/apertus.py index a016bfdde..6cbc918df 100644 --- a/python/sglang/srt/models/apertus.py +++ b/python/sglang/srt/models/apertus.py @@ -394,7 +394,7 @@ class ApertusModel(nn.Module): layer_self_attn.attn.v_scale = scaling_factor else: raise RuntimeError( - "Self attention has no KV cache scaling " "factor attribute!" + "Self attention has no KV cache scaling factor attribute!" ) diff --git a/python/sglang/srt/models/baichuan.py b/python/sglang/srt/models/baichuan.py index b1968dad1..5e30c7335 100644 --- a/python/sglang/srt/models/baichuan.py +++ b/python/sglang/srt/models/baichuan.py @@ -80,7 +80,6 @@ def _get_alibi_slopes(total_num_heads: int) -> torch.Tensor: class BaiChuanMLP(nn.Module): - def __init__( self, hidden_size: int, @@ -106,8 +105,7 @@ class BaiChuanMLP(nn.Module): ) if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() @@ -223,7 +221,6 @@ class BaiChuanAttention(nn.Module): class BaiChuanDecoderLayer(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -284,7 +281,6 @@ class BaiChuanDecoderLayer(nn.Module): class BaiChuanModel(nn.Module): - def __init__( self, config: PretrainedConfig, diff --git a/python/sglang/srt/models/bailing_moe.py b/python/sglang/srt/models/bailing_moe.py index 3258d849d..96a564127 100644 --- a/python/sglang/srt/models/bailing_moe.py +++ b/python/sglang/srt/models/bailing_moe.py @@ -234,7 +234,9 @@ class BailingMoESparseMoeBlock(nn.Module): self.score_function == "softmax" and self.correction_bias is None ) or ( self.score_function == "sigmoid" and self.correction_bias is not None - ), "score_function and correction_bias should be in 2 combination (softmax, None) or (sigmoid, not None)" + ), ( + "score_function and correction_bias should be in 2 combination (softmax, None) or (sigmoid, not None)" + ) self.topk = TopK( top_k=self.top_k, @@ -684,7 +686,6 @@ class BailingMoEBlock(nn.Module): class BailingMoEModel(nn.Module): - def __init__( self, config: PretrainedConfig, diff --git a/python/sglang/srt/models/bailing_moe_linear.py b/python/sglang/srt/models/bailing_moe_linear.py index 779d3cc71..07da00d68 100644 --- a/python/sglang/srt/models/bailing_moe_linear.py +++ b/python/sglang/srt/models/bailing_moe_linear.py @@ -164,7 +164,6 @@ def weight_loader_with_alias(alias: str): class BailingMLP(nn.Module): - def __init__( self, hidden_size: int, @@ -233,7 +232,6 @@ class BailingMoEGate(nn.Module): class BailingMoE(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -296,7 +294,9 @@ class BailingMoE(nn.Module): self.score_function == "softmax" and self.correction_bias is None ) or ( self.score_function == "sigmoid" and self.correction_bias is not None - ), "score_function and correction_bias should be in 2 combination (softmax, None) or (sigmoid, not None)" + ), ( + "score_function and correction_bias should be in 2 combination (softmax, None) or (sigmoid, not None)" + ) self.topk = TopK( top_k=self.top_k, @@ -503,12 +503,12 @@ class BailingMoELinearAttention(nn.Module): self.group_norm_size = getattr(config, "group_norm_size", 1) self.rms_norm_eps = float(getattr(config, "rms_norm_eps", 1e-5)) - assert ( - self.tp_size <= self.group_norm_size - ), "tp_size must be less than or equal to group_norm_size that can use local rms norm" - assert ( - self.group_norm_size % self.tp_size == 0 - ), "group_norm_size must be divisible by tp_size" + assert self.tp_size <= self.group_norm_size, ( + "tp_size must be less than or equal to group_norm_size that can use local rms norm" + ) + assert self.group_norm_size % self.tp_size == 0, ( + "group_norm_size must be divisible by tp_size" + ) self.g_norm = BailingGroupRMSNormGate( hidden_size=self.hidden_inner_size // self.tp_size, eps=self.rms_norm_eps, @@ -622,7 +622,6 @@ class BailingMoELinearAttention(nn.Module): class BailingMoEAttention(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -731,7 +730,6 @@ class BailingMoEAttention(nn.Module): class BailingMoELinearDecoderLayer(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -918,7 +916,6 @@ class BailingMoELinearDecoderLayer(nn.Module): class BailingMoELinearModel(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -943,9 +940,9 @@ class BailingMoELinearModel(nn.Module): f"Layer config: {num_linear} linear attention layers, {num_full} full attention layers" ) - assert ( - self.num_layers % self.layer_group_size == 0 - ), f"num_layers={self.num_layers} must be divided by layer_group_size={self.layer_group_size}" + assert self.num_layers % self.layer_group_size == 0, ( + f"num_layers={self.num_layers} must be divided by layer_group_size={self.layer_group_size}" + ) if self.pp_group.is_first_rank: self.word_embeddings = VocabParallelEmbedding( @@ -1041,7 +1038,6 @@ class BailingMoELinearModel(nn.Module): class BailingMoELinearForCausalLM(nn.Module): - packed_modules_mapping = { "fused_qkv_a_proj_with_mqa": ["q_a_proj", "kv_a_proj_with_mqa"], "gate_up_proj": ["gate_proj", "up_proj"], @@ -1492,7 +1488,6 @@ class BailingMoELinearForCausalLM(nn.Module): weight_loader(param, loaded_weight, shard_id) break else: - for mapping in expert_params_mapping: param_name, weight_name, expert_id, shard_id = mapping if weight_name not in name: @@ -1514,7 +1509,6 @@ class BailingMoELinearForCausalLM(nn.Module): ) break else: - if name.endswith(".bias") and name not in params_dict: continue if "slope" in name: @@ -1571,7 +1565,6 @@ class BailingMoELinearForCausalLM(nn.Module): cached_a_proj.pop(q_a_proj_name) cached_a_proj.pop(kv_a_proj_name) else: - if name not in params_dict: name = name.replace(".dense.", ".o_proj.") if name not in params_dict: diff --git a/python/sglang/srt/models/bailing_moe_v3.py b/python/sglang/srt/models/bailing_moe_v3.py index 6b09ba954..a73c2e0d5 100644 --- a/python/sglang/srt/models/bailing_moe_v3.py +++ b/python/sglang/srt/models/bailing_moe_v3.py @@ -494,7 +494,9 @@ class BailingMoE(nn.Module): self.score_function == "softmax" and self.correction_bias is None ) or ( self.score_function == "sigmoid" and self.correction_bias is not None - ), "score_function and correction_bias should be in 2 combination (softmax, None) or (sigmoid, not None)" + ), ( + "score_function and correction_bias should be in 2 combination (softmax, None) or (sigmoid, not None)" + ) self._enable_a2a_moe = not get_moe_a2a_backend().is_none() @@ -1156,7 +1158,9 @@ class BailingMoELinearModel(nn.Module): assert ( isinstance(self.layer_group_size, list) or self.num_layers % self.layer_group_size == 0 - ), f"num_layers={self.num_layers} must be divided by layer_group_size={self.layer_group_size}" + ), ( + f"num_layers={self.num_layers} must be divided by layer_group_size={self.layer_group_size}" + ) if self.pp_group.is_first_rank: self.word_embeddings = VocabParallelEmbedding( diff --git a/python/sglang/srt/models/bert.py b/python/sglang/srt/models/bert.py index 154900912..714d0b43e 100644 --- a/python/sglang/srt/models/bert.py +++ b/python/sglang/srt/models/bert.py @@ -23,7 +23,6 @@ BertConfig = None class BertEmbedding(nn.Module): - def __init__(self, config: BertConfig): super().__init__() @@ -77,7 +76,6 @@ class BertEmbedding(nn.Module): class BertPooler(nn.Module): - def __init__(self, config: BertConfig): super().__init__() self.dense = nn.Linear(config.hidden_size, config.hidden_size) @@ -96,7 +94,6 @@ class BertPooler(nn.Module): class BertEncoder(nn.Module): - def __init__( self, config: BertConfig, @@ -127,7 +124,6 @@ class BertEncoder(nn.Module): class BertLayer(nn.Module): - def __init__( self, config: BertConfig, @@ -173,7 +169,6 @@ class BertLayer(nn.Module): class BertAttention(nn.Module): - def __init__( self, hidden_size: int, @@ -208,7 +203,6 @@ class BertAttention(nn.Module): class BertSelfAttention(nn.Module): - def __init__( self, hidden_size: int, @@ -264,7 +258,6 @@ class BertSelfAttention(nn.Module): class BertSelfOutput(nn.Module): - def __init__( self, hidden_size: int, @@ -291,7 +284,6 @@ class BertSelfOutput(nn.Module): class BertIntermediate(nn.Module): - def __init__( self, hidden_size: int, @@ -317,7 +309,6 @@ class BertIntermediate(nn.Module): class BertOutput(nn.Module): - def __init__( self, hidden_size: int, @@ -347,7 +338,6 @@ class BertOutput(nn.Module): class BertModel(nn.Module): - def __init__( self, *, @@ -411,7 +401,6 @@ class BertModel(nn.Module): if not self.use_bert_pooler and "pooler" in name: continue for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: continue name = name.replace(weight_name, param_name) @@ -436,7 +425,6 @@ class Contriever(BertModel): class BertForSequenceClassification(nn.Module): - def __init__( self, *, diff --git a/python/sglang/srt/models/clip.py b/python/sglang/srt/models/clip.py index b93a0f714..3f97a3ea3 100644 --- a/python/sglang/srt/models/clip.py +++ b/python/sglang/srt/models/clip.py @@ -49,7 +49,6 @@ def prepare_clip_attention_mask( class CLIPVisionEmbeddings(nn.Module): - def __init__(self, config: CLIPVisionConfig): super().__init__() self.config = config @@ -141,7 +140,6 @@ class CLIPTextEmbeddings(nn.Module): class CLIPMLP(nn.Module): - def __init__( self, config, @@ -240,7 +238,6 @@ class CLIPAttention(nn.Module): class CLIPEncoderLayer(nn.Module): - def __init__( self, config: CLIPVisionConfig, @@ -428,7 +425,6 @@ class CLIPTextModel(nn.Module): class CLIPVisionTransformer(nn.Module): - def __init__( self, config: CLIPVisionConfig, diff --git a/python/sglang/srt/models/deepseek.py b/python/sglang/srt/models/deepseek.py index cce2d5aee..6bb1996b6 100644 --- a/python/sglang/srt/models/deepseek.py +++ b/python/sglang/srt/models/deepseek.py @@ -67,7 +67,6 @@ else: class DeepseekMLP(nn.Module): - def __init__( self, hidden_size: int, @@ -95,8 +94,7 @@ class DeepseekMLP(nn.Module): ) if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() @@ -108,7 +106,6 @@ class DeepseekMLP(nn.Module): class DeepseekMoE(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -228,7 +225,6 @@ class DeepseekMoE(nn.Module): class DeepseekAttention(nn.Module): - def __init__( self, hidden_size: int, @@ -314,7 +310,6 @@ class DeepseekAttention(nn.Module): class DeepseekDecoderLayer(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -386,7 +381,6 @@ class DeepseekDecoderLayer(nn.Module): class DeepseekModel(nn.Module): - fall_back_to_pt_during_load = False def __init__( @@ -440,7 +434,6 @@ class DeepseekModel(nn.Module): class DeepseekForCausalLM(nn.Module): - def __init__( self, config: PretrainedConfig, diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py index 41ebed0ff..707addf2b 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py @@ -486,9 +486,9 @@ class DeepseekMHAForwardMixin: if isinstance(backend, TboAttnBackend): # if enable tbo, get primary backend backend = backend.primary kv_indices = backend.forward_metadata.page_table_1_flattened - assert ( - kv_indices is not None - ), "page_table_1_flattened should have been generated for FP8 MHA path" + assert kv_indices is not None, ( + "page_table_1_flattened should have been generated for FP8 MHA path" + ) if _use_aiter_gfx95: # ROCm (gfx950) stores the FP8 MLA KV in the raw diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha_rocm.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha_rocm.py index 1ae596e02..6213cfb86 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha_rocm.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha_rocm.py @@ -53,7 +53,6 @@ if _use_aiter_gfx95: class DeepseekMHARocmForwardMixin: - def forward_normal_rocm_prepare( self: DeepseekV2AttentionMLA, positions: torch.Tensor, diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py index 34c0369b9..c32b314a2 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py @@ -17,7 +17,6 @@ if TYPE_CHECKING: class DeepseekMLACpuForwardMixin: - def init_mla_fused_rope_cpu_forward(self: DeepseekV2AttentionMLA): assert hasattr(self, "has_fused_proj") and hasattr(self, "is_packed_weight") @@ -61,9 +60,7 @@ class DeepseekMLACpuForwardMixin: self.fused_qkv_a_proj_with_mqa.quant_method.quant_config.weight_block_size == self.q_b_proj.quant_method.quant_config.weight_block_size ) - self.weight_block_size = ( - self.fused_qkv_a_proj_with_mqa.quant_method.quant_config.weight_block_size - ) + self.weight_block_size = self.fused_qkv_a_proj_with_mqa.quant_method.quant_config.weight_block_size def forward_absorb_fused_mla_rope_cpu_prepare( self: DeepseekV2AttentionMLA, @@ -72,9 +69,9 @@ class DeepseekMLACpuForwardMixin: forward_batch: ForwardBatch, zero_allocator: BumpAllocator, ): - assert self.q_lora_rank is not None and use_intel_amx_backend( - self - ), "forward_absorb_fused_mla_rope_cpu_prepare requires q_lora_rank is not None and use_intel_amx_backend" + assert self.q_lora_rank is not None and use_intel_amx_backend(self), ( + "forward_absorb_fused_mla_rope_cpu_prepare requires q_lora_rank is not None and use_intel_amx_backend" + ) q_input, k_input, v_input = ( torch.ops.sgl_kernel.qkv_proj_with_rope_fused_weight( @@ -126,9 +123,9 @@ class DeepseekMLACpuForwardMixin: zero_allocator, gate=None, ): - assert self.q_lora_rank is not None and use_intel_amx_backend( - self - ), "forward_absorb_fused_mla_rope_cpu_core requires q_lora_rank is not None and use_intel_amx_backend" + assert self.q_lora_rank is not None and use_intel_amx_backend(self), ( + "forward_absorb_fused_mla_rope_cpu_core requires q_lora_rank is not None and use_intel_amx_backend" + ) attn_output = self.attn_mqa(q_input, k_input, v_input, forward_batch) attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank) diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py index 01fd424c1..98d367bc8 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py @@ -30,7 +30,6 @@ if _is_hip: class DeepseekMLAFusedRopeRocmForwardMixin: - def init_mla_fused_rope_rocm_forward(self: DeepseekV2AttentionMLA): self.rocm_fused_decode_mla = get_bool_env_var( "SGLANG_ROCM_FUSED_DECODE_MLA", "false" diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py index e7fe7f7a5..667d54856 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py @@ -341,7 +341,6 @@ def _fused_rope_cat_and_cache( class DeepseekMLARocmForwardMixin: - def forward_absorb_rocm_prepare( self: DeepseekV2AttentionMLA, positions: torch.Tensor, diff --git a/python/sglang/srt/models/deepseek_janus_pro.py b/python/sglang/srt/models/deepseek_janus_pro.py index d8298e61f..5aafa026f 100644 --- a/python/sglang/srt/models/deepseek_janus_pro.py +++ b/python/sglang/srt/models/deepseek_janus_pro.py @@ -887,9 +887,9 @@ class VisionTransformer(nn.Module): if global_pool is not None: assert global_pool in ("", "avg", "token", "map") if global_pool == "map" and self.attn_pool is None: - assert ( - False - ), "Cannot currently add attention pooling in reset_classifier()." + assert False, ( + "Cannot currently add attention pooling in reset_classifier()." + ) elif global_pool != "map " and self.attn_pool is not None: self.attn_pool = None # remove attention pooling self.global_pool = global_pool @@ -987,7 +987,6 @@ def model_name_to_cls(cls_name): cls = CLIPVisionTower elif "VQ" in cls_name: - cls = VQ_models[cls_name] elif "vision_head" in cls_name: cls = vision_head @@ -1056,9 +1055,9 @@ def create_siglip_vit( ckpt_path: str = "", **kwargs, ): - assert ( - model_name in SigLIP_MODEL_CONFIG.keys() - ), f"model name should be in {SigLIP_MODEL_CONFIG.keys()}" + assert model_name in SigLIP_MODEL_CONFIG.keys(), ( + f"model name should be in {SigLIP_MODEL_CONFIG.keys()}" + ) vision_cfg = SigLIPVisionCfg(**SigLIP_MODEL_CONFIG[model_name]) @@ -1918,7 +1917,6 @@ class MultiModalityPreTrainedModel(PreTrainedModel): # Copied and adapted from: # https://github.com/deepseek-ai/Janus/tree/main/janus/models/modeling_vlm.py class MultiModalityCausalLM(MultiModalityPreTrainedModel): - def __init__( self, config: MultiModalityConfig, diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index bb03b82a6..e5b9a1707 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -99,7 +99,6 @@ _is_npu = is_npu() class DeepseekModelNextN(nn.Module): - def __init__( self, config: PretrainedConfig, diff --git a/python/sglang/srt/models/deepseek_ocr.py b/python/sglang/srt/models/deepseek_ocr.py index e86dae093..0f7d640a2 100644 --- a/python/sglang/srt/models/deepseek_ocr.py +++ b/python/sglang/srt/models/deepseek_ocr.py @@ -183,7 +183,6 @@ def merge_multimodal_embeddings( class MlpProjector(nn.Module): - def __init__( self, projector_type, @@ -458,9 +457,9 @@ class Attention(nn.Module): self.use_rel_pos = use_rel_pos if self.use_rel_pos: - assert ( - input_size is not None - ), "Input size must be provided if using relative positional encoding." + assert input_size is not None, ( + "Input size must be provided if using relative positional encoding." + ) # initialize relative positional embeddings self.rel_pos_h = nn.Parameter(torch.zeros(2 * input_size[0] - 1, head_dim)) self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, head_dim)) @@ -958,7 +957,6 @@ class NoTPAttention(torch.nn.Module): xqkv = xqkv.view(bsz, seqlen, 3, self.num_heads, self.head_dim) if self.use_flash_attention: - xq, xk, xv = torch.split(xqkv, 1, dim=2) xq = xq.squeeze(2) xk = xk.squeeze(2) @@ -1602,7 +1600,7 @@ class DeepseekOCRForCausalLM(nn.Module): if pixel_values is not None: if not isinstance(pixel_values, (torch.Tensor, list)): raise ValueError( - "Incorrect type of pixel values. " f"Got type: {type(pixel_values)}" + f"Incorrect type of pixel values. Got type: {type(pixel_values)}" ) if not isinstance(images_spatial_crop, (torch.Tensor, list)): @@ -1613,7 +1611,7 @@ class DeepseekOCRForCausalLM(nn.Module): if not isinstance(images_crop, (torch.Tensor, list)): raise ValueError( - "Incorrect type of image crop. " f"Got type: {type(images_crop)}" + f"Incorrect type of image crop. Got type: {type(images_crop)}" ) return [pixel_values, images_crop, images_spatial_crop] diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 65c880d48..3470b65c8 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -303,8 +303,7 @@ class DeepseekV2MLP(nn.Module): self.down_proj.weight = self.down_proj.weight_packed if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() self.use_fused_clamp_act_mul = _is_hip @@ -563,7 +562,6 @@ class MoEGate(nn.Module): class DeepseekV2MoE(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -785,13 +783,15 @@ class DeepseekV2MoE(nn.Module): self.shared_experts._enable_nvfp4_gemm_swiglu_fusion = True self.shared_experts.down_proj._accepts_prequantized_fp4 = True self._shared_expert_tp1 = _shared_expert_use_tp1 - is_packed_weight = hasattr( - self.shared_experts.gate_up_proj.quant_method, "quant_config" - ) and self.shared_experts.gate_up_proj.quant_method.quant_config.get_name() in { - "awq", - "awq_marlin", - "moe_wna16", - } + is_packed_weight = ( + hasattr(self.shared_experts.gate_up_proj.quant_method, "quant_config") + and self.shared_experts.gate_up_proj.quant_method.quant_config.get_name() + in { + "awq", + "awq_marlin", + "moe_wna16", + } + ) shared_gate_up_weight = getattr( self.shared_experts.gate_up_proj, "weight", None ) @@ -823,9 +823,7 @@ class DeepseekV2MoE(nn.Module): self.shared_experts.gate_up_proj.quant_method.quant_config.weight_block_size == self.shared_experts.down_proj.quant_method.quant_config.weight_block_size ) - self.shared_experts_weight_block_size = ( - self.shared_experts.gate_up_proj.quant_method.quant_config.weight_block_size - ) + self.shared_experts_weight_block_size = self.shared_experts.gate_up_proj.quant_method.quant_config.weight_block_size self.top_k = config.num_experts_per_tok @@ -1740,7 +1738,6 @@ class DeepseekV2AttentionMLA( DeepseekMLAFusedRopeRocmForwardMixin, DeepseekMLACpuForwardMixin, ): - def __init__( self, config: PretrainedConfig, @@ -2106,18 +2103,18 @@ class DeepseekV2AttentionMLA( not get_attn_tp_context().input_scattered and hidden_states[0].shape[0] == 0 ): - assert ( - not self.o_proj.reduce_results - ), "short-circuiting allreduce will lead to hangs" + assert not self.o_proj.reduce_results, ( + "short-circuiting allreduce will lead to hangs" + ) return hidden_states[0] else: if ( not get_attn_tp_context().input_scattered and hidden_states.shape[0] == 0 ): - assert ( - not self.o_proj.reduce_results - ), "short-circuiting allreduce will lead to hangs" + assert not self.o_proj.reduce_results, ( + "short-circuiting allreduce will lead to hangs" + ) return hidden_states, None, forward_batch, None attn_forward_method = self.dispatch_attn_forward_method(forward_batch) @@ -2301,7 +2298,6 @@ class DeepseekV2AttentionMLA( class DeepseekV2DecoderLayer(nn.Module): - def __init__( self, config: PretrainedConfig, diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 6733e56c8..68b4d1b08 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -596,9 +596,9 @@ def deepseek_v4_attention_with_output( finally: forward_batch.out_cache_loc = original_out_cache_loc - assert ( - output[:real_num_tokens].numel() == ret.numel() - ), f"Output tensor element mismatch: {output[:real_num_tokens].numel()} != {ret.numel()}" + assert output[:real_num_tokens].numel() == ret.numel(), ( + f"Output tensor element mismatch: {output[:real_num_tokens].numel()} != {ret.numel()}" + ) output[:real_num_tokens].view(ret.shape).copy_(ret) return @@ -610,7 +610,6 @@ bcg_deepseek_v4_attention_with_output = eager_on_graph(True)( class MqaAttentionBase(nn.Module): - def __init__( self, config: DeepSeekV4Config, @@ -740,9 +739,9 @@ class MqaAttentionBase(nn.Module): if fp8: from sglang.srt.layers import deep_gemm_wrapper - assert hasattr( - self.wo_a, "weight_scale_inv" - ), "FP8 quant_config must create weight_scale_inv" + assert hasattr(self.wo_a, "weight_scale_inv"), ( + "FP8 quant_config must create weight_scale_inv" + ) self.wo_a.weight_scale_inv.format_ue8m0 = ( deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 ) @@ -3910,9 +3909,9 @@ class DeepseekV4ForCausalLM(nn.Module): ) bucket = cache_wqkv_a_weight.setdefault(param_name, {}) shard_key = "q" if is_q else "kv" - assert ( - shard_key not in bucket - ), f"duplicate shard {shard_key} for {param_name}" + assert shard_key not in bucket, ( + f"duplicate shard {shard_key} for {param_name}" + ) bucket[shard_key] = _clone_if_runai_streamed_tensor( loaded_weight ) @@ -4027,9 +4026,9 @@ EntryClass = [DeepseekV4ForCausalLM] def _dequant_fp8(weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: from einops import rearrange - assert ( - weight.dtype == torch.float8_e4m3fn - ), f"expected fp8_e4m3fn, got {weight.dtype}" + assert weight.dtype == torch.float8_e4m3fn, ( + f"expected fp8_e4m3fn, got {weight.dtype}" + ) assert scale.dtype in ( torch.float8_e8m0fnu, torch.float32, diff --git a/python/sglang/srt/models/deepseek_v4_dspark.py b/python/sglang/srt/models/deepseek_v4_dspark.py index a90806318..8862bb7bb 100644 --- a/python/sglang/srt/models/deepseek_v4_dspark.py +++ b/python/sglang/srt/models/deepseek_v4_dspark.py @@ -93,7 +93,6 @@ def apply_rotary_emb( class DSparkAttention(MqaAttentionBase): - def __init__( self, config: DeepSeekV4Config, @@ -116,9 +115,9 @@ class DSparkAttention(MqaAttentionBase): wo_b_reduce_results=True, rope_original_seq_len=0, ) - assert ( - self.compress_ratio == 0 - ), "DSpark draft attention requires compress_ratio == 0." + assert self.compress_ratio == 0, ( + "DSpark draft attention requires compress_ratio == 0." + ) self.window_size = int( getattr(config, "sliding_window", None) or config.window_size ) @@ -354,7 +353,6 @@ def _resolve_dspark_pool() -> DeepSeekV4TokenToKVPool: class MarkovW2ShardGeometry(msgspec.Struct, frozen=True): - tp_size: int org_vocab_start: int org_vocab_end: int @@ -363,7 +361,6 @@ class MarkovW2ShardGeometry(msgspec.Struct, frozen=True): class DSparkV4MarkovHead(nn.Module): - markov_head_type = "vanilla" def __init__(self, *, vocab_size: int, markov_rank: int) -> None: @@ -531,7 +528,6 @@ def build_dspark_v4_confidence_head( class DSparkV4Stage(DeepseekV4DecoderLayer): - def __init__( self, config: DeepSeekV4Config, @@ -1018,12 +1014,12 @@ class DeepseekV4ForCausalLMDSpark(nn.Module): stage_id, rest = parts[1], parts[2] if rest.startswith("markov_head."): - return f"markov_head.{rest[len('markov_head.'):]}" + return f"markov_head.{rest[len('markov_head.') :]}" if rest.startswith("confidence_head."): if self.confidence_head is None: return None - return f"confidence_head.{rest[len('confidence_head.'):]}" + return f"confidence_head.{rest[len('confidence_head.') :]}" mapped_rest = rest mapped_rest = mapped_rest.replace("attn.", "self_attn.", 1) @@ -1069,12 +1065,12 @@ class DeepseekV4ForCausalLMDSpark(nn.Module): return None if rest.startswith("markov_head."): - return f"markov_head.{rest[len('markov_head.'):]}" + return f"markov_head.{rest[len('markov_head.') :]}" if rest.startswith("confidence_head."): if self.confidence_head is None: return None - return f"confidence_head.{rest[len('confidence_head.'):]}" + return f"confidence_head.{rest[len('confidence_head.') :]}" mapped_rest = rest if mapped_rest.startswith("attn."): diff --git a/python/sglang/srt/models/deepseek_v4_nextn.py b/python/sglang/srt/models/deepseek_v4_nextn.py index 39c064d9c..df937a1ba 100644 --- a/python/sglang/srt/models/deepseek_v4_nextn.py +++ b/python/sglang/srt/models/deepseek_v4_nextn.py @@ -221,7 +221,6 @@ class DeepseekV4ModelNextN(nn.Module): class DeepseekV4ForCausalLMNextN(DeepseekV4ForCausalLM): - def __init__( self, config: PretrainedConfig, diff --git a/python/sglang/srt/models/deepseek_vl2.py b/python/sglang/srt/models/deepseek_vl2.py index 98492deac..b1d11e523 100644 --- a/python/sglang/srt/models/deepseek_vl2.py +++ b/python/sglang/srt/models/deepseek_vl2.py @@ -156,7 +156,6 @@ class DeepseekVL2MlpProjector(nn.Module): class DeepseekVL2ForCausalLM(nn.Module): - @staticmethod def shared_experts_fusion_disable_reason(hf_config, quant_config): language_config = hf_config.language_config diff --git a/python/sglang/srt/models/dflash.py b/python/sglang/srt/models/dflash.py index 3a88ad392..59f907bb4 100644 --- a/python/sglang/srt/models/dflash.py +++ b/python/sglang/srt/models/dflash.py @@ -134,8 +134,7 @@ def _get_dflash_layer_attention_params( config, default=AttentionType.DECODER ) raise ValueError( - "Unsupported DFLASH draft layer type. " - f"layer_types[{layer_id}]={layer_type!r}." + f"Unsupported DFLASH draft layer type. layer_types[{layer_id}]={layer_type!r}." ) diff --git a/python/sglang/srt/models/dots3_common/dots_omni_audio.py b/python/sglang/srt/models/dots3_common/dots_omni_audio.py index c2c2d2c47..76ba93184 100644 --- a/python/sglang/srt/models/dots3_common/dots_omni_audio.py +++ b/python/sglang/srt/models/dots3_common/dots_omni_audio.py @@ -918,8 +918,7 @@ class DotsEncoderWithMask(nn.Module): while time_step * SAMPLE_RATE < audio_waveform.shape[0]: segments.append( audio_waveform[ - time_step - * SAMPLE_RATE : (time_step + self.chunk_seconds) + time_step * SAMPLE_RATE : (time_step + self.chunk_seconds) * SAMPLE_RATE ] ) diff --git a/python/sglang/srt/models/dots3_common/modeling.py b/python/sglang/srt/models/dots3_common/modeling.py index 360bd114f..ad9ce1de6 100644 --- a/python/sglang/srt/models/dots3_common/modeling.py +++ b/python/sglang/srt/models/dots3_common/modeling.py @@ -735,14 +735,16 @@ class Dots3AttentionMLA(nn.Module): self.num_heads = attn_config.num_attention_heads assert self.num_heads % attn_tp_size == 0 self.num_local_heads = self.num_heads // attn_tp_size - assert ( - attn_config.num_attention_heads == attn_config.num_key_value_heads - ), "Dots3 Only supports equal number of query and key value heads." + assert attn_config.num_attention_heads == attn_config.num_key_value_heads, ( + "Dots3 Only supports equal number of query and key value heads." + ) self.attention_gate_type = attn_config.attention_gate_type assert self.attention_gate_type in { "headwise", "elementwise", - }, f"Unsupported attention_gate_type: {self.attention_gate_type}. Expected 'headwise' or 'elementwise'." + }, ( + f"Unsupported attention_gate_type: {self.attention_gate_type}. Expected 'headwise' or 'elementwise'." + ) self.g_proj_local_dim = self.num_local_heads * ( 1 if self.attention_gate_type == "headwise" else self.v_head_dim ) @@ -832,9 +834,9 @@ class Dots3AttentionMLA(nn.Module): # Optional NSA (Native Sparse Attention) indexer. if self.use_nsa: - assert ( - self.q_lora_rank is not None - ), "Dots3 NSA requires q_lora_rank to be set in the config." + assert self.q_lora_rank is not None, ( + "Dots3 NSA requires q_lora_rank to be set in the config." + ) self.indexer = Indexer( hidden_size=self.hidden_size, index_n_heads=config.index_n_heads, @@ -1073,9 +1075,9 @@ class Dots3AttentionMLA(nn.Module): zero_allocator: BumpAllocator, ): if hidden_states.shape[0] == 0: - assert ( - not self.o_proj.reduce_results - ), "short-circuiting allreduce will lead to hangs" + assert not self.o_proj.reduce_results, ( + "short-circuiting allreduce will lead to hangs" + ) return hidden_states, None, forward_batch, None attn_forward_method = self.dispatch_attn_forward_method(forward_batch) @@ -1853,9 +1855,9 @@ class Dots3LanguageModelForCausalLM(nn.Module): # for quark model load # Always fuse q_a_proj/kv_a_proj_with_mqa/g_proj when loading Dots3. self.fuse_qkv_a_g_proj = True - assert ( - config.q_lora_rank is not None - ), "Dots3 requires q_lora_rank to enable fused_qkv_a_g_proj_with_mqa loading." + assert config.q_lora_rank is not None, ( + "Dots3 requires q_lora_rank to enable fused_qkv_a_g_proj_with_mqa loading." + ) if self.fuse_qkv_a_g_proj: self.packed_modules_mapping["fused_qkv_a_g_proj_with_mqa"] = [ "q_a_proj", @@ -2039,12 +2041,16 @@ class Dots3LanguageModelForCausalLM(nn.Module): assert ( self.quant_config is not None and self.quant_config.weight_block_size is not None - ), "Dots3 MLA kv_b_proj only supports FP8 block quantization with weight_block_size=(128, 128)." + ), ( + "Dots3 MLA kv_b_proj only supports FP8 block quantization with weight_block_size=(128, 128)." + ) weight_block_size = tuple(self.quant_config.weight_block_size) assert weight_block_size == ( 128, 128, - ), f"Dots3 MLA kv_b_proj only supports FP8 block_size=(128, 128), got {weight_block_size}." + ), ( + f"Dots3 MLA kv_b_proj only supports FP8 block_size=(128, 128), got {weight_block_size}." + ) block_scale = self_attn.kv_b_proj.weight_scale_inv if not ( @@ -2062,9 +2068,9 @@ class Dots3LanguageModelForCausalLM(nn.Module): torch.bfloat16, ) else: - assert ( - w.dtype == torch.bfloat16 - ), f"Dots3 MLA kv_b_proj only supports BF16 or FP8(128x128), got dtype={w.dtype}." + assert w.dtype == torch.bfloat16, ( + f"Dots3 MLA kv_b_proj only supports BF16 or FP8(128x128), got dtype={w.dtype}." + ) w_kc, w_vc = w.unflatten( 0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim) @@ -2207,9 +2213,9 @@ class Dots3LanguageModelForCausalLM(nn.Module): # Always fuse q_a_proj/kv_a_proj_with_mqa/g_proj when loading Dots3. fuse_qkv_a_g_proj = True - assert ( - self.config.q_lora_rank is not None - ), "Dots3 requires q_lora_rank to enable fused_qkv_a_g_proj_with_mqa loading." + assert self.config.q_lora_rank is not None, ( + "Dots3 requires q_lora_rank to enable fused_qkv_a_g_proj_with_mqa loading." + ) cached_a_proj = {} if fuse_qkv_a_g_proj else None attn_tp_rank = get_parallel().attn_tp_rank attn_tp_size = get_parallel().attn_tp_size @@ -2217,9 +2223,9 @@ class Dots3LanguageModelForCausalLM(nn.Module): def shard_g_proj_for_attention_tp( weight: torch.Tensor, cat_dim: int, is_scale: bool ): - assert ( - weight.ndim > cat_dim - ), f"weight.ndim={weight.ndim}, cat_dim={cat_dim}" + assert weight.ndim > cat_dim, ( + f"weight.ndim={weight.ndim}, cat_dim={cat_dim}" + ) dim_size = weight.shape[cat_dim] if not is_scale: assert dim_size % attn_tp_size == 0, ( diff --git a/python/sglang/srt/models/dspark.py b/python/sglang/srt/models/dspark.py index be3930638..21861e3a7 100644 --- a/python/sglang/srt/models/dspark.py +++ b/python/sglang/srt/models/dspark.py @@ -84,7 +84,6 @@ def run_markov_block( class VanillaMarkov(nn.Module): - markov_head_type = "vanilla" def __init__(self, *, vocab_size: int, markov_rank: int) -> None: @@ -219,7 +218,6 @@ class Nemotron35VanillaMarkov(VanillaMarkov): class GatedMarkovHead(VanillaMarkov): - markov_head_type = "gated" def __init__(self, *, vocab_size: int, markov_rank: int, hidden_size: int) -> None: @@ -260,7 +258,6 @@ class GatedMarkovHead(VanillaMarkov): class RNNHead(VanillaMarkov): - markov_head_type = "rnn" def __init__(self, *, vocab_size: int, markov_rank: int, hidden_size: int) -> None: @@ -413,7 +410,6 @@ def build_nemotron_35_markov_head(config, quant_config, prefix: str) -> nn.Modul class DSparkConfidenceHead(nn.Module): - def __init__( self, *, @@ -483,7 +479,6 @@ _DSPARK_SKIPPED_WEIGHT_PREFIXES = ("lm_head.", "rotary_emb.") class DSparkDraftMixin: - supports_pre_gather_target_hidden_projection = True def __init__(self, config, quant_config=None, prefix: str = "") -> None: @@ -859,7 +854,6 @@ class DSparkDraftMixin: class DSparkDraftModel(DSparkDraftMixin, DFlashDraftModel): - def prune_to_ctx_kv_injection(self) -> None: self.markov_head = None self.confidence_head = None diff --git a/python/sglang/srt/models/ernie45_moe_vl.py b/python/sglang/srt/models/ernie45_moe_vl.py index 7ad2c26a3..3cd69b1ee 100644 --- a/python/sglang/srt/models/ernie45_moe_vl.py +++ b/python/sglang/srt/models/ernie45_moe_vl.py @@ -227,7 +227,6 @@ class Ernie4_5_VLMoeMoE(nn.Module): layer_id >= vision_moe_layer_start_index and layer_id <= vision_moe_layer_end_index ): - self.vision_experts_gate = ReplicatedLinear( config.hidden_size, config.moe_num_experts[1], diff --git a/python/sglang/srt/models/ernie45_vl.py b/python/sglang/srt/models/ernie45_vl.py index c798fed18..343837fde 100644 --- a/python/sglang/srt/models/ernie45_vl.py +++ b/python/sglang/srt/models/ernie45_vl.py @@ -54,7 +54,6 @@ logger = logging.getLogger(__name__) class Ernie4_5_VisionMLP(nn.Module): - def __init__( self, in_features: int, @@ -86,7 +85,6 @@ class Ernie4_5_VisionMLP(nn.Module): class Ernie4_5_VisionBlock(nn.Module): - def __init__( self, dim: int, @@ -145,7 +143,6 @@ class Ernie4_5_VisionBlock(nn.Module): class Ernie4_5_VisionPatchEmbed(nn.Module): - def __init__( self, patch_size: int = 14, @@ -352,7 +349,6 @@ class VariableResolutionResamplerModel(nn.Module): class Ernie4_5_VisionRotaryEmbedding(nn.Module): - def __init__(self, dim: int, theta: float = 10000.0) -> None: super().__init__() self.inv_freq = 1.0 / theta ** ( @@ -368,7 +364,6 @@ class Ernie4_5_VisionRotaryEmbedding(nn.Module): class Ernie4_5_VisionTransformer(nn.Module): - def __init__( self, vision_config: PretrainedConfig, @@ -715,9 +710,9 @@ class Ernie4_5_VLMoeForConditionalGeneration(nn.Module): self._set_visual_token_mask(input_ids, forward_batch) - assert ( - input_ids.numel() == positions.shape[-1] - ), f"input_ids {input_ids.shape} and position_ids {positions.shape} should have the same length" + assert input_ids.numel() == positions.shape[-1], ( + f"input_ids {input_ids.shape} and position_ids {positions.shape} should have the same length" + ) hidden_states = general_mm_embed_routine( input_ids=input_ids, diff --git a/python/sglang/srt/models/exaone.py b/python/sglang/srt/models/exaone.py index 441cfafde..d6be6a096 100644 --- a/python/sglang/srt/models/exaone.py +++ b/python/sglang/srt/models/exaone.py @@ -69,8 +69,7 @@ class ExaoneGatedMLP(nn.Module): ) if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() diff --git a/python/sglang/srt/models/exaone4.py b/python/sglang/srt/models/exaone4.py index d01dad06c..63ee1b342 100644 --- a/python/sglang/srt/models/exaone4.py +++ b/python/sglang/srt/models/exaone4.py @@ -69,8 +69,7 @@ class Exaone4GatedMLP(nn.Module): ) if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() diff --git a/python/sglang/srt/models/exaone_moe.py b/python/sglang/srt/models/exaone_moe.py index 93e8de9fb..61c6cd255 100755 --- a/python/sglang/srt/models/exaone_moe.py +++ b/python/sglang/srt/models/exaone_moe.py @@ -112,8 +112,7 @@ class ExaoneMoEMLP(nn.Module): ) if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() diff --git a/python/sglang/srt/models/falcon_h1.py b/python/sglang/srt/models/falcon_h1.py index 8e3060cfb..e49f964ed 100644 --- a/python/sglang/srt/models/falcon_h1.py +++ b/python/sglang/srt/models/falcon_h1.py @@ -70,8 +70,7 @@ class FalconH1MLP(nn.Module): ) if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() self.layer_id = layer_id @@ -96,7 +95,6 @@ class FalconH1MLP(nn.Module): class FalconH1HybridAttentionDecoderLayer(nn.Module): - def __init__( self, config: FalconH1Config, @@ -275,16 +273,18 @@ class FalconH1HybridAttentionDecoderLayer(nn.Module): # B vector 2 * d_ssm -> 2 * d_ssm + (n_group * d_state) mup_vector[ :, - (2 * self.d_ssm) - // self.tp_size : (2 * self.d_ssm + self.groups_time_state_size) + (2 * self.d_ssm) // self.tp_size : ( + 2 * self.d_ssm + self.groups_time_state_size + ) // self.tp_size, ] *= self.zxbcdt_multipliers[2] # C vector 2 * d_ssm + (n_group * d_state) # -> 2 * d_ssm + 2 * (n_group * d_state) mup_vector[ :, - (2 * self.d_ssm + self.groups_time_state_size) - // self.tp_size : (2 * self.d_ssm + 2 * self.groups_time_state_size) + (2 * self.d_ssm + self.groups_time_state_size) // self.tp_size : ( + 2 * self.d_ssm + 2 * self.groups_time_state_size + ) // self.tp_size, ] *= self.zxbcdt_multipliers[3] # dt vector 2 * d_ssm + 2 * (n_group * d_state) @@ -521,7 +521,6 @@ class FalconH1ForCausalLM(nn.Module): params_dict = dict(self.named_parameters()) loaded_params: Set[str] = set() for name, loaded_weight in weights: - if "rotary_emb.inv_freq" in name: continue diff --git a/python/sglang/srt/models/gemma2_reward.py b/python/sglang/srt/models/gemma2_reward.py index 8c8eda22b..d0485c9cb 100644 --- a/python/sglang/srt/models/gemma2_reward.py +++ b/python/sglang/srt/models/gemma2_reward.py @@ -53,9 +53,9 @@ class Gemma2ForSequenceClassification(nn.Module): input_embeds: torch.Tensor = None, get_embedding: bool = True, ) -> EmbeddingPoolerOutput: - assert ( - get_embedding - ), "Gemma2ForSequenceClassification is only used for embedding" + assert get_embedding, ( + "Gemma2ForSequenceClassification is only used for embedding" + ) hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) last_token_hidden = self.pooler(hidden_states, forward_batch).embeddings diff --git a/python/sglang/srt/models/gemma4_causal.py b/python/sglang/srt/models/gemma4_causal.py index 4ebe925c5..26917a2de 100644 --- a/python/sglang/srt/models/gemma4_causal.py +++ b/python/sglang/srt/models/gemma4_causal.py @@ -981,9 +981,9 @@ class Gemma4TextModel(PreTrainedModel): ) hidden_states = input_embeds else: - assert ( - pp_proxy_tensors is not None - ), "pp_proxy_tensors is required on non-first PP ranks" + assert pp_proxy_tensors is not None, ( + "pp_proxy_tensors is required on non-first PP ranks" + ) hidden_states = pp_proxy_tensors["hidden_states"] # PLE inputs were computed on rank 0 and forwarded along the # pipeline; non-PLE models simply omit the key. diff --git a/python/sglang/srt/models/gemma4_vision.py b/python/sglang/srt/models/gemma4_vision.py index d71b0226c..9e7548654 100644 --- a/python/sglang/srt/models/gemma4_vision.py +++ b/python/sglang/srt/models/gemma4_vision.py @@ -204,7 +204,6 @@ class Gemma4VisionAttention(nn.Module): if is_cuda(): major, _ = get_device_capability() if major == 9: - if get_platform().is_blackwell: return "triton_attn" return "fa3" diff --git a/python/sglang/srt/models/glm4_moe.py b/python/sglang/srt/models/glm4_moe.py index 7f7a6392a..8f2791be4 100644 --- a/python/sglang/srt/models/glm4_moe.py +++ b/python/sglang/srt/models/glm4_moe.py @@ -472,13 +472,15 @@ class Glm4MoeSparseMoeBlock(nn.Module): else {} ), ) - is_packed_weight = hasattr( - self.shared_experts.gate_up_proj.quant_method, "quant_config" - ) and self.shared_experts.gate_up_proj.quant_method.quant_config.get_name() in { - "awq", - "awq_marlin", - "moe_wna16", - } + is_packed_weight = ( + hasattr(self.shared_experts.gate_up_proj.quant_method, "quant_config") + and self.shared_experts.gate_up_proj.quant_method.quant_config.get_name() + in { + "awq", + "awq_marlin", + "moe_wna16", + } + ) self.shared_experts_is_int8 = ( not is_packed_weight and self.shared_experts.gate_up_proj.weight.dtype == torch.int8 @@ -500,9 +502,7 @@ class Glm4MoeSparseMoeBlock(nn.Module): self.shared_experts.gate_up_proj.quant_method.quant_config.weight_block_size == self.shared_experts.down_proj.quant_method.quant_config.weight_block_size ) - self.shared_experts_weight_block_size = ( - self.shared_experts.gate_up_proj.quant_method.quant_config.weight_block_size - ) + self.shared_experts_weight_block_size = self.shared_experts.gate_up_proj.quant_method.quant_config.weight_block_size self.top_k = config.num_experts_per_tok diff --git a/python/sglang/srt/models/glm4v_moe.py b/python/sglang/srt/models/glm4v_moe.py index 9c976fa80..5f9d2ed51 100644 --- a/python/sglang/srt/models/glm4v_moe.py +++ b/python/sglang/srt/models/glm4v_moe.py @@ -104,9 +104,9 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): if is_shared_experts_fusion_disabled(): return self.num_fused_shared_experts = self.config.n_shared_experts - assert ( - self.num_fused_shared_experts == 1 - ), "Only 1 fused shared expert is supported for Glm4vMoeForConditionalGeneration" + assert self.num_fused_shared_experts == 1, ( + "Only 1 fused shared expert is supported for Glm4vMoeForConditionalGeneration" + ) log_info_on_rank0(logger, "Shared experts fusion optimization enabled.") def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]], is_nextn=False): diff --git a/python/sglang/srt/models/gpt2.py b/python/sglang/srt/models/gpt2.py index 613690f2f..21e5ad59b 100644 --- a/python/sglang/srt/models/gpt2.py +++ b/python/sglang/srt/models/gpt2.py @@ -41,7 +41,6 @@ from sglang.srt.utils import add_prefix class GPT2Attention(nn.Module): - def __init__( self, layer_id: int, @@ -95,7 +94,6 @@ class GPT2Attention(nn.Module): class GPT2MLP(nn.Module): - def __init__( self, intermediate_size: int, @@ -133,7 +131,6 @@ class GPT2MLP(nn.Module): class GPT2Block(nn.Module): - def __init__( self, layer_id: int, @@ -182,7 +179,6 @@ class GPT2Block(nn.Module): class GPT2Model(nn.Module): - def __init__( self, config: GPT2Config, @@ -230,7 +226,6 @@ class GPT2Model(nn.Module): class GPT2LMHeadModel(nn.Module): - def __init__( self, config: GPT2Config, diff --git a/python/sglang/srt/models/gpt_bigcode.py b/python/sglang/srt/models/gpt_bigcode.py index 47ac77609..bc63c8656 100644 --- a/python/sglang/srt/models/gpt_bigcode.py +++ b/python/sglang/srt/models/gpt_bigcode.py @@ -42,7 +42,6 @@ from sglang.srt.utils import add_prefix class GPTBigCodeAttention(nn.Module): - def __init__( self, layer_id: int, @@ -114,7 +113,6 @@ class GPTBigCodeAttention(nn.Module): class GPTBigMLP(nn.Module): - def __init__( self, intermediate_size: int, @@ -150,7 +148,6 @@ class GPTBigMLP(nn.Module): class GPTBigCodeBlock(nn.Module): - def __init__( self, layer_id: int, @@ -193,7 +190,6 @@ class GPTBigCodeBlock(nn.Module): class GPTBigCodeModel(nn.Module): - def __init__( self, config: GPTBigCodeConfig, diff --git a/python/sglang/srt/models/gpt_j.py b/python/sglang/srt/models/gpt_j.py index 72c736724..c90dec402 100644 --- a/python/sglang/srt/models/gpt_j.py +++ b/python/sglang/srt/models/gpt_j.py @@ -49,7 +49,6 @@ from sglang.srt.utils import add_prefix class GPTJAttention(nn.Module): - def __init__( self, layer_id: int, @@ -118,7 +117,6 @@ class GPTJAttention(nn.Module): class GPTJMLP(nn.Module): - def __init__( self, intermediate_size: int, @@ -151,7 +149,6 @@ class GPTJMLP(nn.Module): class GPTJBlock(nn.Module): - def __init__( self, layer_id: int, @@ -194,7 +191,6 @@ class GPTJBlock(nn.Module): class GPTJModel(nn.Module): - def __init__( self, config: GPTJConfig, @@ -242,7 +238,6 @@ class GPTJModel(nn.Module): class GPTJForCausalLM(nn.Module): - def __init__( self, config: GPTJConfig, diff --git a/python/sglang/srt/models/gpt_oss.py b/python/sglang/srt/models/gpt_oss.py index 3c6f2b736..0db01f79e 100644 --- a/python/sglang/srt/models/gpt_oss.py +++ b/python/sglang/srt/models/gpt_oss.py @@ -224,8 +224,7 @@ class GptOssSparseMoeBlock(nn.Module): ) extra_kwargs = { # for moe gate_up_proj and down_proj and their bias loading - "use_weight_loader_fused": quant_config_name - != "mxfp4" + "use_weight_loader_fused": quant_config_name != "mxfp4" } self.experts = experts_type( @@ -977,9 +976,9 @@ class GptOssForCausalLM(nn.Module): original_intermediate_size = getattr( self.config, "original_intermediate_size", intermediate_size ) - assert ( - intermediate_size % mxfp4_block == 0 - ), f"{intermediate_size=} must be divisible by {mxfp4_block=}" + assert intermediate_size % mxfp4_block == 0, ( + f"{intermediate_size=} must be divisible by {mxfp4_block=}" + ) intermediate_size_block = intermediate_size // mxfp4_block per_rank_intermediate_size_block = math.ceil( diff --git a/python/sglang/srt/models/granite.py b/python/sglang/srt/models/granite.py index 45a3f239e..c5102822d 100644 --- a/python/sglang/srt/models/granite.py +++ b/python/sglang/srt/models/granite.py @@ -109,8 +109,7 @@ class GraniteMLP(nn.Module): ) if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() diff --git a/python/sglang/srt/models/granitemoe.py b/python/sglang/srt/models/granitemoe.py index 17989711b..b38814ae9 100644 --- a/python/sglang/srt/models/granitemoe.py +++ b/python/sglang/srt/models/granitemoe.py @@ -115,7 +115,6 @@ class GraniteMoeMoE(nn.Module): class GraniteMoeSharedMLP(nn.Module): - def __init__( self, config: GraniteConfig, @@ -155,7 +154,6 @@ class GraniteMoeSharedMLP(nn.Module): class GraniteMoeAttention(nn.Module): - def __init__( self, config: GraniteConfig, @@ -254,7 +252,6 @@ class GraniteMoeAttention(nn.Module): class GraniteMoeDecoderLayer(nn.Module): - def __init__( self, config: GraniteConfig, @@ -333,7 +330,6 @@ class GraniteMoeDecoderLayer(nn.Module): class GraniteMoeModel(nn.Module): - def __init__( self, config: GraniteConfig, @@ -389,7 +385,6 @@ class GraniteMoeModel(nn.Module): class GraniteMoeForCausalLM(nn.Module): - # Legacy and current HF expert / router names with otherwise shared layout hf_to_sglang_mapper = WeightsMapper( orig_to_new_suffix={ diff --git a/python/sglang/srt/models/granitemoehybrid.py b/python/sglang/srt/models/granitemoehybrid.py index ed6f37df5..b6345e7bf 100644 --- a/python/sglang/srt/models/granitemoehybrid.py +++ b/python/sglang/srt/models/granitemoehybrid.py @@ -179,7 +179,6 @@ class GraniteMoeHybridAttention(nn.Module): ) if config.position_embedding_type == "rope": - self.rotary_emb = get_rope( head_size=self.head_dim, rotary_dim=self.head_dim, # its not in the config diff --git a/python/sglang/srt/models/hunyuan.py b/python/sglang/srt/models/hunyuan.py index ef082e429..5649a42c2 100644 --- a/python/sglang/srt/models/hunyuan.py +++ b/python/sglang/srt/models/hunyuan.py @@ -74,7 +74,6 @@ def _get_cla_factor(config: PretrainedConfig) -> int: class HunYuanMLP(nn.Module): - def __init__( self, hidden_size: int, @@ -103,8 +102,7 @@ class HunYuanMLP(nn.Module): ) if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() @@ -116,7 +114,6 @@ class HunYuanMLP(nn.Module): class HunYuanSparseMoeBlock(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -244,7 +241,6 @@ def check_head_dim(config): class HunYuanAttention(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -385,7 +381,6 @@ class HunYuanAttention(nn.Module): class HunYuanDecoderLayer(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -483,7 +478,6 @@ class HunYuanDecoderLayer(nn.Module): class HunYuanModel(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -797,7 +791,7 @@ class HunYuanMoEV1ForCausalLM(nn.Module): layer_self_attn.attn._kv_scale = scaling_factor else: raise RuntimeError( - "Self attention has no KV cache scaling " "factor attribute!" + "Self attention has no KV cache scaling factor attribute!" ) diff --git a/python/sglang/srt/models/hunyuan_v3_nextn.py b/python/sglang/srt/models/hunyuan_v3_nextn.py index 1b2388391..f831d54ef 100644 --- a/python/sglang/srt/models/hunyuan_v3_nextn.py +++ b/python/sglang/srt/models/hunyuan_v3_nextn.py @@ -39,7 +39,6 @@ logger = logging.getLogger(__name__) class HYV3ModelNextN(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -116,7 +115,6 @@ class HYV3ModelNextN(nn.Module): class HYV3ForCausalLMNextN(nn.Module): - def __init__( self, config: PretrainedConfig, diff --git a/python/sglang/srt/models/idefics2.py b/python/sglang/srt/models/idefics2.py index d54aac976..f833ce730 100644 --- a/python/sglang/srt/models/idefics2.py +++ b/python/sglang/srt/models/idefics2.py @@ -37,7 +37,6 @@ from sglang.srt.utils import add_prefix, is_npu class Idefics2VisionMLP(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -70,7 +69,6 @@ class Idefics2VisionMLP(nn.Module): class Idefics2EncoderLayer(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -241,7 +239,6 @@ class Idefics2VisionEmbeddings(nn.Module): ) for batch_idx, p_attn_mask in enumerate(patch_attention_mask): - if tgt_sizes is not None: nb_patches_h = tgt_sizes[batch_idx][0] nb_patches_w = tgt_sizes[batch_idx][1] @@ -284,7 +281,6 @@ class Idefics2VisionEmbeddings(nn.Module): class Idefics2VisionTransformer(nn.Module): - def __init__( self, config: PretrainedConfig, diff --git a/python/sglang/srt/models/inkling_common/dense_mlp.py b/python/sglang/srt/models/inkling_common/dense_mlp.py index 66d3edadc..35748c8e5 100644 --- a/python/sglang/srt/models/inkling_common/dense_mlp.py +++ b/python/sglang/srt/models/inkling_common/dense_mlp.py @@ -329,9 +329,9 @@ class InklingBatchDenseMLP(nn.Module, FusedMoELoadingMixin): """ assert x.ndim in (2, 3), f"{x.shape=}" assert gammas.ndim in (2, 3), f"{gammas.shape=}" - assert ( - gammas.size(-1) == self.n_shared_experts - ), f"{gammas.shape=} {self.n_shared_experts=}" + assert gammas.size(-1) == self.n_shared_experts, ( + f"{gammas.shape=} {self.n_shared_experts=}" + ) if self._fp4_strategy.serves_fp4: return self._forward_fp4(x, gammas, use_reduce_scatter) @@ -387,9 +387,9 @@ class InklingBatchDenseMLP(nn.Module, FusedMoELoadingMixin): silu_and_mul_triton, ) - assert ( - self.inference_moe_w13_interleaved - ), "silu_and_mul_triton requires interleaved w13" + assert self.inference_moe_w13_interleaved, ( + "silu_and_mul_triton requires interleaved w13" + ) y_st_2f = y_st2f.view(-1, y_st2f.size(-1)) y_st_f = silu_and_mul_triton(y_st_2f, gammas_st.reshape(-1)) return y_st_f.view(*y_st2f.shape[:-1], y_st2f.size(-1) // 2) diff --git a/python/sglang/srt/models/inkling_common/kernels/sconv.py b/python/sglang/srt/models/inkling_common/kernels/sconv.py index 37b7c0926..84b2f7658 100644 --- a/python/sglang/srt/models/inkling_common/kernels/sconv.py +++ b/python/sglang/srt/models/inkling_common/kernels/sconv.py @@ -60,9 +60,9 @@ def _metadata_out( ) for name, shape, dtype in spec: t = out[name] - assert ( - tuple(t.shape) == shape and t.dtype == dtype and t.is_contiguous() - ), f"{name}: got {tuple(t.shape)}/{t.dtype}, want {shape}/{dtype} contiguous" + assert tuple(t.shape) == shape and t.dtype == dtype and t.is_contiguous(), ( + f"{name}: got {tuple(t.shape)}/{t.dtype}, want {shape}/{dtype} contiguous" + ) return out @@ -1173,9 +1173,9 @@ def save_intermediate_conv_windows( if hidden_states.dim() == 2: hidden_states = hidden_states.view(batch_size, -1, hidden_states.shape[-1]) - assert ( - hidden_states.dim() == 3 - ), f"unexpected hidden_states shape {hidden_states.shape}" + assert hidden_states.dim() == 3, ( + f"unexpected hidden_states shape {hidden_states.shape}" + ) assert hidden_states.shape[0] == batch_size assert hidden_states.shape[2] == D assert intermediate_out.shape[1] == draft_token_num diff --git a/python/sglang/srt/models/inkling_common/moe.py b/python/sglang/srt/models/inkling_common/moe.py index 7f23b8cc5..f5b6c52e5 100644 --- a/python/sglang/srt/models/inkling_common/moe.py +++ b/python/sglang/srt/models/inkling_common/moe.py @@ -512,7 +512,9 @@ def make_forward_inputs_2d( return hidden_states, topk_weights, topk_ids, top_k, num_experts -def run_moe_preprocess(topk_ids: torch.Tensor, num_experts: int) -> tuple[ +def run_moe_preprocess( + topk_ids: torch.Tensor, num_experts: int +) -> tuple[ torch.Tensor, torch.Tensor, torch.Tensor, @@ -562,9 +564,9 @@ def activation( use_interleaved: bool = True, ): if activation_type == "silu_and_mul": - assert ( - gateup_output.is_contiguous() - ), f"{gateup_output.shape=} {gateup_output.stride()=}" + assert gateup_output.is_contiguous(), ( + f"{gateup_output.shape=} {gateup_output.stride()=}" + ) assert gateup_output.ndim == 2, f"{gateup_output.shape=}" out_dtype = None if gateup_output.numel() == 0: diff --git a/python/sglang/srt/models/inkling_common/quantization/quant.py b/python/sglang/srt/models/inkling_common/quantization/quant.py index dbb1d187d..9efb274b5 100644 --- a/python/sglang/srt/models/inkling_common/quantization/quant.py +++ b/python/sglang/srt/models/inkling_common/quantization/quant.py @@ -50,9 +50,9 @@ class InklingNvfp4MoEMethod(InklingMoEMethodBase): from sglang.srt.models.inkling_common.dense_mlp import InklingBatchDenseMLP - assert isinstance( - layer, InklingBatchDenseMLP - ), "InklingNvfp4MoEMethod is only used for InklingBatchDenseMLP (shared experts)" + assert isinstance(layer, InklingBatchDenseMLP), ( + "InklingNvfp4MoEMethod is only used for InklingBatchDenseMLP (shared experts)" + ) w13_up_dim = 2 * intermediate_size_per_partition diff --git a/python/sglang/srt/models/internlm2.py b/python/sglang/srt/models/internlm2.py index ee9807ac6..d376d96c0 100644 --- a/python/sglang/srt/models/internlm2.py +++ b/python/sglang/srt/models/internlm2.py @@ -70,8 +70,7 @@ class InternLM2MLP(nn.Module): ) if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() diff --git a/python/sglang/srt/models/interns1.py b/python/sglang/srt/models/interns1.py index d64b767c8..e6eebacfd 100644 --- a/python/sglang/srt/models/interns1.py +++ b/python/sglang/srt/models/interns1.py @@ -186,7 +186,6 @@ class InternS1ForConditionalGeneration(nn.Module): name = "vision_model." + name[len("model.vision_tower.") :] if name.startswith("vision_model.encoder.layer"): - name = name.replace(r".layer.", r".layers.") name = name.replace(r".attention.", r".attn.attn.") name = name.replace(r".projection_layer.", r".proj.") diff --git a/python/sglang/srt/models/interns1pro.py b/python/sglang/srt/models/interns1pro.py index df3b4d03b..0bc524452 100644 --- a/python/sglang/srt/models/interns1pro.py +++ b/python/sglang/srt/models/interns1pro.py @@ -115,9 +115,9 @@ class InternS1ProTextDecoderLayer(Qwen3MoeDecoderLayer): # update with group router self.router_n_groups = getattr(config, "router_n_groups", -1) if self.router_n_groups > 0: - assert ( - config.num_experts_per_tok % self.router_n_groups == 0 - ), f"{config.num_experts_per_tok} cannot be divided by {self.router_n_groups}" + assert config.num_experts_per_tok % self.router_n_groups == 0, ( + f"{config.num_experts_per_tok} cannot be divided by {self.router_n_groups}" + ) self.mlp.topk = TopK( top_k=config.num_experts_per_tok, renormalize=config.norm_topk_prob, @@ -131,9 +131,7 @@ class InternS1ProTextDecoderLayer(Qwen3MoeDecoderLayer): def get_group_offsets(router_n_groups: int, group_size: int, device: str): group_offsets = ( torch.arange(router_n_groups, device=device) * group_size - ).view( - 1, -1, 1 - ) # [1, n_groups, 1] + ).view(1, -1, 1) # [1, n_groups, 1] return group_offsets def _custom_routing_function( @@ -146,9 +144,9 @@ class InternS1ProTextDecoderLayer(Qwen3MoeDecoderLayer): """Group router""" routing_weights = torch.softmax(gating_output, dim=-1, dtype=torch.float32) if self.router_n_groups > 0: - assert ( - routing_weights.shape[-1] % self.router_n_groups == 0 - ), f"{routing_weights.shape[-1]} cannot be divided by {self.router_n_groups}" + assert routing_weights.shape[-1] % self.router_n_groups == 0, ( + f"{routing_weights.shape[-1]} cannot be divided by {self.router_n_groups}" + ) per_group_top_k = topk // self.router_n_groups group_size = routing_weights.shape[-1] // self.router_n_groups group_offsets = self.get_group_offsets( @@ -189,7 +187,6 @@ class InternS1ProTextModel(Qwen3MoeLLMModel): class InternS1ProForConditionalGeneration(Qwen3VLMoeForConditionalGeneration): - def __init__( self, config: PretrainedConfig, diff --git a/python/sglang/srt/models/internvl.py b/python/sglang/srt/models/internvl.py index d6880e2a1..4ab7ea3b0 100644 --- a/python/sglang/srt/models/internvl.py +++ b/python/sglang/srt/models/internvl.py @@ -234,7 +234,6 @@ NORM2FN = { class InternVisionEncoderLayer(nn.Module): - def __init__( self, config: PretrainedConfig, diff --git a/python/sglang/srt/models/iquest_loopcoder.py b/python/sglang/srt/models/iquest_loopcoder.py index 62488a432..eb98bc4b2 100644 --- a/python/sglang/srt/models/iquest_loopcoder.py +++ b/python/sglang/srt/models/iquest_loopcoder.py @@ -90,9 +90,9 @@ class LoopGateProjection(nn.Module): """ num_heads, num_tokens, head_dim = query.shape - assert ( - num_heads == self.num_heads - ), f"Expected {self.num_heads} heads, got {num_heads}" + assert num_heads == self.num_heads, ( + f"Expected {self.num_heads} heads, got {num_heads}" + ) query_flat = query.reshape(-1, head_dim) diff --git a/python/sglang/srt/models/kimi_k25.py b/python/sglang/srt/models/kimi_k25.py index debbf587b..54208e7d7 100644 --- a/python/sglang/srt/models/kimi_k25.py +++ b/python/sglang/srt/models/kimi_k25.py @@ -87,7 +87,6 @@ def apply_rope( class MoonViTEncoderLayer(nn.Module): - def __init__( self, num_heads: int, @@ -226,7 +225,6 @@ def get_1d_sincos_pos_embed(embed_dim, t_size, cls_token=False): class Learnable2DInterpPosEmbDivided_fixed(nn.Module): - def __init__( self, height: int, @@ -370,7 +368,6 @@ class Rope2DPosEmbRepeated(nn.Module): class MoonVision3dPatchEmbed(nn.Module): - def __init__( self, out_dim: int, @@ -382,14 +379,14 @@ class MoonVision3dPatchEmbed(nn.Module): pos_emb_type: str = "divided_fixed", ): super().__init__() - assert isinstance( - patch_size, int | Sequence - ), f"Invalid patch_size type: {type(patch_size)}" + assert isinstance(patch_size, int | Sequence), ( + f"Invalid patch_size type: {type(patch_size)}" + ) if isinstance(patch_size, int): patch_size = (patch_size, patch_size) - assert ( - len(patch_size) == 2 - ), f"Expected patch_size to be a tuple of 2, got {patch_size}" + assert len(patch_size) == 2, ( + f"Expected patch_size to be a tuple of 2, got {patch_size}" + ) self.patch_size = patch_size self.proj = Conv2dLayer( @@ -436,9 +433,9 @@ class MoonViT3dEncoder(nn.Module): ) -> None: super().__init__() - assert ( - video_attn_type == "spatial_temporal" - ), f'video_attn_type must be "spatial_temporal", got {video_attn_type}' + assert video_attn_type == "spatial_temporal", ( + f'video_attn_type must be "spatial_temporal", got {video_attn_type}' + ) self.video_attn_type = video_attn_type self.rope_2d = Rope2DPosEmbRepeated( block_cfg["hidden_dim"] // block_cfg["num_heads"], 512, 512 diff --git a/python/sglang/srt/models/kimi_linear.py b/python/sglang/srt/models/kimi_linear.py index d1ab4a1f6..1b010b42f 100644 --- a/python/sglang/srt/models/kimi_linear.py +++ b/python/sglang/srt/models/kimi_linear.py @@ -687,9 +687,9 @@ class KimiLinearModel(nn.Module): self.norm = PPMissingLayer() world_size = get_parallel().tp_size - assert ( - config.num_attention_heads % world_size == 0 - ), "num_attention_heads must be divisible by world_size" + assert config.num_attention_heads % world_size == 0, ( + "num_attention_heads must be divisible by world_size" + ) def forward( self, diff --git a/python/sglang/srt/models/kimi_vl.py b/python/sglang/srt/models/kimi_vl.py index 5cef3f977..92aae68fa 100644 --- a/python/sglang/srt/models/kimi_vl.py +++ b/python/sglang/srt/models/kimi_vl.py @@ -88,7 +88,6 @@ class MaxImageTokenMeta: class KimiVLMultiModalProjector(nn.Module): - def __init__(self, config: KimiVLConfig): super().__init__() @@ -357,7 +356,7 @@ def get_spec_layer_idx_from_weight_name( ): layer_idx = config.num_hidden_layers for i in range(config.num_nextn_predict_layers): - if weight_name.startswith(f"model.layers.{layer_idx+i}."): + if weight_name.startswith(f"model.layers.{layer_idx + i}."): return layer_idx + i return None diff --git a/python/sglang/srt/models/kimi_vl_moonvit.py b/python/sglang/srt/models/kimi_vl_moonvit.py index ac4276584..e166b382d 100644 --- a/python/sglang/srt/models/kimi_vl_moonvit.py +++ b/python/sglang/srt/models/kimi_vl_moonvit.py @@ -103,7 +103,6 @@ def apply_rope( class Learnable2DInterpPosEmb(nn.Module): - def __init__( self, height: int, width: int, dim: int, interpolation_mode: str = "bicubic" ) -> None: @@ -159,7 +158,6 @@ class Learnable2DInterpPosEmb(nn.Module): class MoonVisionPatchEmbed(nn.Module): - def __init__( self, out_dim: int, @@ -169,14 +167,14 @@ class MoonVisionPatchEmbed(nn.Module): pos_emb_width: int = 14, ): super().__init__() - assert isinstance( - patch_size, (int, Sequence) - ), f"Invalid patch_size type: {type(patch_size)}" + assert isinstance(patch_size, (int, Sequence)), ( + f"Invalid patch_size type: {type(patch_size)}" + ) if isinstance(patch_size, int): patch_size = (patch_size, patch_size) - assert ( - len(patch_size) == 2 - ), f"Expected patch_size to be a tuple of 2, got {patch_size}" + assert len(patch_size) == 2, ( + f"Expected patch_size to be a tuple of 2, got {patch_size}" + ) self.patch_size = patch_size self.proj = Conv2dLayer( @@ -402,7 +400,6 @@ class MLP2(nn.Module): class MoonVitEncoderLayer(nn.Module): - def __init__( self, num_heads: int, @@ -475,7 +472,6 @@ class MoonVitEncoderLayer(nn.Module): class MoonVitEncoder(nn.Module): - def __init__( self, hidden_dim: int, @@ -604,7 +600,6 @@ def tpool_patch_merger( class MoonVitVLProjector(nn.Module): - def __init__( self, in_channels: int, diff --git a/python/sglang/srt/models/llada2.py b/python/sglang/srt/models/llada2.py index 626914dfe..1e4e804c4 100644 --- a/python/sglang/srt/models/llada2.py +++ b/python/sglang/srt/models/llada2.py @@ -256,7 +256,9 @@ class LLaDA2MoeSparseMoeBlock(nn.Module): self.score_function == "softmax" and self.correction_bias is None ) or ( self.score_function == "sigmoid" and self.correction_bias is not None - ), "score_function and correction_bias should be in 2 combination (softmax, None) or (sigmoid, not None)" + ), ( + "score_function and correction_bias should be in 2 combination (softmax, None) or (sigmoid, not None)" + ) self.topk = TopK( top_k=self.top_k, @@ -707,7 +709,6 @@ class LLaDA2MoeBlock(nn.Module): class LLaDA2MoeModel(nn.Module): - def __init__( self, config: PretrainedConfig, diff --git a/python/sglang/srt/models/llama.py b/python/sglang/srt/models/llama.py index c3349251d..ad6aeb8bc 100644 --- a/python/sglang/srt/models/llama.py +++ b/python/sglang/srt/models/llama.py @@ -103,8 +103,7 @@ class LlamaMLP(nn.Module): ) if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() @@ -485,7 +484,7 @@ class LlamaModel(nn.Module): layer_self_attn.attn.v_scale = scaling_factor else: raise RuntimeError( - "Self attention has no KV cache scaling " "factor attribute!" + "Self attention has no KV cache scaling factor attribute!" ) def get_input_embeddings(self) -> nn.Embedding: diff --git a/python/sglang/srt/models/llama4.py b/python/sglang/srt/models/llama4.py index 8d16ca406..d7d8a26e9 100644 --- a/python/sglang/srt/models/llama4.py +++ b/python/sglang/srt/models/llama4.py @@ -70,7 +70,6 @@ logger = logging.getLogger(__name__) class Llama4MoE(nn.Module): - @torch.compile(dynamic=True, backend=get_compiler_backend()) @staticmethod def custom_routing_function( @@ -195,7 +194,6 @@ def _get_or_create_alt_stream(device_module): class Llama4Attention(nn.Module): - def __init__( self, config: Llama4TextConfig, diff --git a/python/sglang/srt/models/llama_classification.py b/python/sglang/srt/models/llama_classification.py index 438b8b32c..5441af73b 100644 --- a/python/sglang/srt/models/llama_classification.py +++ b/python/sglang/srt/models/llama_classification.py @@ -59,9 +59,9 @@ class LlamaForClassification(nn.Module): input_embeds: torch.Tensor = None, get_embedding: bool = True, ) -> EmbeddingPoolerOutput: - assert ( - get_embedding - ), "LlamaForClassification is only used for embedding. Please add --is-embedding when you launch the server." + assert get_embedding, ( + "LlamaForClassification is only used for embedding. Please add --is-embedding when you launch the server." + ) hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) return score_and_pool( diff --git a/python/sglang/srt/models/llama_embedding.py b/python/sglang/srt/models/llama_embedding.py index ba448f7fc..2b265e084 100644 --- a/python/sglang/srt/models/llama_embedding.py +++ b/python/sglang/srt/models/llama_embedding.py @@ -33,9 +33,9 @@ class LlamaEmbeddingModel(nn.Module): input_embeds: torch.Tensor = None, get_embedding: bool = True, ) -> EmbeddingPoolerOutput: - assert ( - get_embedding - ), "LlamaEmbeddingModel / MistralModel is only used for embedding" + assert get_embedding, ( + "LlamaEmbeddingModel / MistralModel is only used for embedding" + ) hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) return self.pooler(hidden_states, forward_batch) diff --git a/python/sglang/srt/models/llama_reward.py b/python/sglang/srt/models/llama_reward.py index 0eb5995f3..b263c4ae0 100644 --- a/python/sglang/srt/models/llama_reward.py +++ b/python/sglang/srt/models/llama_reward.py @@ -58,9 +58,9 @@ class LlamaForSequenceClassification(nn.Module): input_embeds: torch.Tensor = None, get_embedding: bool = True, ) -> EmbeddingPoolerOutput: - assert ( - get_embedding - ), "LlamaForSequenceClassification is only used for embedding" + assert get_embedding, ( + "LlamaForSequenceClassification is only used for embedding" + ) hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) last_token_hidden = self.pooler(hidden_states, forward_batch).embeddings @@ -110,9 +110,9 @@ class LlamaForSequenceClassificationWithNormal_Weights(LlamaForSequenceClassific input_embeds: torch.Tensor = None, get_embedding: bool = True, ) -> EmbeddingPoolerOutput: - assert ( - get_embedding - ), "LlamaForSequenceClassification is only used for embedding" + assert get_embedding, ( + "LlamaForSequenceClassification is only used for embedding" + ) hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) logits = self.score(hidden_states) weights = self.weights(hidden_states) diff --git a/python/sglang/srt/models/longcat_flash.py b/python/sglang/srt/models/longcat_flash.py index 870c94079..e2571ba28 100644 --- a/python/sglang/srt/models/longcat_flash.py +++ b/python/sglang/srt/models/longcat_flash.py @@ -183,8 +183,7 @@ class LongcatFlashMLP(nn.Module): ) if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() @@ -247,7 +246,6 @@ class LongcatFlashRouter(nn.Module): class LongcatFlashMoE(nn.Module): - def __init__( self, config: LongcatFlashConfig, @@ -357,7 +355,6 @@ class LongcatFlashMoE(nn.Module): class LongcatFlashDecoderLayer(nn.Module): - def __init__( self, config: LongcatFlashConfig, diff --git a/python/sglang/srt/models/longcat_flash_nextn.py b/python/sglang/srt/models/longcat_flash_nextn.py index a352e52d5..0a799b56d 100644 --- a/python/sglang/srt/models/longcat_flash_nextn.py +++ b/python/sglang/srt/models/longcat_flash_nextn.py @@ -107,7 +107,6 @@ logger = logging.getLogger(__name__) class LongcatFlashDenseDecoderLayer(nn.Module): - def __init__( self, config: LongcatFlashConfig, @@ -284,7 +283,6 @@ class LongcatFlashModelNextN(nn.Module): class LongcatFlashForCausalLMNextN(LongcatFlashForCausalLM): - def __init__( self, config: LongcatFlashConfig, diff --git a/python/sglang/srt/models/midashenglm.py b/python/sglang/srt/models/midashenglm.py index f64fec5e5..2f10d3a30 100644 --- a/python/sglang/srt/models/midashenglm.py +++ b/python/sglang/srt/models/midashenglm.py @@ -33,9 +33,9 @@ _Tuple2: TypeAlias = int | tuple[int, int] | Sequence[int] def _resolve_tuple2(x: _Tuple2) -> tuple[int, int]: if isinstance(x, collections.abc.Sequence): - assert ( - len(x) == 2 - ), f"Expected a sequence of length 2, got {x} with length {len(x)}" + assert len(x) == 2, ( + f"Expected a sequence of length 2, got {x} with length {len(x)}" + ) return cast(tuple[int, int], tuple(x)) return (x, x) @@ -378,9 +378,9 @@ class DashengAudioTransformer(nn.Module): t = x.shape[-1] input_splits = x.split(target_length_in_patches, dim=-1) if x_length is not None: - assert len(x_length) == len( - x - ), "batchsizes of input x and x_length need to be same" + assert len(x_length) == len(x), ( + "batchsizes of input x and x_length need to be same" + ) assert x_length.ndim == 1, "Lengths are of size (B,)" scaled_lengths = (x_length / (self.hop_length * 4)).long() mask = self._to_mask(max_length=t, lengths=scaled_lengths) diff --git a/python/sglang/srt/models/mimo_mtp.py b/python/sglang/srt/models/mimo_mtp.py index d25b87c59..b0dade237 100644 --- a/python/sglang/srt/models/mimo_mtp.py +++ b/python/sglang/srt/models/mimo_mtp.py @@ -23,7 +23,6 @@ from sglang.srt.runtime_context import get_parallel class MiMoMultiTokenPredictorLayer(nn.Module): - def __init__( self, config: PretrainedConfig, diff --git a/python/sglang/srt/models/mimo_v2.py b/python/sglang/srt/models/mimo_v2.py index 315adf08a..eb02bca01 100644 --- a/python/sglang/srt/models/mimo_v2.py +++ b/python/sglang/srt/models/mimo_v2.py @@ -318,8 +318,7 @@ class MiMoV2MLP(nn.Module): ) if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() @@ -372,7 +371,6 @@ class MoEGate(nn.Module): class MiMoV2MoE(nn.Module): - def __init__( self, config: MiMoV2Config, @@ -1133,7 +1131,7 @@ class MiMoV2Model(nn.Module): layer_self_attn.attn.v_scale = scaling_factor else: raise RuntimeError( - "Self attention has no KV cache scaling " "factor attribute!" + "Self attention has no KV cache scaling factor attribute!" ) @@ -1235,9 +1233,9 @@ class MiMoV2ForCausalLM(nn.Module, AudioEncoderMixin): return self._routed_experts_weights_of_layer.value def get_input_embedding(self, input_ids: torch.Tensor) -> torch.Tensor: - assert ( - self.model is not None - ), "get_input_embedding() is not available in encoder_only mode" + assert self.model is not None, ( + "get_input_embedding() is not available in encoder_only mode" + ) return self.model.get_input_embedding(input_ids) def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs): @@ -1351,9 +1349,9 @@ class MiMoV2ForCausalLM(nn.Module, AudioEncoderMixin): input_embeds: torch.Tensor = None, pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> torch.Tensor: - assert ( - not self.config.encoder_only - ), "forward() should not be called in encoder_only mode" + assert not self.config.encoder_only, ( + "forward() should not be called in encoder_only mode" + ) if self._is_multimodal: hidden_states, hidden_states_before_norm = general_mm_embed_routine( @@ -1624,15 +1622,15 @@ class MiMoV2ForCausalLM(nn.Module, AudioEncoderMixin): ) def get_embed_and_head(self): - assert ( - self.model is not None and self.lm_head is not None - ), "get_embed_and_head() is not available in encoder_only mode" + assert self.model is not None and self.lm_head is not None, ( + "get_embed_and_head() is not available in encoder_only mode" + ) return self.model.embed_tokens.weight, self.lm_head.weight def set_embed_and_head(self, embed, head): - assert ( - self.model is not None and self.lm_head is not None - ), "set_embed_and_head() is not available in encoder_only mode" + assert self.model is not None and self.lm_head is not None, ( + "set_embed_and_head() is not available in encoder_only mode" + ) del self.model.embed_tokens.weight del self.lm_head.weight self.model.embed_tokens.weight = embed diff --git a/python/sglang/srt/models/mimo_v2_nextn.py b/python/sglang/srt/models/mimo_v2_nextn.py index 1eaa1f453..6c3943260 100644 --- a/python/sglang/srt/models/mimo_v2_nextn.py +++ b/python/sglang/srt/models/mimo_v2_nextn.py @@ -238,7 +238,6 @@ class MiMoV2ModelNextN(nn.Module): class MiMoV2MTP(MiMoV2ForCausalLM): - def __init__( self, config: PretrainedConfig, @@ -320,7 +319,6 @@ class MiMoV2MTP(MiMoV2ForCausalLM): continue for param_name, weight_name, shard_id in stacked_params_mapping: - if f".{weight_name}." not in name: continue if "mtp_block" not in name: diff --git a/python/sglang/srt/models/minicpm.py b/python/sglang/srt/models/minicpm.py index 3ec0fff8f..b28ab4392 100644 --- a/python/sglang/srt/models/minicpm.py +++ b/python/sglang/srt/models/minicpm.py @@ -73,8 +73,7 @@ class MiniCPMMLP(nn.Module): ) if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() diff --git a/python/sglang/srt/models/minicpm3.py b/python/sglang/srt/models/minicpm3.py index b9374dfe2..4ff7f9fa4 100644 --- a/python/sglang/srt/models/minicpm3.py +++ b/python/sglang/srt/models/minicpm3.py @@ -72,8 +72,7 @@ class MiniCPM3MLP(nn.Module): ) if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() @@ -94,7 +93,6 @@ def input_to_float8(x, dtype=torch.float8_e4m3fn): class MiniCPM3AttentionMLA(nn.Module): - def __init__( self, config: PretrainedConfig, diff --git a/python/sglang/srt/models/minicpmo.py b/python/sglang/srt/models/minicpmo.py index f9781dbf2..b9f8c4286 100644 --- a/python/sglang/srt/models/minicpmo.py +++ b/python/sglang/srt/models/minicpmo.py @@ -1188,7 +1188,6 @@ class MiniCPMWhisperEncoderLayer(nn.Module): # Copied from from transformers.models.whisper.modeling_whisper.WhisperEncoder and add use_cache for streaming inference class MiniCPMWhisperEncoder(WhisperEncoder): - def __init__(self, config: WhisperConfig): super().__init__(config) self.layers = nn.ModuleList( @@ -1355,9 +1354,9 @@ class MiniCPMWhisperEncoder(WhisperEncoder): # check if head_mask has a correct number of layers specified if desired if head_mask is not None: - assert head_mask.size()[0] == ( - len(self.layers) - ), f"The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}." + assert head_mask.size()[0] == (len(self.layers)), ( + f"The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}." + ) for idx, encoder_layer in enumerate(self.layers): if output_hidden_states: @@ -1457,9 +1456,9 @@ class MiniCPMO(MiniCPMBaseModel): logger.info("TTS is disabled for now") if self.config.init_tts: # print("tts enabled") - assert ( - _tts_deps - ), "please make sure vector_quantize_pytorch and vocos are installed." + assert _tts_deps, ( + "please make sure vector_quantize_pytorch and vocos are installed." + ) self.tts = self.init_tts_module() def init_tts_module(self): @@ -1878,7 +1877,6 @@ class MiniCPMO(MiniCPMBaseModel): params_dict = dict(self.named_parameters()) for name, loaded_weight in weights: - if "rotary_emb.inv_freq~" in name or "projector" in name: continue if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: diff --git a/python/sglang/srt/models/minicpmv.py b/python/sglang/srt/models/minicpmv.py index 1dc9d6de0..b8fd89a27 100644 --- a/python/sglang/srt/models/minicpmv.py +++ b/python/sglang/srt/models/minicpmv.py @@ -269,7 +269,6 @@ class BaseResampler(nn.Module): class Resampler2_5(BaseResampler): - def __init__( self, num_queries: int, @@ -346,9 +345,7 @@ class Resampler2_5(BaseResampler): key_padding_mask[i, patch_len[i] :] = True pos_embed = torch.nn.utils.rnn.pad_sequence( pos_embed, batch_first=True, padding_value=0.0 - ).permute( - 1, 0, 2 - ) # BLD => L * B * D + ).permute(1, 0, 2) # BLD => L * B * D x, _ = self.kv_proj(x) # B * L * D x = self.ln_kv(x).permute(1, 0, 2) # L * B * D @@ -369,7 +366,6 @@ class Resampler2_5(BaseResampler): class Resampler4_5(BaseResampler): - def __init__( self, num_queries: int, @@ -522,9 +518,7 @@ class Resampler4_5(BaseResampler): pos_embed_2d = torch.nn.utils.rnn.pad_sequence( pos_embed_2d, batch_first=True, padding_value=0.0 - ).permute( - 1, 0, 2 - ) # BLD => L * B * D + ).permute(1, 0, 2) # BLD => L * B * D k = x v = x + pos_embed_2d @@ -554,14 +548,10 @@ class Resampler4_5(BaseResampler): k = torch.nn.utils.rnn.pad_sequence( merge_k, batch_first=True, padding_value=0.0 - ).permute( - 1, 0, 2 - ) # L*(end-start) + ).permute(1, 0, 2) # L*(end-start) v = torch.nn.utils.rnn.pad_sequence( merge_v, batch_first=True, padding_value=0.0 - ).permute( - 1, 0, 2 - ) # L*(end-start) + ).permute(1, 0, 2) # L*(end-start) key_padding_mask = torch.nn.utils.rnn.pad_sequence( merge_key_padding_mask, batch_first=True, padding_value=True ).squeeze(-1) @@ -725,8 +715,7 @@ class MiniCPMBaseModel(nn.Module): ) if not isinstance(image_embeds, (torch.Tensor, list)): raise ValueError( - f"Incorrect type of image embeds. " - f"Got type: {type(image_embeds)}" + f"Incorrect type of image embeds. Got type: {type(image_embeds)}" ) if isinstance(image_embeds, list): diff --git a/python/sglang/srt/models/minimax_m2.py b/python/sglang/srt/models/minimax_m2.py index 8d959e8d8..54037429c 100644 --- a/python/sglang/srt/models/minimax_m2.py +++ b/python/sglang/srt/models/minimax_m2.py @@ -299,15 +299,15 @@ class MiniMaxM2RMSNormTP(nn.Module): # Align with QKVParallelLinear pattern if self.attn_tp_size >= num_heads: - assert ( - self.attn_tp_size % num_heads == 0 - ), f"attn_tp_size ({self.attn_tp_size}) must be divisible by num_heads ({num_heads})" + assert self.attn_tp_size % num_heads == 0, ( + f"attn_tp_size ({self.attn_tp_size}) must be divisible by num_heads ({num_heads})" + ) self.num_heads = 1 self.num_head_replicas = self.attn_tp_size // num_heads else: - assert ( - num_heads % self.attn_tp_size == 0 - ), f"num_heads ({num_heads}) must be divisible by attn_tp_size ({self.attn_tp_size})" + assert num_heads % self.attn_tp_size == 0, ( + f"num_heads ({num_heads}) must be divisible by attn_tp_size ({self.attn_tp_size})" + ) self.num_heads = num_heads // self.attn_tp_size self.num_head_replicas = 1 @@ -854,9 +854,9 @@ class MiniMaxM2Attention(nn.Module): forward_batch: ForwardBatch, ): if hidden_states.shape[0] == 0: - assert ( - not self.o_proj.reduce_results - ), "short-circuiting allreduce will lead to hangs" + assert not self.o_proj.reduce_results, ( + "short-circuiting allreduce will lead to hangs" + ) return hidden_states, forward_batch, None qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) @@ -873,9 +873,9 @@ class MiniMaxM2Attention(nn.Module): forward_batch: ForwardBatch, ): if hidden_states.shape[0] == 0: - assert ( - not self.o_proj.reduce_results - ), "short-circuiting allreduce will lead to hangs" + assert not self.o_proj.reduce_results, ( + "short-circuiting allreduce will lead to hangs" + ) return hidden_states, forward_batch, None qkv, _ = self.qkv_proj(hidden_states) if self.use_qk_norm: diff --git a/python/sglang/srt/models/minimax_m3.py b/python/sglang/srt/models/minimax_m3.py index 5e9c0f42f..786fb11ee 100644 --- a/python/sglang/srt/models/minimax_m3.py +++ b/python/sglang/srt/models/minimax_m3.py @@ -1071,9 +1071,9 @@ class MiniMaxM3Attention(nn.Module): ): """NPU qkv projection + fused norm/RoPE/split; returns (None, fb, inner_state).""" if hidden_states.shape[0] == 0: - assert ( - not self.o_proj.reduce_results - ), "short-circuiting allreduce will lead to hangs" + assert not self.o_proj.reduce_results, ( + "short-circuiting allreduce will lead to hangs" + ) return hidden_states, forward_batch, None qkv, _ = self.qkv_proj(hidden_states) @@ -1612,9 +1612,9 @@ class MiniMaxM3SparseForCausalLM(nn.Module): if is_shared_experts_fusion_disabled(): return self.num_fused_shared_experts = self.config.n_shared_experts - assert ( - self.num_fused_shared_experts == 1 - ), "Only 1 fused shared expert is supported for MiniMax-M3" + assert self.num_fused_shared_experts == 1, ( + "Only 1 fused shared expert is supported for MiniMax-M3" + ) log_info_on_rank0(logger, "Shared experts fusion optimization enabled.") def set_eagle3_layers_to_capture(self, layer_ids: Optional[list[int]] = None): diff --git a/python/sglang/srt/models/minimax_m3_vl.py b/python/sglang/srt/models/minimax_m3_vl.py index 2fdff0370..5eed7654c 100644 --- a/python/sglang/srt/models/minimax_m3_vl.py +++ b/python/sglang/srt/models/minimax_m3_vl.py @@ -172,9 +172,9 @@ class MiniMaxM3SparseForConditionalGeneration(nn.Module): if is_shared_experts_fusion_disabled(): return self.num_fused_shared_experts = self.config.text_config.n_shared_experts - assert ( - self.num_fused_shared_experts == 1 - ), "Only 1 fused shared expert is supported" + assert self.num_fused_shared_experts == 1, ( + "Only 1 fused shared expert is supported" + ) log_info_on_rank0(logger, "Shared experts fusion optimization enabled.") @classmethod diff --git a/python/sglang/srt/models/minimax_vl_common.py b/python/sglang/srt/models/minimax_vl_common.py index e5a24013f..333076efd 100644 --- a/python/sglang/srt/models/minimax_vl_common.py +++ b/python/sglang/srt/models/minimax_vl_common.py @@ -93,9 +93,9 @@ class MiniMaxVLMultiModalProjector(nn.Module): tp_size=tp_size, tp_rank=tp_rank, ) - assert ( - projector_hidden_act == "gelu" - ), f"Only gelu activation is supported, got {projector_hidden_act}" + assert projector_hidden_act == "gelu", ( + f"Only gelu activation is supported, got {projector_hidden_act}" + ) self.act = get_act_fn(projector_hidden_act) self.linear_2 = RowParallelLinear( mid_size, @@ -148,9 +148,9 @@ class MiniMaxVLPatchMerger(nn.Module): tp_size=tp_size, tp_rank=tp_rank, ) - assert ( - projector_hidden_act == "gelu" - ), f"Only gelu activation is supported, got {projector_hidden_act}" + assert projector_hidden_act == "gelu", ( + f"Only gelu activation is supported, got {projector_hidden_act}" + ) self.act = get_act_fn(projector_hidden_act) self.linear_2 = RowParallelLinear( mid_size, @@ -230,9 +230,9 @@ class CLIPVisionEmbeddings(nn.Module): if self.patch_embedding.weight.dtype != pixel_values.dtype: self.patch_embedding = self.patch_embedding.to(pixel_values.dtype) - assert ( - pixel_values.dim() == 2 - ), f"pixel_values must be 2D, got {pixel_values.dim()}D" + assert pixel_values.dim() == 2, ( + f"pixel_values must be 2D, got {pixel_values.dim()}D" + ) pixel_values = pixel_values.reshape( pixel_values.shape[0], self.input_num_channels, @@ -285,9 +285,9 @@ class CLIPEncoderLayer(nn.Module): tp_rank=tp_rank, ) hidden_act = getattr(config, "hidden_act", "gelu") - assert ( - hidden_act == "gelu" - ), f"Only gelu activation is supported, got {hidden_act}" + assert hidden_act == "gelu", ( + f"Only gelu activation is supported, got {hidden_act}" + ) self.act = get_act_fn(hidden_act) self.fc2 = RowParallelLinear( config.intermediate_size, @@ -430,9 +430,9 @@ class MiniMaxVLVisionTransformer(nn.Module): workspace_buffer=workspace_buffer, ) - assert ( - self.config.position_embedding_type == "rope" - ), "Only rope position embedding is supported" + assert self.config.position_embedding_type == "rope", ( + "Only rope position embedding is supported" + ) assert self.config.rope_mode == "3d", "Only 3D RoPE is supported" rope_theta = getattr(config, "rope_theta") assert rope_theta is not None, "rope_theta must be set" @@ -672,9 +672,9 @@ class MiniMaxVLVisionTransformer(nn.Module): cu_seq_len = self._compute_cu_seq_len(grid_thw, hidden_states.device) rotary_pos_emb = self._get_rope_embed_3d(grid_thw, self.spatial_merge_size) - assert ( - rotary_pos_emb.device == hidden_states.device - ), "rotary_pos_emb and hidden_states must be on the same device" + assert rotary_pos_emb.device == hidden_states.device, ( + "rotary_pos_emb and hidden_states must be on the same device" + ) max_seqlen: Optional[int] = None sequence_lengths: Optional[torch.Tensor] = None diff --git a/python/sglang/srt/models/mistral_eagle.py b/python/sglang/srt/models/mistral_eagle.py index 434856026..1e22022f9 100644 --- a/python/sglang/srt/models/mistral_eagle.py +++ b/python/sglang/srt/models/mistral_eagle.py @@ -65,9 +65,9 @@ class MistralEagleModel(nn.Module): super().__init__() self.config = config self.vocab_size = config.vocab_size - assert ( - get_pp_group().world_size == 1 - ), "MistralForCausalLMEagle currently does not support pipeline parallelism" + assert get_pp_group().world_size == 1, ( + "MistralForCausalLMEagle currently does not support pipeline parallelism" + ) self.pp_group = get_pp_group() self.embed_tokens = VocabParallelEmbedding( config.vocab_size, diff --git a/python/sglang/srt/models/mixtral.py b/python/sglang/srt/models/mixtral.py index e797b7daf..5e34952cd 100644 --- a/python/sglang/srt/models/mixtral.py +++ b/python/sglang/srt/models/mixtral.py @@ -336,7 +336,6 @@ class MixtralModel(nn.Module): class MixtralForCausalLM(nn.Module): - def __init__( self, config: MixtralConfig, diff --git a/python/sglang/srt/models/mllama.py b/python/sglang/srt/models/mllama.py index 2735fe344..b76ff7b59 100644 --- a/python/sglang/srt/models/mllama.py +++ b/python/sglang/srt/models/mllama.py @@ -85,7 +85,6 @@ class ColumnParallelConv2dPatch(torch.nn.Module): class MllamaPrecomputedAspectRatioEmbedding(nn.Module): - def __init__(self, config: config_mllama.MllamaVisionConfig, is_gated: bool = True): super().__init__() self.max_num_tiles = config.max_num_tiles @@ -870,7 +869,6 @@ class MllamaForConditionalGeneration(nn.Module): # pixel_values: shape (bs, num_image, num_tiles, 3, image_res, image_res) max_num_images = max_num_tiles = bs = 0 for i, mm_input in enumerate(forward_batch.mm_inputs): - if not forward_batch.encoder_cached[i] and mm_input is not None: pixel_values = torch.cat( [item.feature for item in mm_input.mm_items], dim=0 diff --git a/python/sglang/srt/models/mllama4.py b/python/sglang/srt/models/mllama4.py index 28ddcdc87..34beef383 100644 --- a/python/sglang/srt/models/mllama4.py +++ b/python/sglang/srt/models/mllama4.py @@ -48,7 +48,6 @@ logger = logging.getLogger(__name__) class Llama4VisionMLP(nn.Module): - def __init__( self, input_size: int, @@ -115,7 +114,6 @@ def pixel_shuffle(input_tensor, shuffle_ratio): class Llama4VisionPixelShuffleMLP(nn.Module): - def __init__( self, config, @@ -155,7 +153,6 @@ def apply_position_embedding(q, k, freqs_ci, shape): class Llama4VisionEncoderLayer(nn.Module): - def __init__( self, config: Llama4VisionConfig, @@ -226,7 +223,6 @@ class Llama4VisionEncoderLayer(nn.Module): class Llama4VisionEncoder(nn.Module): - def __init__( self, config: Llama4VisionConfig, @@ -272,7 +268,6 @@ class Llama4VisionEncoder(nn.Module): class Llama4UnfoldConvolution(nn.Module): - def __init__( self, config: Llama4VisionConfig, @@ -350,7 +345,6 @@ class Llama4VisionRotaryEmbedding(nn.Module): class Llama4VisionModel(nn.Module): - def __init__( self, config: Llama4VisionConfig, diff --git a/python/sglang/srt/models/moss_vl.py b/python/sglang/srt/models/moss_vl.py index 55a57111e..45c6792ce 100644 --- a/python/sglang/srt/models/moss_vl.py +++ b/python/sglang/srt/models/moss_vl.py @@ -1192,7 +1192,6 @@ class MossVLForCausalLM(nn.Module): class MossVLForConditionalGeneration(nn.Module): - def __init__(self, config, quant_config=None, prefix: str = ""): super().__init__() self.config = config diff --git a/python/sglang/srt/models/muse_glimmer.py b/python/sglang/srt/models/muse_glimmer.py index a1ddf6fbc..07d2599f3 100644 --- a/python/sglang/srt/models/muse_glimmer.py +++ b/python/sglang/srt/models/muse_glimmer.py @@ -522,7 +522,6 @@ class MuseGlimmerVisionPatchEmbedder(nn.Module): class MuseGlimmerVisionRotaryEmbedding(nn.Module): - def __init__(self, head_dim: int, theta: float): super().__init__() spatial_dim = head_dim // 2 @@ -611,7 +610,6 @@ class MuseGlimmerVisionEncoderLayer(nn.Module): class MuseGlimmerVisionModel(nn.Module): - def __init__( self, config, diff --git a/python/sglang/srt/models/nemotron_h_mtp.py b/python/sglang/srt/models/nemotron_h_mtp.py index 3c375e90c..5df677188 100644 --- a/python/sglang/srt/models/nemotron_h_mtp.py +++ b/python/sglang/srt/models/nemotron_h_mtp.py @@ -231,9 +231,9 @@ class NemotronHMultiTokenPredictor(nn.Module): self.mtp_start_layer_idx = config.num_hidden_layers self.num_mtp_layers = getattr(config, "num_nextn_predict_layers", 1) - assert ( - self.num_mtp_layers == 1 - ), "Only one MTP layer is supported for NemotronH-MTP" + assert self.num_mtp_layers == 1, ( + "Only one MTP layer is supported for NemotronH-MTP" + ) self.pattern_str = config.mtp_hybrid_override_pattern self.pattern_len = len(self.pattern_str) @@ -280,9 +280,9 @@ class NemotronHMultiTokenPredictor(nn.Module): ) def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: - assert ( - self.embed_tokens is not None - ), "embed_tokens not initialized - must be shared from target model" + assert self.embed_tokens is not None, ( + "embed_tokens not initialized - must be shared from target model" + ) return self.embed_tokens(input_ids) def forward( diff --git a/python/sglang/srt/models/nemotron_nas.py b/python/sglang/srt/models/nemotron_nas.py index 60341c74f..f45e16591 100644 --- a/python/sglang/srt/models/nemotron_nas.py +++ b/python/sglang/srt/models/nemotron_nas.py @@ -58,7 +58,6 @@ def _find_multiple(n: int, k: int) -> int: class DeciLMDecoderLayer(nn.Module): - def __init__( self, config: LlamaConfig, diff --git a/python/sglang/srt/models/nvila.py b/python/sglang/srt/models/nvila.py index dce329f77..30505a46a 100644 --- a/python/sglang/srt/models/nvila.py +++ b/python/sglang/srt/models/nvila.py @@ -327,9 +327,9 @@ def merge_features_for_dynamic_s2( ) ) - assert block_cnt == len( - image_features - ), f"The number of blocks ({block_cnt}) does not match length of image_features ({len(image_features)})!" + assert block_cnt == len(image_features), ( + f"The number of blocks ({block_cnt}) does not match length of image_features ({len(image_features)})!" + ) return image_features_each_image, new_block_sizes diff --git a/python/sglang/srt/models/olmo.py b/python/sglang/srt/models/olmo.py index 15106fd76..e1ff3e5d2 100644 --- a/python/sglang/srt/models/olmo.py +++ b/python/sglang/srt/models/olmo.py @@ -233,7 +233,6 @@ class OlmoDecoderLayer(nn.Module): class OlmoModel(nn.Module): - def __init__( self, config: OlmoConfig, diff --git a/python/sglang/srt/models/olmo2.py b/python/sglang/srt/models/olmo2.py index 85d5f507e..f606efaf0 100644 --- a/python/sglang/srt/models/olmo2.py +++ b/python/sglang/srt/models/olmo2.py @@ -322,7 +322,6 @@ class Olmo2DecoderLayer(nn.Module): class Olmo2Model(nn.Module): - def __init__( self, config: PretrainedConfig, diff --git a/python/sglang/srt/models/olmoe.py b/python/sglang/srt/models/olmoe.py index df7774b8a..c66156b33 100644 --- a/python/sglang/srt/models/olmoe.py +++ b/python/sglang/srt/models/olmoe.py @@ -107,7 +107,6 @@ class OlmoeMoE(nn.Module): class OlmoeAttention(nn.Module): - def __init__( self, layer_id: int, @@ -196,7 +195,6 @@ class OlmoeAttention(nn.Module): class OlmoeDecoderLayer(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -261,7 +259,6 @@ class OlmoeDecoderLayer(nn.Module): class OlmoeModel(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -311,7 +308,6 @@ class OlmoeModel(nn.Module): class OlmoeForCausalLM(nn.Module): - fall_back_to_pt_during_load = False def __init__( diff --git a/python/sglang/srt/models/opt.py b/python/sglang/srt/models/opt.py index 0a6cc37f3..3ac98bfd2 100644 --- a/python/sglang/srt/models/opt.py +++ b/python/sglang/srt/models/opt.py @@ -72,7 +72,6 @@ def get_activation(name="relu"): class OPTLearnedPositionalEmbedding(nn.Embedding): - def __init__(self, num_embeddings: int, embedding_dim: int): # OPT is set up so that if padding_idx is specified then offset the # embedding ids by 2 and adjust num_embeddings appropriately. Other @@ -85,7 +84,6 @@ class OPTLearnedPositionalEmbedding(nn.Embedding): class OPTAttention(nn.Module): - def __init__( self, embed_dim: int, @@ -143,7 +141,6 @@ class OPTAttention(nn.Module): class OPTDecoderLayer(nn.Module): - def __init__( self, config: OPTConfig, @@ -220,7 +217,6 @@ class OPTDecoderLayer(nn.Module): class OPTDecoder(nn.Module): - def __init__( self, config: OPTConfig, @@ -324,7 +320,6 @@ class OPTDecoder(nn.Module): class OPTModel(nn.Module): - def __init__( self, config: OPTConfig, @@ -380,7 +375,7 @@ class OPTModel(nn.Module): layer_self_attn.attn.v_scale = scaling_factor else: raise RuntimeError( - "Self attention has no KV cache scaling " "factor attribute!" + "Self attention has no KV cache scaling factor attribute!" ) diff --git a/python/sglang/srt/models/orion.py b/python/sglang/srt/models/orion.py index 7eb2bc7cb..9d34fa921 100644 --- a/python/sglang/srt/models/orion.py +++ b/python/sglang/srt/models/orion.py @@ -64,8 +64,7 @@ class OrionMLP(nn.Module): ) if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() diff --git a/python/sglang/srt/models/paddleocr_vl.py b/python/sglang/srt/models/paddleocr_vl.py index 1e375ad16..cb007079e 100644 --- a/python/sglang/srt/models/paddleocr_vl.py +++ b/python/sglang/srt/models/paddleocr_vl.py @@ -172,7 +172,6 @@ class Projector(nn.Module): class SiglipVisionEmbeddings(nn.Module): - def __init__(self, config): super().__init__() self.config = config @@ -298,7 +297,6 @@ class SiglipVisionEmbeddings(nn.Module): class SigLIPRotaryEmbedding(nn.Module): - def __init__(self, dim: int, theta: float = 10000.0) -> None: super().__init__() self.dim = dim @@ -322,7 +320,6 @@ class SigLIPRotaryEmbedding(nn.Module): class SiglipMLP(nn.Module): - def __init__( self, config, @@ -360,7 +357,6 @@ class SiglipMLP(nn.Module): class SiglipEncoderLayer(nn.Module): - def __init__( self, config, @@ -418,7 +414,6 @@ class SiglipEncoderLayer(nn.Module): class SiglipEncoder(nn.Module): - def __init__( self, config, @@ -481,7 +476,6 @@ class SiglipEncoder(nn.Module): class SiglipVisionTransformer(nn.Module): - def __init__( self, config, @@ -574,7 +568,6 @@ class SiglipVisionModel(nn.Module): class PaddleOCRVLForConditionalGeneration(Ernie4_5_ForCausalLM): - def __init__(self, *, config, quant_config=None, prefix: str = ""): super().__init__(config=config, quant_config=quant_config, prefix=prefix) config = self.config diff --git a/python/sglang/srt/models/parakeet.py b/python/sglang/srt/models/parakeet.py index b77e3648a..8630d892b 100644 --- a/python/sglang/srt/models/parakeet.py +++ b/python/sglang/srt/models/parakeet.py @@ -97,7 +97,7 @@ class ProjectedParakeet(nn.Module): if name.startswith("sound_encoder."): target_name = name[len("sound_encoder.") :] elif name.startswith("sound_projection."): - target_name = f"projection.{name[len('sound_projection.'):]}" + target_name = f"projection.{name[len('sound_projection.') :]}" else: continue diff --git a/python/sglang/srt/models/persimmon.py b/python/sglang/srt/models/persimmon.py index 8f5a2e6c5..28f4cf250 100644 --- a/python/sglang/srt/models/persimmon.py +++ b/python/sglang/srt/models/persimmon.py @@ -28,7 +28,6 @@ from sglang.srt.utils import add_prefix, make_layers class PersimmonMLP(nn.Module): - def __init__( self, config: PersimmonConfig, quant_config: Optional[QuantizationConfig] = None ): @@ -49,7 +48,6 @@ class PersimmonMLP(nn.Module): class PersimmonAttention(nn.Module): - def __init__( self, config: PersimmonConfig, @@ -144,7 +142,6 @@ class PersimmonAttention(nn.Module): class PersimmonDecoderLayer(nn.Module): - def __init__( self, config: PersimmonConfig, @@ -196,7 +193,6 @@ class PersimmonDecoderLayer(nn.Module): class PersimmonModel(nn.Module): - def __init__( self, config: PersimmonConfig, @@ -259,7 +255,6 @@ class PersimmonModel(nn.Module): class PersimmonForCausalLM(nn.Module): - def __init__( self, config: PersimmonConfig, diff --git a/python/sglang/srt/models/phi.py b/python/sglang/srt/models/phi.py index 96867e0ca..1dda797cb 100644 --- a/python/sglang/srt/models/phi.py +++ b/python/sglang/srt/models/phi.py @@ -29,7 +29,6 @@ from sglang.srt.utils import add_prefix, make_layers class PhiAttention(nn.Module): - def __init__( self, config: PhiConfig, @@ -99,7 +98,6 @@ class PhiAttention(nn.Module): class PhiMLP(nn.Module): - def __init__( self, config: PhiConfig, quant_config: Optional[QuantizationConfig] = None ): @@ -128,7 +126,6 @@ class PhiMLP(nn.Module): class PhiLayer(nn.Module): - def __init__( self, config: PhiConfig, @@ -167,7 +164,6 @@ class PhiLayer(nn.Module): class PhiModel(nn.Module): - def __init__( self, config: PhiConfig, diff --git a/python/sglang/srt/models/phi3_small.py b/python/sglang/srt/models/phi3_small.py index e04fb6b9e..15b90ef7b 100644 --- a/python/sglang/srt/models/phi3_small.py +++ b/python/sglang/srt/models/phi3_small.py @@ -51,7 +51,6 @@ def gegelu(input, limit: Optional[float] = None): class Phi3SmallMLP(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -60,9 +59,9 @@ class Phi3SmallMLP(nn.Module): ) -> None: super().__init__() self.config = config - assert ( - self.config.hidden_act == "gegelu" - ), "Only `gegelu` is supported for the 4.7 series of models .." + assert self.config.hidden_act == "gegelu", ( + "Only `gegelu` is supported for the 4.7 series of models .." + ) self.hidden_size = config.hidden_size self.gegelu_limit = config.gegelu_limit self.intermediate_size = config.intermediate_size @@ -90,7 +89,6 @@ class Phi3SmallMLP(nn.Module): class Phi3SmallSelfAttention(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -233,7 +231,6 @@ class Phi3SmallSelfAttention(nn.Module): class Phi3SmallDecoderLayer(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -286,7 +283,6 @@ class Phi3SmallDecoderLayer(nn.Module): class Phi3SmallModel(nn.Module): - def __init__( self, config: Phi3Config, diff --git a/python/sglang/srt/models/phi4mm.py b/python/sglang/srt/models/phi4mm.py index 2b261d4ff..34f277e68 100644 --- a/python/sglang/srt/models/phi4mm.py +++ b/python/sglang/srt/models/phi4mm.py @@ -105,9 +105,9 @@ class Phi4MMImageEncoder(nn.Module): self.base_feat_height_target = self.base_feat_height_target // 2 # with_hd_transform and with_learnable_separator should have same value - assert ( - self.use_hd_transform == self.with_learnable_separator - ), "use_hd_transform and with_learnable_separator should have same value" + assert self.use_hd_transform == self.with_learnable_separator, ( + "use_hd_transform and with_learnable_separator should have same value" + ) assert self.use_hd_transform, "learnable separator is only for hd transform" # 1024 * 4, merge spatial to channel dimension self.glb_GN = nn.Parameter( @@ -210,9 +210,11 @@ class Phi4MMImageEncoder(nn.Module): assert ( base_feat_height == base_feat_height_target and base_feat_width == base_feat_height_target - ), f'base_feat_height: {base_feat_height},"\ + ), ( + f'base_feat_height: {base_feat_height},"\ f" base_feat_width: {base_feat_width}, "\ f"expect {base_feat_height_target} features for hd transform' + ) # bs x max_num_crops x (24x24) x C img_features = img_features.view( @@ -359,10 +361,10 @@ class Phi4MMImageEncoder(nn.Module): ) # temp_len = int((h*w+1)*144 + 1 + (h+1)*12) - assert ( - temp_len == output_imgs[-1].shape[1] - ), f'temp_len: {temp_len}, output_imgs[-1].shape[1]: "\ + assert temp_len == output_imgs[-1].shape[1], ( + f'temp_len: {temp_len}, output_imgs[-1].shape[1]: "\ "{output_imgs[-1].shape[1]}' + ) output_len.append(temp_len) diff --git a/python/sglang/srt/models/phi4mm_audio.py b/python/sglang/srt/models/phi4mm_audio.py index 10ccbbc09..410386352 100644 --- a/python/sglang/srt/models/phi4mm_audio.py +++ b/python/sglang/srt/models/phi4mm_audio.py @@ -383,9 +383,9 @@ class TransformerEncoderBase(abc.ABC, nn.Module): if nemo_conv_settings: default_nemo_conv_settings.update(nemo_conv_settings) for i in ["subsampling_factor", "feat_in", "feat_out"]: - assert ( - i not in nemo_conv_settings - ), "{i} should be specified outside of the NeMo dictionary" + assert i not in nemo_conv_settings, ( + "{i} should be specified outside of the NeMo dictionary" + ) self.embed = NemoConvSubsampling( **default_nemo_conv_settings, @@ -403,9 +403,9 @@ class TransformerEncoderBase(abc.ABC, nn.Module): else None ) if self.relative_attention_bias_type == "t5": - assert ( - self.num_heads % self.attention_group_size == 0 - ), "attention_group_size must divide n_head" + assert self.num_heads % self.attention_group_size == 0, ( + "attention_group_size must divide n_head" + ) self.relative_attention_bias_layer = T5RelativeAttentionLogitBias( self.num_heads // self.attention_group_size, max_distance=relative_attention_bias_args.get( @@ -477,8 +477,7 @@ class TransformerEncoderBase(abc.ABC, nn.Module): ) if len(left_chunk) != len(chunk_size): raise ValueError( - "The length of left_chunk must be the same as length of " - "chunk_size." + "The length of left_chunk must be the same as length of chunk_size." ) left_chunk_train_eff = left_chunk[chunk_size_index] else: @@ -837,9 +836,9 @@ class ConformerEncoder(TransformerEncoderBase): self.replication_pad_for_subsample_embedding: bool = ( replication_pad_for_subsample_embedding ) - assert ( - self.num_heads % attention_group_size == 0 - ), "attention_group_size must divide n_head" + assert self.num_heads % attention_group_size == 0, ( + "attention_group_size must divide n_head" + ) self.num_heads_k = self.num_heads // attention_group_size self.encoders = MultiSequential( @@ -1114,9 +1113,9 @@ class AudioEmbedding(nn.Module): self.qformer = None if kwargs.get("use_conv_downsample", False): - assert ( - self.qformer is None - ), "don't support use qformer and conv downsample together" + assert self.qformer is None, ( + "don't support use qformer and conv downsample together" + ) nemo_conv_settings = kwargs.get("nemo_conv_settings", {}) default_nemo_conv_settings = { "subsampling": "dw_striding", @@ -1132,9 +1131,9 @@ class AudioEmbedding(nn.Module): if nemo_conv_settings: default_nemo_conv_settings.update(nemo_conv_settings) for i in ["subsampling_factor", "feat_in", "feat_out"]: - assert ( - i not in nemo_conv_settings - ), "{i} should be specified outside of the NeMo dictionary" + assert i not in nemo_conv_settings, ( + "{i} should be specified outside of the NeMo dictionary" + ) self.conv_ds = NemoConvSubsampling( **default_nemo_conv_settings, @@ -1231,7 +1230,7 @@ class AudioEmbedding(nn.Module): audio_set_tensor = self.audio_projection_for_vision(audio_features) else: raise ValueError( - f"audio_projection_mode = {audio_projection_mode} not " "implemented" + f"audio_projection_mode = {audio_projection_mode} not implemented" ) return audio_set_tensor diff --git a/python/sglang/srt/models/phi4mm_utils.py b/python/sglang/srt/models/phi4mm_utils.py index e6bf35ebf..47e4b835d 100644 --- a/python/sglang/srt/models/phi4mm_utils.py +++ b/python/sglang/srt/models/phi4mm_utils.py @@ -242,8 +242,9 @@ class GLUPointWiseConv(nn.Module): x[:, self.output_dim : self.output_dim * 2, :] + self.b2 ) else: - x = (x[:, 0 : self.output_dim, :]) * ( - x[:, self.output_dim : self.output_dim * 2, :] + x = ( + (x[:, 0 : self.output_dim, :]) + * (x[:, self.output_dim : self.output_dim * 2, :]) ) else: if self.bias_in_glu: @@ -1111,7 +1112,7 @@ class NemoConvSubsampling(torch.nn.Module): and subsampling_conv_chunking_factor % 2 != 0 ): raise ValueError( - "subsampling_conv_chunking_factor should be -1, 1, or a " "power of 2" + "subsampling_conv_chunking_factor should be -1, 1, or a power of 2" ) self.subsampling_conv_chunking_factor = subsampling_conv_chunking_factor @@ -1578,7 +1579,7 @@ class NemoConvSubsampling(torch.nn.Module): and subsampling_conv_chunking_factor % 2 != 0 ): raise ValueError( - "subsampling_conv_chunking_factor should be -1, 1, or a " "power of 2" + "subsampling_conv_chunking_factor should be -1, 1, or a power of 2" ) self.subsampling_conv_chunking_factor = subsampling_conv_chunking_factor diff --git a/python/sglang/srt/models/phimoe.py b/python/sglang/srt/models/phimoe.py index 65246ede4..f9f2ab06f 100644 --- a/python/sglang/srt/models/phimoe.py +++ b/python/sglang/srt/models/phimoe.py @@ -31,7 +31,6 @@ from sglang.srt.utils import add_prefix, make_layers class PhiMoEConfig(PretrainedConfig): - model_type = "phimoe" def __init__( @@ -229,7 +228,6 @@ class PhiMoE(nn.Module): class PhiMoEAttention(nn.Module): - def __init__( self, hidden_size: int, @@ -325,7 +323,6 @@ class PhiMoEAttention(nn.Module): class PhiMoEDecoderLayer(nn.Module): - def __init__( self, config: PhiMoEConfig, @@ -396,7 +393,6 @@ class PhiMoEDecoderLayer(nn.Module): class PhiMoEModel(nn.Module): - def __init__( self, config: PhiMoEConfig, @@ -449,7 +445,6 @@ class PhiMoEModel(nn.Module): class PhiMoEForCausalLM(nn.Module): - def __init__( self, config: PhiMoEConfig, diff --git a/python/sglang/srt/models/qwen.py b/python/sglang/srt/models/qwen.py index f6fd92128..a8f373507 100644 --- a/python/sglang/srt/models/qwen.py +++ b/python/sglang/srt/models/qwen.py @@ -73,8 +73,7 @@ class QWenMLP(nn.Module): ) if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() diff --git a/python/sglang/srt/models/qwen2.py b/python/sglang/srt/models/qwen2.py index 228ac550d..49eaecf8b 100644 --- a/python/sglang/srt/models/qwen2.py +++ b/python/sglang/srt/models/qwen2.py @@ -86,8 +86,7 @@ class Qwen2MLP(nn.Module): ) if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() @@ -465,7 +464,7 @@ class Qwen2Model(nn.Module): layer_self_attn.attn.v_scale = scaling_factor else: raise RuntimeError( - "Self attention has no KV cache scaling " "factor attribute!" + "Self attention has no KV cache scaling factor attribute!" ) diff --git a/python/sglang/srt/models/qwen2_5_vl.py b/python/sglang/srt/models/qwen2_5_vl.py index 3b58a0282..6b805a113 100644 --- a/python/sglang/srt/models/qwen2_5_vl.py +++ b/python/sglang/srt/models/qwen2_5_vl.py @@ -228,7 +228,6 @@ class Qwen2_5_VLMLP(nn.Module): class Qwen2_5_VisionBlock(nn.Module): - def __init__( self, dim: int, @@ -306,7 +305,6 @@ class Qwen2_5_VisionBlock(nn.Module): class Qwen2_5_VisionPatchMerger(nn.Module): - def __init__( self, dim: int, @@ -366,7 +364,6 @@ class Qwen2_5_VisionPatchMerger(nn.Module): class Qwen2_5_VisionTransformer(nn.Module, RotaryPosMixin): - def __init__( self, vision_config: Qwen2_5_VLVisionConfig, @@ -797,7 +794,6 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module): if current_dim == expected_dim: return pixel_values if current_dim != raw_patch_dim: - return pixel_values assert pixel_values.dim() == 2, pixel_values.dim() diff --git a/python/sglang/srt/models/qwen2_classification.py b/python/sglang/srt/models/qwen2_classification.py index a7f2b4fda..a1dc41e08 100644 --- a/python/sglang/srt/models/qwen2_classification.py +++ b/python/sglang/srt/models/qwen2_classification.py @@ -1,4 +1,4 @@ -# Copyright 2023-2024 SGLang Team +# Copyright 2023-2024 SGLang Team # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at @@ -57,9 +57,9 @@ class Qwen2ForSequenceClassification(nn.Module): input_embeds: torch.Tensor = None, get_embedding: bool = True, ) -> EmbeddingPoolerOutput: - assert ( - get_embedding - ), "Qwen2ForSequenceClassification is only used for embedding" + assert get_embedding, ( + "Qwen2ForSequenceClassification is only used for embedding" + ) hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) return score_and_pool( diff --git a/python/sglang/srt/models/qwen2_vl.py b/python/sglang/srt/models/qwen2_vl.py index 1e83b2364..828ca2aee 100644 --- a/python/sglang/srt/models/qwen2_vl.py +++ b/python/sglang/srt/models/qwen2_vl.py @@ -94,7 +94,6 @@ class Qwen2VLVideoInputs(TypedDict): class Qwen2VisionMLP(nn.Module): - def __init__( self, in_features: int, @@ -126,7 +125,6 @@ class Qwen2VisionMLP(nn.Module): class Qwen2VisionBlock(nn.Module): - def __init__( self, dim: int, @@ -183,7 +181,6 @@ class Qwen2VisionBlock(nn.Module): class Qwen2VisionPatchEmbed(nn.Module): - def __init__( self, patch_size: int = 14, @@ -209,7 +206,6 @@ class Qwen2VisionPatchEmbed(nn.Module): class Qwen2VisionPatchMerger(nn.Module): - def __init__( self, d_model: int, @@ -256,7 +252,6 @@ class Qwen2VisionPatchMerger(nn.Module): class Qwen2VisionRotaryEmbedding(nn.Module): - def __init__(self, dim: int, theta: float = 10000.0) -> None: super().__init__() self.dim = dim @@ -291,7 +286,6 @@ class Qwen2VisionRotaryEmbedding(nn.Module): class Qwen2VisionTransformer(nn.Module): - def __init__( self, vision_config: Qwen2VLVisionConfig, diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 2268b835d..924b70ce8 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -559,9 +559,9 @@ class Qwen3_5GatedDeltaNet(nn.Module): cpu_split_sizes.append( int(target_size_sim * split_sizes[i] / split_size_sum) ) - assert ( - sum(cpu_split_sizes) == target_size_sim - ), f"Padding the loaded weight failed due to sizes are not divisible cleanly from {cpu_split_sizes} to {target_size_sim}" + assert sum(cpu_split_sizes) == target_size_sim, ( + f"Padding the loaded weight failed due to sizes are not divisible cleanly from {cpu_split_sizes} to {target_size_sim}" + ) chunks = loaded_weight.split(cpu_split_sizes, dim=split_dim) else: chunks = loaded_weight.split(split_sizes, dim=split_dim) diff --git a/python/sglang/srt/models/qwen3_5_mtp.py b/python/sglang/srt/models/qwen3_5_mtp.py index 439c4f6ea..f116e12c9 100644 --- a/python/sglang/srt/models/qwen3_5_mtp.py +++ b/python/sglang/srt/models/qwen3_5_mtp.py @@ -83,7 +83,6 @@ def _mtp_quant_config(quant_config): class Qwen3_5ForCausalLMMTP(nn.Module): - @staticmethod def shared_experts_fusion_disable_reason(hf_config, quant_config): return Qwen3_5ForCausalLM.shared_experts_fusion_disable_reason( diff --git a/python/sglang/srt/models/qwen3_next.py b/python/sglang/srt/models/qwen3_next.py index 3aca4ec7c..43da47e00 100644 --- a/python/sglang/srt/models/qwen3_next.py +++ b/python/sglang/srt/models/qwen3_next.py @@ -505,7 +505,6 @@ def _apply_qwen3_next_mlp( class Qwen3HybridLinearDecoderLayer(nn.Module): - def __init__( self, config: Qwen3NextConfig, @@ -596,7 +595,6 @@ class Qwen3HybridLinearDecoderLayer(nn.Module): class Qwen3HybridAttentionDecoderLayer(nn.Module): - def __init__( self, config: Qwen3NextConfig, @@ -1145,9 +1143,7 @@ class Qwen3NextForCausalLM(nn.Module): params_dict = dict(self.named_parameters()) loaded_params: Set[str] = set() for name, loaded_weight in weights: - if is_mtp: - if "mtp" not in name: continue @@ -1240,9 +1236,9 @@ class Qwen3NextForCausalLM(nn.Module): # continue if name.endswith("_scale") and name not in params_dict: - assert ( - abs(loaded_weight.item() - 1.0) < 1e-6 - ), f"Expected 1.0, got {loaded_weight.item()} in skipped {name}" + assert abs(loaded_weight.item() - 1.0) < 1e-6, ( + f"Expected 1.0, got {loaded_weight.item()} in skipped {name}" + ) continue param = params_dict[name] weight_loader = getattr( diff --git a/python/sglang/srt/models/qwen3_next_mtp.py b/python/sglang/srt/models/qwen3_next_mtp.py index dd36afb51..f39b989ba 100644 --- a/python/sglang/srt/models/qwen3_next_mtp.py +++ b/python/sglang/srt/models/qwen3_next_mtp.py @@ -39,7 +39,6 @@ logger = logging.getLogger(__name__) class Qwen3NextForCausalLMMTP(Qwen3NextForCausalLM): - def __init__( self, config: PretrainedConfig, diff --git a/python/sglang/srt/models/qwen3_omni_moe.py b/python/sglang/srt/models/qwen3_omni_moe.py index 18ae8475e..1b72f9aac 100644 --- a/python/sglang/srt/models/qwen3_omni_moe.py +++ b/python/sglang/srt/models/qwen3_omni_moe.py @@ -376,7 +376,6 @@ class Qwen3OmniMoeAudioEncoder(PreTrainedModel): class Qwen3OmniMoeVisionPatchMerger(nn.Module): - def __init__( self, dim: int, diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py index 2b7f1cc06..ad7b5ce01 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/python/sglang/srt/models/qwen3_vl.py @@ -122,7 +122,6 @@ def _resolve_vision_tp( class Qwen3_VisionMLP(nn.Module): - def __init__( self, in_features: int, @@ -202,7 +201,6 @@ class Qwen3VLVisionPatchEmbed(nn.Module): class Qwen3_VisionBlock(nn.Module): - def __init__( self, dim: int, @@ -278,7 +276,6 @@ class Qwen3_VisionBlock(nn.Module): class Qwen3VLMoeVisionPatchMerger(nn.Module): - def __init__( self, dim: int, @@ -343,7 +340,6 @@ class Qwen3VLMoeVisionPatchMerger(nn.Module): class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin): - def __init__( self, vision_config: Qwen3VLVisionConfig, @@ -1098,7 +1094,9 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin): loaded_params.add(name) return loaded_params - def _prepare_graph_inputs(self, x: torch.Tensor, grid_thw: torch.Tensor) -> tuple[ + def _prepare_graph_inputs( + self, x: torch.Tensor, grid_thw: torch.Tensor + ) -> tuple[ torch.Tensor, torch.Tensor, torch.Tensor, @@ -1136,7 +1134,6 @@ cached_get_processor = lru_cache(get_processor) class Qwen3LLMModel(Qwen3Model): - def __init__( self, *, @@ -1148,7 +1145,9 @@ class Qwen3LLMModel(Qwen3Model): if not self.pp_group.is_first_rank: assert self.start_layer >= len( config.vision_config.deepstack_visual_indexes - ), "start_layer should be greater than or equal to len(deepstack_visual_indexes)" + ), ( + "start_layer should be greater than or equal to len(deepstack_visual_indexes)" + ) self.hidden_size = config.hidden_size self.deepstack_embed_to_decoder_layer = range( @@ -1376,9 +1375,9 @@ class Qwen3VLForConditionalGeneration(nn.Module): self.capture_aux_hidden_states = False def separate_deepstack_embeds(self, embedding): - assert ( - embedding.shape[-1] % (1 + self.num_deepstack_embeddings) == 0 - ), f"hidden_state of {embedding.shape} should be divisible by ({1 + self.num_deepstack_embeddings})" + assert embedding.shape[-1] % (1 + self.num_deepstack_embeddings) == 0, ( + f"hidden_state of {embedding.shape} should be divisible by ({1 + self.num_deepstack_embeddings})" + ) separate_index = self.config.hidden_size input_embeds = embedding[:, :separate_index] diff --git a/python/sglang/srt/models/radio.py b/python/sglang/srt/models/radio.py index df295f202..20edc607b 100644 --- a/python/sglang/srt/models/radio.py +++ b/python/sglang/srt/models/radio.py @@ -202,9 +202,9 @@ class ViTPatchGenerator(nn.Module): def forward_video(self, x: torch.Tensor, temporal_patch_size: int) -> torch.Tensor: """Embed video frames with temporal compression via tubelet grouping.""" - assert ( - self.video_embedder is not None - ), "video_embedder is required for temporal compression" + assert self.video_embedder is not None, ( + "video_embedder is required for temporal compression" + ) T = temporal_patch_size num_frames = x.shape[0] @@ -258,9 +258,9 @@ class ViTPatchGenerator(nn.Module): if src_embed.shape != targ_embed.shape: src_size = int(math.sqrt(src_embed.shape[1])) - assert ( - src_size**2 == src_embed.shape[1] - ), "Unable to interpolate non-square embedding" + assert src_size**2 == src_embed.shape[1], ( + "Unable to interpolate non-square embedding" + ) src_embed = rearrange( src_embed, "b (h w) c -> b c h w", h=src_size, w=src_size @@ -281,9 +281,9 @@ class ViTPatchGenerator(nn.Module): if src_proj_weight.shape != targ_proj_weight.shape: src_patch_size = int(math.sqrt(src_proj_weight.shape[1] // 3)) - assert (src_patch_size**2) * 3 == src_proj_weight.shape[ - 1 - ], "Unable to interpolate non-square patch size" + assert (src_patch_size**2) * 3 == src_proj_weight.shape[1], ( + "Unable to interpolate non-square patch size" + ) src_proj_weight = rearrange( src_proj_weight, diff --git a/python/sglang/srt/models/registry.py b/python/sglang/srt/models/registry.py index 71fdea29f..b83e5bb07 100644 --- a/python/sglang/srt/models/registry.py +++ b/python/sglang/srt/models/registry.py @@ -114,14 +114,14 @@ def import_model_classes(package_name: str, strict: bool = False): entry, list ): # To support multiple model classes in one module for tmp in entry: - assert ( - tmp.__name__ not in model_arch_name_to_cls - ), f"Duplicated model implementation for {tmp.__name__}" + assert tmp.__name__ not in model_arch_name_to_cls, ( + f"Duplicated model implementation for {tmp.__name__}" + ) model_arch_name_to_cls[tmp.__name__] = tmp else: - assert ( - entry.__name__ not in model_arch_name_to_cls - ), f"Duplicated model implementation for {entry.__name__}" + assert entry.__name__ not in model_arch_name_to_cls, ( + f"Duplicated model implementation for {entry.__name__}" + ) model_arch_name_to_cls[entry.__name__] = entry return model_arch_name_to_cls diff --git a/python/sglang/srt/models/roberta.py b/python/sglang/srt/models/roberta.py index c81590320..0fbf1c2b4 100644 --- a/python/sglang/srt/models/roberta.py +++ b/python/sglang/srt/models/roberta.py @@ -36,7 +36,6 @@ class RobertaClassificationHead(nn.Module): class RobertaEmbedding(nn.Module): - def __init__(self, config: RobertaConfig): super().__init__() self.size = config.hidden_size @@ -170,7 +169,6 @@ class XLMRobertaBaseModel(nn.Module): if self.pooler is None and "pooler" in name: continue for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: continue name = name.replace(weight_name, param_name) @@ -305,9 +303,9 @@ class XLMRobertaForSequenceClassification(nn.Module): input_embeds: torch.Tensor = None, get_embedding: bool = True, ) -> torch.Tensor: - assert ( - get_embedding - ), "XLMRobertaForSequenceClassification is only used for rerank" + assert get_embedding, ( + "XLMRobertaForSequenceClassification is only used for rerank" + ) hidden_states = self.roberta( input_ids, positions, forward_batch, input_embeds, get_embedding diff --git a/python/sglang/srt/models/siglip.py b/python/sglang/srt/models/siglip.py index 6f8dec4d6..52d072cee 100644 --- a/python/sglang/srt/models/siglip.py +++ b/python/sglang/srt/models/siglip.py @@ -20,7 +20,6 @@ from sglang.srt.utils import add_prefix # Adapted from transformers.models.siglip.modeling_siglip.SiglipVisionTransformer class SiglipVisionEmbeddings(nn.Module): - def __init__(self, config: SiglipVisionConfig, use_data_parallel: bool = False): super().__init__() self.config = config @@ -63,7 +62,6 @@ class SiglipVisionEmbeddings(nn.Module): # Copied from sglang.srt.models.clip.CLIPMLP class SiglipMLP(nn.Module): - def __init__( self, config, @@ -102,7 +100,6 @@ class SiglipMLP(nn.Module): # Copied from sglang.srt.models.clip.CLIPEncoderLayer class SiglipEncoderLayer(nn.Module): - def __init__( self, config: SiglipVisionConfig, @@ -233,7 +230,6 @@ class SiglipEncoder(nn.Module): # Adapted from transformers.models.siglip.modeling_siglip.SiglipVisionTransformer class SiglipVisionTransformer(nn.Module): - def __init__( self, config: SiglipVisionConfig, diff --git a/python/sglang/srt/models/solar.py b/python/sglang/srt/models/solar.py index d66f22c7a..bf54128f5 100644 --- a/python/sglang/srt/models/solar.py +++ b/python/sglang/srt/models/solar.py @@ -60,7 +60,6 @@ from sglang.srt.utils.hf_transformers_utils import get_rope_config class SolarMLP(nn.Module): - def __init__( self, hidden_size: int, @@ -87,8 +86,7 @@ class SolarMLP(nn.Module): ) if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() @@ -100,7 +98,6 @@ class SolarMLP(nn.Module): class SolarAttention(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -186,7 +183,6 @@ class SolarAttention(nn.Module): class SolarDecoderLayer(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -263,7 +259,6 @@ class SolarDecoderLayer(nn.Module): class SolarModel(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -380,12 +375,11 @@ class SolarModel(nn.Module): layer_self_attn.attn.v_scale = scaling_factor else: raise RuntimeError( - "Self attention has no KV cache scaling " "factor attribute!" + "Self attention has no KV cache scaling factor attribute!" ) class SolarForCausalLM(nn.Module): - packed_modules_mapping = { "qkv_proj": [ ("q_proj", "q"), @@ -475,12 +469,10 @@ class SolarForCausalLM(nn.Module): params_dict = dict(self.named_parameters()) for name, loaded_weight in weights: - is_packed = False for packed_name, sources in self.packed_modules_mapping.items(): for src_name, shard_id in sources: if src_name in name: - model_param_name = name.replace(src_name, packed_name) if model_param_name in params_dict: diff --git a/python/sglang/srt/models/starcoder2.py b/python/sglang/srt/models/starcoder2.py index 22250f63f..5b33f704a 100644 --- a/python/sglang/srt/models/starcoder2.py +++ b/python/sglang/srt/models/starcoder2.py @@ -52,7 +52,6 @@ from sglang.srt.utils import add_prefix, make_layers class Starcoder2Attention(nn.Module): - def __init__( self, config: Starcoder2Config, @@ -134,7 +133,6 @@ class Starcoder2Attention(nn.Module): class Starcoder2MLP(nn.Module): - def __init__( self, config: Starcoder2Config, @@ -169,7 +167,6 @@ class Starcoder2MLP(nn.Module): class Starcoder2DecoderLayer(nn.Module): - def __init__( self, config: Starcoder2Config, @@ -219,7 +216,6 @@ class Starcoder2DecoderLayer(nn.Module): class Starcoder2Model(nn.Module): - def __init__( self, config: Starcoder2Config, @@ -276,7 +272,6 @@ class Starcoder2Model(nn.Module): class Starcoder2ForCausalLM(nn.Module): - def __init__( self, config: Starcoder2Config, diff --git a/python/sglang/srt/models/step3_vl.py b/python/sglang/srt/models/step3_vl.py index f659e926e..755156bbb 100644 --- a/python/sglang/srt/models/step3_vl.py +++ b/python/sglang/srt/models/step3_vl.py @@ -94,8 +94,7 @@ class Step3TextMLP(nn.Module): ) if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() @@ -604,7 +603,6 @@ class Step3VisionAttention(nn.Module): class Step3VisionEmbeddings(nn.Module): - def __init__(self, config: Step3VisionEncoderConfig): super().__init__() self.config = config @@ -731,7 +729,6 @@ class Step3VisionEncoder(nn.Module): class Step3VLForConditionalGeneration(nn.Module): - def __init__( self, config: Step3VLConfig, diff --git a/python/sglang/srt/models/step3p5_mtp.py b/python/sglang/srt/models/step3p5_mtp.py index 070bbe1f6..d3d059e1e 100644 --- a/python/sglang/srt/models/step3p5_mtp.py +++ b/python/sglang/srt/models/step3p5_mtp.py @@ -41,7 +41,6 @@ def get_spec_layer_idx_from_weight_name( class SharedHead(nn.Module): - def __init__( self, config, diff --git a/python/sglang/srt/models/step3p7.py b/python/sglang/srt/models/step3p7.py index 5dd0816f7..82159ec4c 100644 --- a/python/sglang/srt/models/step3p7.py +++ b/python/sglang/srt/models/step3p7.py @@ -25,7 +25,6 @@ from sglang.srt.utils import add_prefix class Step3p7ForConditionalGeneration(nn.Module): - # NVFP4 checkpoints (e.g. huangyu-nv/step3p7-nvfp4-moe-only-kvfp8) use # "model.language_model." prefix, while sglang parameters are named # "language_model.model.". This mapper remaps the quantization ignore diff --git a/python/sglang/srt/models/torch_native_llama.py b/python/sglang/srt/models/torch_native_llama.py index 68bdbbf56..1ee258199 100644 --- a/python/sglang/srt/models/torch_native_llama.py +++ b/python/sglang/srt/models/torch_native_llama.py @@ -126,8 +126,7 @@ class LlamaMLP(nn.Module): self.down_proj = torch.nn.Linear(intermediate_size, hidden_size, bias=False) if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() diff --git a/python/sglang/srt/models/transformers.py b/python/sglang/srt/models/transformers.py index 343f601d9..af5e65b48 100644 --- a/python/sglang/srt/models/transformers.py +++ b/python/sglang/srt/models/transformers.py @@ -1077,9 +1077,9 @@ class TransformersBase(nn.Module): ) if get_embedding: - assert ( - self.pooler is not None - ), "pooling is not enabled for this model class" + assert self.pooler is not None, ( + "pooling is not enabled for this model class" + ) return self.pooler(hidden_states, forward_batch) assert self.logits_processor is not None and self.lm_head is not None @@ -1100,7 +1100,6 @@ class TransformersBase(nn.Module): class CausalMixin: - def __init__(self, *args, prefix: str = "", **kwargs): super().__init__(*args, prefix=prefix, **kwargs) @@ -1128,7 +1127,6 @@ class CausalMixin: class EmbeddingMixin: - def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.ignore_unexpected_prefixes.append("lm_head.") @@ -1141,7 +1139,6 @@ class EmbeddingMixin: class MoEMixin: - def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) diff --git a/python/sglang/srt/models/unlimited_ocr.py b/python/sglang/srt/models/unlimited_ocr.py index 8c0cafd44..63b09489f 100644 --- a/python/sglang/srt/models/unlimited_ocr.py +++ b/python/sglang/srt/models/unlimited_ocr.py @@ -181,7 +181,7 @@ class UnlimitedOCRForCausalLM(nn.Module): if pixel_values is not None: if not isinstance(pixel_values, (torch.Tensor, list)): raise ValueError( - "Incorrect type of pixel values. " f"Got type: {type(pixel_values)}" + f"Incorrect type of pixel values. Got type: {type(pixel_values)}" ) if not isinstance(images_spatial_crop, (torch.Tensor, list)): raise ValueError( @@ -190,7 +190,7 @@ class UnlimitedOCRForCausalLM(nn.Module): ) if not isinstance(images_crop, (torch.Tensor, list)): raise ValueError( - "Incorrect type of image crop. " f"Got type: {type(images_crop)}" + f"Incorrect type of image crop. Got type: {type(images_crop)}" ) return [pixel_values, images_crop, images_spatial_crop] diff --git a/python/sglang/srt/models/utils.py b/python/sglang/srt/models/utils.py index b9b18a6dd..27c7addd1 100644 --- a/python/sglang/srt/models/utils.py +++ b/python/sglang/srt/models/utils.py @@ -382,9 +382,9 @@ def compute_cu_seqlens_from_grid_numpy(grid_thw: torch.Tensor) -> torch.Tensor: Returns: cu_seqlens: 1D int32 tensor on CPU, shape [N + 1] """ - assert ( - grid_thw.device.type == "cpu" - ), "compute_cu_seqlens_from_grid_numpy expects a CPU tensor" + assert grid_thw.device.type == "cpu", ( + "compute_cu_seqlens_from_grid_numpy expects a CPU tensor" + ) arr = grid_thw.numpy() cu_seqlens = np.repeat(arr[:, 1] * arr[:, 2], arr[:, 0]).cumsum( diff --git a/python/sglang/srt/models/whisper.py b/python/sglang/srt/models/whisper.py index d5135406c..d8299f272 100644 --- a/python/sglang/srt/models/whisper.py +++ b/python/sglang/srt/models/whisper.py @@ -42,9 +42,9 @@ class WhisperAttention(torch.nn.Module): self.is_encoder = is_encoder tp_size = get_parallel().tp_size - assert ( - num_heads % tp_size == 0 - ), f"num_heads ({num_heads}) must be divisible by tp_size ({tp_size})" + assert num_heads % tp_size == 0, ( + f"num_heads ({num_heads}) must be divisible by tp_size ({tp_size})" + ) self.num_heads = num_heads // tp_size if (head_dim * num_heads) != embed_dim: @@ -262,7 +262,6 @@ class WhisperDecoderLayer(torch.nn.Module): class WhisperEncoder(torch.nn.Module): - def __init__( self, config: WhisperConfig, quant_config: Optional[QuantizationConfig] = None ): @@ -314,7 +313,6 @@ class WhisperEncoder(torch.nn.Module): class WhisperDecoder(torch.nn.Module): - def __init__( self, config: WhisperConfig, quant_config: Optional[QuantizationConfig] = None ): @@ -362,7 +360,6 @@ class WhisperDecoder(torch.nn.Module): class WhisperForConditionalGeneration(torch.nn.Module): - def __init__( self, config: WhisperConfig, quant_config: Optional[QuantizationConfig] = None ): diff --git a/python/sglang/srt/models/xverse.py b/python/sglang/srt/models/xverse.py index d4366c494..6fd1222ad 100644 --- a/python/sglang/srt/models/xverse.py +++ b/python/sglang/srt/models/xverse.py @@ -72,8 +72,7 @@ class XverseMLP(nn.Module): ) if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() diff --git a/python/sglang/srt/models/xverse_moe.py b/python/sglang/srt/models/xverse_moe.py index 964a6ba4c..898cf3abc 100644 --- a/python/sglang/srt/models/xverse_moe.py +++ b/python/sglang/srt/models/xverse_moe.py @@ -52,7 +52,6 @@ from sglang.srt.utils.hf_transformers_utils import get_rope_config class XverseMLP(nn.Module): - def __init__( self, hidden_size: int, @@ -80,8 +79,7 @@ class XverseMLP(nn.Module): ) if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() @@ -93,7 +91,6 @@ class XverseMLP(nn.Module): class XverseMoE(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -195,7 +192,6 @@ class XverseMoE(nn.Module): class XverseAttention(nn.Module): - def __init__( self, hidden_size: int, @@ -281,7 +277,6 @@ class XverseAttention(nn.Module): class XverseDecoderLayer(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -352,7 +347,6 @@ class XverseDecoderLayer(nn.Module): class XverseModel(nn.Module): - fall_back_to_pt_during_load = False def __init__( @@ -401,7 +395,6 @@ class XverseModel(nn.Module): class XverseMoeForCausalLM(nn.Module): - def __init__( self, config: PretrainedConfig, diff --git a/python/sglang/srt/models/zaya.py b/python/sglang/srt/models/zaya.py index 1f63f2019..8f9b19ba1 100644 --- a/python/sglang/srt/models/zaya.py +++ b/python/sglang/srt/models/zaya.py @@ -386,9 +386,9 @@ class CCA(nn.Module): # Full (global) head counts retained for weight loading and shape asserts. self.num_q_heads_full = int(cca_num_q_heads) self.num_k_heads_full = int(cca_num_k_heads) - assert ( - self.num_q_heads_full % self.num_k_heads_full == 0 - ), "num_q_heads must be a multiple of num_k_heads" + assert self.num_q_heads_full % self.num_k_heads_full == 0, ( + "num_q_heads must be a multiple of num_k_heads" + ) self.gqa_groups = self.num_q_heads_full // self.num_k_heads_full # Head-parallel TP requires both head counts to be divisible by tp_size. @@ -558,16 +558,16 @@ class CCA(nn.Module): [loaded_weight[q_start:q_end], loaded_weight[k_start:k_end]], dim=0, ) - assert ( - sliced.shape == param.data.shape - ), f"conv shard shape mismatch: {sliced.shape} vs {param.data.shape}" + assert sliced.shape == param.data.shape, ( + f"conv shard shape mismatch: {sliced.shape} vs {param.data.shape}" + ) param.data.copy_(sliced) def temp_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: sliced = loaded_weight[k_temp_start:k_temp_end] - assert ( - sliced.shape == param.data.shape - ), f"temp shard shape mismatch: {sliced.shape} vs {param.data.shape}" + assert sliced.shape == param.data.shape, ( + f"temp shard shape mismatch: {sliced.shape} vs {param.data.shape}" + ) param.data.copy_(sliced) set_weight_attrs(self.conv_qk[0].weight, {"weight_loader": conv_row_loader}) @@ -1147,9 +1147,9 @@ class ZayaBlock(nn.Module): f"number of experts {self.num_moe_experts}" ) - assert ( - config.activation_func == "swiglu" - ), "ZayaBlock only supports SwiGLU activation" + assert config.activation_func == "swiglu", ( + "ZayaBlock only supports SwiGLU activation" + ) assert config.gated_linear_unit, "ZayaBlock requires gated_linear_unit=True" self.router = ZayaRouter( diff --git a/python/sglang/srt/multimodal/cache/identity.py b/python/sglang/srt/multimodal/cache/identity.py index 6aee84559..66568e7cf 100644 --- a/python/sglang/srt/multimodal/cache/identity.py +++ b/python/sglang/srt/multimodal/cache/identity.py @@ -298,8 +298,7 @@ def _canonicalize(value: Any) -> Any: "items": items, } raise ValueError( - "Unsupported value in multimodal cache identity: " - f"{_qualified_type_name(value)}" + f"Unsupported value in multimodal cache identity: {_qualified_type_name(value)}" ) diff --git a/python/sglang/srt/multimodal/evs/evs_module.py b/python/sglang/srt/multimodal/evs/evs_module.py index b62ce08de..f337175d3 100644 --- a/python/sglang/srt/multimodal/evs/evs_module.py +++ b/python/sglang/srt/multimodal/evs/evs_module.py @@ -72,9 +72,9 @@ class EVSEmbeddingResult(EmbeddingResult): input_ids, filler_token_id ) input_ids = input_ids[extend_prefix_len : extend_prefix_len + extend_seq_len] - assert ( - len(input_ids) == extend_seq_len - ), f"Input ids length changed after redistribution, got {len(input_ids)} != {extend_seq_len}" + assert len(input_ids) == extend_seq_len, ( + f"Input ids length changed after redistribution, got {len(input_ids)} != {extend_seq_len}" + ) return input_ids, offsets @@ -84,9 +84,9 @@ class EVSConfig: spatial_merge_size: int = 1 def __post_init__(self): - assert ( - self.video_pruning_rate >= 0.0 and self.video_pruning_rate < 1.0 - ), f"Video pruning rate must be between 0.0 and 1.0, got {self.video_pruning_rate=}" + assert self.video_pruning_rate >= 0.0 and self.video_pruning_rate < 1.0, ( + f"Video pruning rate must be between 0.0 and 1.0, got {self.video_pruning_rate=}" + ) class EVS(torch.nn.Module, ABC): diff --git a/python/sglang/srt/multimodal/evs/evs_processor.py b/python/sglang/srt/multimodal/evs/evs_processor.py index 98e731afa..100fad141 100644 --- a/python/sglang/srt/multimodal/evs/evs_processor.py +++ b/python/sglang/srt/multimodal/evs/evs_processor.py @@ -72,7 +72,7 @@ class EVSProcessor: return evs_config = evs_model.create_evs_config(hf_config) logger.info( - f"""[EVS] {evs_config} {'enabled' if evs_config.video_pruning_rate > 0.0 else 'disabled'} for model={evs_model.__name__}; model_config={config_name}""" + f"""[EVS] {evs_config} {"enabled" if evs_config.video_pruning_rate > 0.0 else "disabled"} for model={evs_model.__name__}; model_config={config_name}""" ) if evs_config.video_pruning_rate > 0.0: self.evs_config = evs_config diff --git a/python/sglang/srt/multimodal/media_artifacts/base.py b/python/sglang/srt/multimodal/media_artifacts/base.py index f93890f18..762c6343f 100644 --- a/python/sglang/srt/multimodal/media_artifacts/base.py +++ b/python/sglang/srt/multimodal/media_artifacts/base.py @@ -147,7 +147,7 @@ class MediaArtifactCacheMixin: """ if modality != Modality.IMAGE: raise NotImplementedError( - f"{modality.name.lower()} artifact decoding " "requires a model adapter" + f"{modality.name.lower()} artifact decoding requires a model adapter" ) data = snapshot.data if isinstance(data, torch.Tensor): @@ -171,7 +171,7 @@ class MediaArtifactCacheMixin: """ if modality != Modality.IMAGE: raise NotImplementedError( - f"{modality.name.lower()} artifact identity " "requires a model adapter" + f"{modality.name.lower()} artifact identity requires a model adapter" ) return snapshot_media(source) diff --git a/python/sglang/srt/multimodal/mm_utils.py b/python/sglang/srt/multimodal/mm_utils.py index d7548ba20..97cab3284 100644 --- a/python/sglang/srt/multimodal/mm_utils.py +++ b/python/sglang/srt/multimodal/mm_utils.py @@ -130,8 +130,9 @@ def select_best_resolution(original_size, possible_resolutions): for width, height in possible_resolutions: # Calculate the downscaled size to keep the aspect ratio scale = min(width / original_width, height / original_height) - downscaled_width, downscaled_height = int(original_width * scale), int( - original_height * scale + downscaled_width, downscaled_height = ( + int(original_width * scale), + int(original_height * scale), ) # Calculate effective and wasted resolutions diff --git a/python/sglang/srt/multimodal/processors/base_processor.py b/python/sglang/srt/multimodal/processors/base_processor.py index 30440de8b..6c2367b86 100644 --- a/python/sglang/srt/multimodal/processors/base_processor.py +++ b/python/sglang/srt/multimodal/processors/base_processor.py @@ -1754,9 +1754,9 @@ class BaseMultimodalProcessor(ABC): and not raw_audios and not raw_videos ): - assert isinstance( - base_output.input_ids, list - ), f"expected list[int] input_ids, got {type(base_output.input_ids)}" + assert isinstance(base_output.input_ids, list), ( + f"expected list[int] input_ids, got {type(base_output.input_ids)}" + ) try: counts = self.resolve_image_token_counts(raw_images) image_placeholder_token_id = mm_tokens.image_token_id diff --git a/python/sglang/srt/multimodal/processors/ernie45_vl.py b/python/sglang/srt/multimodal/processors/ernie45_vl.py index ca76deba7..37a629c4d 100644 --- a/python/sglang/srt/multimodal/processors/ernie45_vl.py +++ b/python/sglang/srt/multimodal/processors/ernie45_vl.py @@ -145,9 +145,9 @@ def smart_nframes( Returns: int: the number of frames for video used for model inputs. """ - assert not ( - "fps" in ele and "nframes" in ele - ), "Only accept either `fps` or `nframes`" + assert not ("fps" in ele and "nframes" in ele), ( + "Only accept either `fps` or `nframes`" + ) if "nframes" in ele: nframes = round_by_factor(ele["nframes"], FRAME_FACTOR) else: @@ -426,9 +426,9 @@ class Ernie4_5_VLImageProcessor(SGLangBaseProcessor): ) mrope_positions = mrope_positions.squeeze(1) - assert ( - input_ids.shape[0] == mrope_positions.shape[-1] - ), "input_ids and mrope_positions should have the same length" + assert input_ids.shape[0] == mrope_positions.shape[-1], ( + "input_ids and mrope_positions should have the same length" + ) return MultimodalProcessorOutput( input_ids=input_ids.tolist(), diff --git a/python/sglang/srt/multimodal/processors/internvl.py b/python/sglang/srt/multimodal/processors/internvl.py index 05c3e2206..d2f9a85b8 100644 --- a/python/sglang/srt/multimodal/processors/internvl.py +++ b/python/sglang/srt/multimodal/processors/internvl.py @@ -562,7 +562,7 @@ class InternVLProcessor(BaseMultimodalProcessor): + (self.VIDEO_CONTEXT_TOKEN * ctx_cnt) + self.IMG_END ) - frame_lines.append(f"Frame {i+1}: {frame_tokens}") + frame_lines.append(f"Frame {i + 1}: {frame_tokens}") video_tokens = "\n".join(frame_lines) + "\n" input_text_updated = input_text_updated.replace(vid_ph, video_tokens, 1) diff --git a/python/sglang/srt/multimodal/processors/mimo_audio.py b/python/sglang/srt/multimodal/processors/mimo_audio.py index a0f727e49..909c9965b 100644 --- a/python/sglang/srt/multimodal/processors/mimo_audio.py +++ b/python/sglang/srt/multimodal/processors/mimo_audio.py @@ -169,9 +169,9 @@ class MiMoAudioPipeline: Output: (mel-spectrogram tensor [T, n_mels], audio_token_len int). """ self._ensure_audio_dependencies() - assert isinstance( - audio, (str, bytes, tuple) - ), f"audio must be a str, bytes or tuple, but got {type(audio)}" + assert isinstance(audio, (str, bytes, tuple)), ( + f"audio must be a str, bytes or tuple, but got {type(audio)}" + ) if isinstance(audio, tuple): waveform, original_sr = audio else: @@ -256,9 +256,9 @@ class MiMoAudioPipeline: if isinstance(audio, (str, bytes, tuple)): return self.preprocess_audio(audio) - assert ( - audio.shape[1] >= self.audio_channels - ), f"audio must have at least {self.audio_channels} channels, but got {audio.shape[1]}" + assert audio.shape[1] >= self.audio_channels, ( + f"audio must have at least {self.audio_channels} channels, but got {audio.shape[1]}" + ) T = audio.shape[0] audio = audio[:, : self.audio_channels].to(torch.long) padded_T = ( diff --git a/python/sglang/srt/multimodal/processors/mimo_v2.py b/python/sglang/srt/multimodal/processors/mimo_v2.py index 55e761fe8..c833d7937 100644 --- a/python/sglang/srt/multimodal/processors/mimo_v2.py +++ b/python/sglang/srt/multimodal/processors/mimo_v2.py @@ -330,14 +330,14 @@ class MiMoProcessor: self.use_video_timestamps = use_video_timestamps assert self.use_video_timestamps - assert ( - not self.use_video_timestamps or self.rope_type == "rope" - ), "use_video_timestamps only supports 1d rope" + assert not self.use_video_timestamps or self.rope_type == "rope", ( + "use_video_timestamps only supports 1d rope" + ) self.video_audio_interleave_length = video_audio_interleave_length self.use_per_grid_t_timestamps = False - assert ( - self.video_audio_interleave_length == -1 or self.rope_type == "rope" - ), "video_audio_interleave_length != -1 only supports 1d rope" + assert self.video_audio_interleave_length == -1 or self.rope_type == "rope", ( + "video_audio_interleave_length != -1 only supports 1d rope" + ) assert ( self.video_audio_interleave_length == -1 or self.video_audio_interleave_length >= 0 @@ -723,9 +723,9 @@ class MiMoProcessor: else: selected_frame_indices = candidate_indices - assert ( - len(selected_frame_indices) > 0 - ), f"No frames selected for segment {start_time} - {end_time} in all_timestamps {all_timestamps.tolist()}" + assert len(selected_frame_indices) > 0, ( + f"No frames selected for segment {start_time} - {end_time} in all_timestamps {all_timestamps.tolist()}" + ) return selected_frame_indices kwargs = self.prepare_video_kwargs(video_input) @@ -791,9 +791,9 @@ class MiMoProcessor: min_pixels, max_pixels = smart_resize_video(num_frames_sampled, **kwargs) - assert ( - num_frames_seg > 0 - ), f"Sampled frame number must be >0. start_time {video_input.start_time}, end_time {video_input.end_time}, start_time_seg {start_time_seg}, end_time_seg {end_time_seg}. Full timestamps {timestamps_sampled.tolist()}. " + assert num_frames_seg > 0, ( + f"Sampled frame number must be >0. start_time {video_input.start_time}, end_time {video_input.end_time}, start_time_seg {start_time_seg}, end_time_seg {end_time_seg}. Full timestamps {timestamps_sampled.tolist()}. " + ) temporal_padding_factor = ( self.temporal_patch_size * self.temporal_compression_ratio @@ -911,9 +911,9 @@ class MiMoProcessor: // self.temporal_compression_ratio ) - assert ( - len(timestamps) == grid_t * self.temporal_patch_size - ), f"Expected {grid_t} * {self.temporal_patch_size} = {grid_t * self.temporal_patch_size} timestamps, but got {len(timestamps)}" + assert len(timestamps) == grid_t * self.temporal_patch_size, ( + f"Expected {grid_t} * {self.temporal_patch_size} = {grid_t * self.temporal_patch_size} timestamps, but got {len(timestamps)}" + ) if not self.use_video_timestamps: raise NotImplementedError @@ -945,7 +945,7 @@ class MiMoProcessor: if verbose: verbose_str = f"Video (video_thw_grid={thw_grid}, video_meta={video_meta}): [ " for i, ts in enumerate(text_timestamps): - verbose_str += f"{ts} {timestamps.tolist()[i*self.temporal_patch_size*self.temporal_compression_ratio : (i+1)*self.temporal_patch_size*self.temporal_compression_ratio]} {num_media_tokens_per_grid}* " + verbose_str += f"{ts} {timestamps.tolist()[i * self.temporal_patch_size * self.temporal_compression_ratio : (i + 1) * self.temporal_patch_size * self.temporal_compression_ratio]} {num_media_tokens_per_grid}* " verbose_str += "]\n" return { @@ -983,9 +983,9 @@ class MiMoProcessor: # Compute per-grid_t audio-segment boundaries. Tokenizer-free so it # runs identically on the single-node path and the EPD encoder side. grid_t, grid_h, grid_w = thw_grid - assert ( - len(timestamps) == grid_t * self.temporal_patch_size - ), f"Expected {grid_t} * {self.temporal_patch_size} timestamps, got {len(timestamps)}" + assert len(timestamps) == grid_t * self.temporal_patch_size, ( + f"Expected {grid_t} * {self.temporal_patch_size} timestamps, got {len(timestamps)}" + ) if not self.use_video_timestamps: raise NotImplementedError @@ -1239,9 +1239,9 @@ class MiMoProcessor: labels = torch.tensor(labels) if len(is_audio_tokenized) > 0: - assert all(is_audio_tokenized) or not any( - is_audio_tokenized - ), "All audio inputs must be tokenized or not tokenized" + assert all(is_audio_tokenized) or not any(is_audio_tokenized), ( + "All audio inputs must be tokenized or not tokenized" + ) extra["is_audio_tokenized"] = is_audio_tokenized[0] if self.rope_type == "rope": diff --git a/python/sglang/srt/multimodal/processors/mimo_v2_asr.py b/python/sglang/srt/multimodal/processors/mimo_v2_asr.py index 3ebba3d69..812a2fbc2 100644 --- a/python/sglang/srt/multimodal/processors/mimo_v2_asr.py +++ b/python/sglang/srt/multimodal/processors/mimo_v2_asr.py @@ -250,11 +250,14 @@ class MiMoV2ASRProcessor(BaseMultimodalProcessor): loop = asyncio.get_running_loop() try: - input_ids, audio_inputs, position_ids, rope_deltas = ( - await loop.run_in_executor( - self.io_executor, - lambda: self._process_contents(contents), - ) + ( + input_ids, + audio_inputs, + position_ids, + rope_deltas, + ) = await loop.run_in_executor( + self.io_executor, + lambda: self._process_contents(contents), ) except RuntimeError as e: logger.error(f"MiMo ASR processor failed in process_mm_data_async: {e}") diff --git a/python/sglang/srt/multimodal/processors/minicpm.py b/python/sglang/srt/multimodal/processors/minicpm.py index 021ae7e4e..e26a4eb8a 100644 --- a/python/sglang/srt/multimodal/processors/minicpm.py +++ b/python/sglang/srt/multimodal/processors/minicpm.py @@ -210,12 +210,12 @@ class MiniCPMMultimodalProcessor(BaseMultimodalProcessor): if not isinstance(pixel_values, (torch.Tensor, list)): raise ValueError( - "Incorrect type of pixel values. " f"Got type: {type(pixel_values)}" + f"Incorrect type of pixel values. Got type: {type(pixel_values)}" ) if not isinstance(tgt_sizes, (torch.Tensor, list)): raise ValueError( - "Incorrect type of target sizes. " f"Got type: {type(tgt_sizes)}" + f"Incorrect type of target sizes. Got type: {type(tgt_sizes)}" ) if len(pixel_values) != len(tgt_sizes): @@ -232,7 +232,7 @@ class MiniCPMMultimodalProcessor(BaseMultimodalProcessor): # per image if len(pixel_b) != len(tgt_b): raise ValueError( - "Inconsistent N lengths, found: " f"{len(pixel_b)} vs {len(tgt_b)}" + f"Inconsistent N lengths, found: {len(pixel_b)} vs {len(tgt_b)}" ) slices_per_image.append(len(pixel_b)) for pixel_n, tgt_n in zip(pixel_b, tgt_b): diff --git a/python/sglang/srt/multimodal/processors/moss_vl.py b/python/sglang/srt/multimodal/processors/moss_vl.py index 50ea9eaab..3f862f355 100644 --- a/python/sglang/srt/multimodal/processors/moss_vl.py +++ b/python/sglang/srt/multimodal/processors/moss_vl.py @@ -537,9 +537,10 @@ class MossVLImageProcessor(SGLangBaseProcessor): *args, **kwargs, ): - normalized_video_data, temp_video_paths = ( - await self._normalize_video_inputs_async(request_obj.video_data) - ) + ( + normalized_video_data, + temp_video_paths, + ) = await self._normalize_video_inputs_async(request_obj.video_data) try: base_output = await self.load_mm_data( diff --git a/python/sglang/srt/multimodal/processors/qwen3_asr.py b/python/sglang/srt/multimodal/processors/qwen3_asr.py index c77ff5462..5635e599b 100644 --- a/python/sglang/srt/multimodal/processors/qwen3_asr.py +++ b/python/sglang/srt/multimodal/processors/qwen3_asr.py @@ -13,10 +13,7 @@ from sglang.srt.multimodal.processors.base_processor import ( AUDIO_PLACEHOLDER = "<|audio_start|><|audio_pad|><|audio_end|>" DEFAULT_ASR_PROMPT = ( - f"<|im_start|>user\n" - f"{AUDIO_PLACEHOLDER}" - f"<|im_end|>\n" - f"<|im_start|>assistant\n" + f"<|im_start|>user\n{AUDIO_PLACEHOLDER}<|im_end|>\n<|im_start|>assistant\n" ) diff --git a/python/sglang/srt/multimodal/processors/qwen_audio.py b/python/sglang/srt/multimodal/processors/qwen_audio.py index 664bb3e4a..8fb15665a 100644 --- a/python/sglang/srt/multimodal/processors/qwen_audio.py +++ b/python/sglang/srt/multimodal/processors/qwen_audio.py @@ -171,9 +171,9 @@ class Qwen2AudioMultimodalProcessor(BaseMultimodalProcessor): base_output, self.mm_tokens ) - assert ( - "feature_attention_mask" in ret - ), "feature_attention_mask not found in processor output" + assert "feature_attention_mask" in ret, ( + "feature_attention_mask not found in processor output" + ) input_lengths = ret["feature_attention_mask"].sum(dim=-1) input_lengths = (input_lengths - 1) // 2 + 1 output_lengths = (input_lengths - 2) // 2 + 1 diff --git a/python/sglang/srt/multimodal/processors/qwen_vl.py b/python/sglang/srt/multimodal/processors/qwen_vl.py index cac51f22b..875d42de6 100644 --- a/python/sglang/srt/multimodal/processors/qwen_vl.py +++ b/python/sglang/srt/multimodal/processors/qwen_vl.py @@ -98,9 +98,7 @@ if _is_cpu and _is_cpu_amx_available: from sglang.srt.layers.amx_utils import fast_preprocess_cpu - transformers.models.qwen2_vl.image_processing_qwen2_vl_fast.Qwen2VLImageProcessorFast._preprocess = ( - fast_preprocess_cpu - ) + transformers.models.qwen2_vl.image_processing_qwen2_vl_fast.Qwen2VLImageProcessorFast._preprocess = fast_preprocess_cpu except Exception as e: logger.warning( f"Failed to hack Qwen2VLImageProcessorFast with AMX optimization: {e}" @@ -178,9 +176,9 @@ def smart_nframes( Returns: int: the number of frames for video used for model inputs. """ - assert not ( - "fps" in ele and "nframes" in ele - ), "Only accept either `fps` or `nframes`" + assert not ("fps" in ele and "nframes" in ele), ( + "Only accept either `fps` or `nframes`" + ) if "nframes" in ele: nframes = round_by_factor(ele["nframes"], FRAME_FACTOR) else: diff --git a/python/sglang/srt/multimodal/processors/step3_vl.py b/python/sglang/srt/multimodal/processors/step3_vl.py index fe73da628..b6554c876 100644 --- a/python/sglang/srt/multimodal/processors/step3_vl.py +++ b/python/sglang/srt/multimodal/processors/step3_vl.py @@ -27,7 +27,6 @@ ImageWithPatches = tuple[Step3Image, list[Step3Image], list[int] | None] class GPUToTensor(torch.nn.Module): - def forward( self, raw_image: Union[np.ndarray, Image.Image, torch.Tensor] ) -> torch.Tensor: diff --git a/python/sglang/srt/multimodal/processors/unlimited_ocr.py b/python/sglang/srt/multimodal/processors/unlimited_ocr.py index f5c52e728..2ca9f117b 100644 --- a/python/sglang/srt/multimodal/processors/unlimited_ocr.py +++ b/python/sglang/srt/multimodal/processors/unlimited_ocr.py @@ -37,8 +37,7 @@ def _resolve_mode(images_config, num_images: int = 1) -> dict: f"Unknown image_mode '{mode}'. Supported: {', '.join(_IMAGE_MODE_PRESETS)}" ) raise ValueError( - f"Unknown image_mode '{mode}'. " - f"Supported: {', '.join(_IMAGE_MODE_PRESETS)}" + f"Unknown image_mode '{mode}'. Supported: {', '.join(_IMAGE_MODE_PRESETS)}" ) _MULTI_IMAGE_ALLOWED = ("tiny", "small", "base") base_size, image_size, crop_mode = preset diff --git a/python/sglang/srt/multimodal/transport/memory_pool.py b/python/sglang/srt/multimodal/transport/memory_pool.py index a594c1c3e..dc32990ff 100644 --- a/python/sglang/srt/multimodal/transport/memory_pool.py +++ b/python/sglang/srt/multimodal/transport/memory_pool.py @@ -59,7 +59,6 @@ def resolve_consumer_rank( return 0 if consumer_rank is None: try: - # Use the global TP rank. An attention/DCP subgroup rank can alias # another consumer's acknowledgement slot. rank = int(get_parallel().tp_rank) diff --git a/python/sglang/srt/multimodal/vit_cuda_graph_runner.py b/python/sglang/srt/multimodal/vit_cuda_graph_runner.py index a2e6933ba..7a40a45ed 100644 --- a/python/sglang/srt/multimodal/vit_cuda_graph_runner.py +++ b/python/sglang/srt/multimodal/vit_cuda_graph_runner.py @@ -64,9 +64,9 @@ class ViTCudaGraphRunner: # captured before the workspace grew. self.sin_cos_ws: Optional[Tuple[torch.Tensor, torch.Tensor]] = None self._retired_sin_cos_ws: List[Tuple[torch.Tensor, torch.Tensor]] = [] - self._sin_cos_ws_by_graph: Dict[Hashable, Tuple[torch.Tensor, torch.Tensor]] = ( - {} - ) + self._sin_cos_ws_by_graph: Dict[ + Hashable, Tuple[torch.Tensor, torch.Tensor] + ] = {} self.max_context_len = getattr(vit, "max_context_len", None) # Qwen2.5-VL specific viarable. diff --git a/python/sglang/srt/multiplex/multiplexing_mixin.py b/python/sglang/srt/multiplex/multiplexing_mixin.py index 5a4e2ffc7..253d42b8f 100644 --- a/python/sglang/srt/multiplex/multiplexing_mixin.py +++ b/python/sglang/srt/multiplex/multiplexing_mixin.py @@ -31,7 +31,6 @@ logger = logging.getLogger(__name__) class SchedulerMultiplexMixin: - def init_pdmux(self: Scheduler): # The current split prefill batch self.split_prefill_batch: Optional[ScheduleBatch] = None diff --git a/python/sglang/srt/multiplex/pdmux_context.py b/python/sglang/srt/multiplex/pdmux_context.py index 05cde1371..2eaa61508 100644 --- a/python/sglang/srt/multiplex/pdmux_context.py +++ b/python/sglang/srt/multiplex/pdmux_context.py @@ -104,7 +104,12 @@ def divide_sm(total_sms, compute_capability, groups): def initialize_stream_groups(gpu_id: int, config: PDMuxConfig): from sgl_kernel import spatial - global STREAM_GROUPS, SM_COUNTS, SM_GROUP_NUM, CURRENT_STREAM_IDX, CURRENT_STREAM_GROUP + global \ + STREAM_GROUPS, \ + SM_COUNTS, \ + SM_GROUP_NUM, \ + CURRENT_STREAM_IDX, \ + CURRENT_STREAM_GROUP # for pd_multiplexing, Init stream_groups device = torch.cuda.current_device() total_sm_count = spatial.get_sm_available(gpu_id) diff --git a/python/sglang/srt/observability/metrics_collector.py b/python/sglang/srt/observability/metrics_collector.py index 6daebdfec..156454338 100644 --- a/python/sglang/srt/observability/metrics_collector.py +++ b/python/sglang/srt/observability/metrics_collector.py @@ -244,7 +244,6 @@ class SchedulerMetricsCollectorContext: class SchedulerMetricsCollector(_StatLoggerDIMixin): - def __init__( self, labels: Dict[str, str], diff --git a/python/sglang/srt/observability/req_time_stats.py b/python/sglang/srt/observability/req_time_stats.py index 6d32a2591..539634de6 100644 --- a/python/sglang/srt/observability/req_time_stats.py +++ b/python/sglang/srt/observability/req_time_stats.py @@ -1057,9 +1057,9 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): ) if SGLANG_TEST_REQUEST_TIME_STATS: - assert ( - queue_duration >= 0 and forward_duration >= 0 - ), f"queue_duration={queue_duration} < 0 or forward_duration={forward_duration} < 0" + assert queue_duration >= 0 and forward_duration >= 0, ( + f"queue_duration={queue_duration} < 0 or forward_duration={forward_duration} < 0" + ) return f"queue_duration={self.format_duration(queue_duration)}, forward_duration={self.format_duration(forward_duration)}, entry_time={self.format_wallclock(self.wait_queue_entry_time)}" elif self.disagg_mode == DisaggregationMode.PREFILL: @@ -1079,7 +1079,9 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): bootstrap_queue_duration >= 0 and queue_duration >= 0 and forward_duration >= 0 - ), f"bootstrap_queue_duration={bootstrap_queue_duration} < 0 or queue_duration={queue_duration} < 0 or forward_duration={forward_duration} < 0" + ), ( + f"bootstrap_queue_duration={bootstrap_queue_duration} < 0 or queue_duration={queue_duration} < 0 or forward_duration={forward_duration} < 0" + ) if ( self.bootstrap_done_time > 0 @@ -1089,9 +1091,9 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): self.prefill_bootstrap_queue_entry_time, self.bootstrap_done_time ) if SGLANG_TEST_REQUEST_TIME_STATS: - assert ( - bootstrap_duration >= 0 - ), f"bootstrap_duration={bootstrap_duration} < 0" + assert bootstrap_duration >= 0, ( + f"bootstrap_duration={bootstrap_duration} < 0" + ) bootstrap_fields = ( f"bootstrap_duration={self.format_duration(bootstrap_duration)}, " ) @@ -1133,7 +1135,9 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): and transfer_duration >= 0 and queue_duration >= 0 and forward_duration >= 0 - ), f"prealloc_duration={prealloc_duration} < 0 or transfer_duration={transfer_duration} < 0 or queue_duration={queue_duration} < 0 or forward_duration={forward_duration} < 0. {self=}" + ), ( + f"prealloc_duration={prealloc_duration} < 0 or transfer_duration={transfer_duration} < 0 or queue_duration={queue_duration} < 0 or forward_duration={forward_duration} < 0. {self=}" + ) # Break down prealloc_duration into sub-phases if self.bootstrap_done_time > 0: @@ -1144,9 +1148,9 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): self.bootstrap_done_time, self.decode_transfer_queue_entry_time ) if SGLANG_TEST_REQUEST_TIME_STATS: - assert ( - bootstrap_duration >= 0 and alloc_wait_duration >= 0 - ), f"bootstrap_duration={bootstrap_duration} < 0 or alloc_wait_duration={alloc_wait_duration} < 0" + assert bootstrap_duration >= 0 and alloc_wait_duration >= 0, ( + f"bootstrap_duration={bootstrap_duration} < 0 or alloc_wait_duration={alloc_wait_duration} < 0" + ) prealloc_fields = ( f"bootstrap_duration={self.format_duration(bootstrap_duration)}, " f"alloc_wait_duration={self.format_duration(alloc_wait_duration)}, " diff --git a/python/sglang/srt/parser/code_completion_parser.py b/python/sglang/srt/parser/code_completion_parser.py index 510f74468..7491445c1 100644 --- a/python/sglang/srt/parser/code_completion_parser.py +++ b/python/sglang/srt/parser/code_completion_parser.py @@ -58,9 +58,9 @@ completion_templates: dict[str, CompletionTemplate] = {} def register_completion_template(template: CompletionTemplate, override: bool = False): """Register a new completion template.""" if not override: - assert ( - template.name not in completion_templates - ), f"{template.name} has been registered." + assert template.name not in completion_templates, ( + f"{template.name} has been registered." + ) completion_templates[template.name] = template diff --git a/python/sglang/srt/parser/conversation.py b/python/sglang/srt/parser/conversation.py index 96d3a8fa9..585ecdca5 100644 --- a/python/sglang/srt/parser/conversation.py +++ b/python/sglang/srt/parser/conversation.py @@ -511,9 +511,9 @@ matching_function_registry: List[Callable] = [] def register_conv_template(template: Conversation, override: bool = False): """Register a new conversation template.""" if not override: - assert ( - template.name not in chat_templates - ), f"{template.name} has been registered." + assert template.name not in chat_templates, ( + f"{template.name} has been registered." + ) chat_templates[template.name] = template diff --git a/python/sglang/srt/parser/template_detection.py b/python/sglang/srt/parser/template_detection.py index 766243bfd..f14aa467d 100644 --- a/python/sglang/srt/parser/template_detection.py +++ b/python/sglang/srt/parser/template_detection.py @@ -153,27 +153,32 @@ REASONING_MODE_RULES = ( DetectionRule( name="force_reasoning_pattern", value=ReasoningToggleConfig(special_case="always"), - predicate=lambda ctx: ctx.has_pattern(r"<\|im_start\|>assistant\\n\\n") - and not ctx.has_text("enable_thinking") - and not ctx.has_text("thinking"), + predicate=lambda ctx: ( + ctx.has_pattern(r"<\|im_start\|>assistant\\n\\n") + and not ctx.has_text("enable_thinking") + and not ctx.has_text("thinking") + ), ), DetectionRule( name="mistral_reasoning_effort", value=ReasoningToggleConfig(special_case="mistral"), - predicate=lambda ctx: ctx.has_text("reasoning_effort") - and ctx.has_text("[THINK]"), + predicate=lambda ctx: ( + ctx.has_text("reasoning_effort") and ctx.has_text("[THINK]") + ), ), DetectionRule( name="explicit_enable_thinking_default_false", value=ReasoningToggleConfig( toggle_param="enable_thinking", default_enabled=False ), - predicate=lambda ctx: ctx.has_pattern( - r"{%\s*if\s+not\s+enable_thinking\s+is\s+defined\s*%}.*?" - r"{%\s*set\s+enable_thinking\s*=\s*(?:false|False)\s*%}", - re.DOTALL, - ) - or _has_toggle_default_assignment(ctx, "enable_thinking", False), + predicate=lambda ctx: ( + ctx.has_pattern( + r"{%\s*if\s+not\s+enable_thinking\s+is\s+defined\s*%}.*?" + r"{%\s*set\s+enable_thinking\s*=\s*(?:false|False)\s*%}", + re.DOTALL, + ) + or _has_toggle_default_assignment(ctx, "enable_thinking", False) + ), ), DetectionRule( name="nemotron_3_super_low_effort", @@ -182,58 +187,65 @@ REASONING_MODE_RULES = ( default_enabled=True, effort_kwarg="low_effort", ), - predicate=lambda ctx: ctx.has_text("low_effort") - and ctx.has_text("truncate_history_thinking"), + predicate=lambda ctx: ( + ctx.has_text("low_effort") and ctx.has_text("truncate_history_thinking") + ), ), DetectionRule( name="enable_thinking_default_true", value=ReasoningToggleConfig( toggle_param="enable_thinking", default_enabled=True ), - predicate=lambda ctx: ctx.has_pattern( - r"{%\s*if\s+not\s+enable_thinking\s+is\s+defined\s*%}.*?" - r"{%\s*set\s+enable_thinking\s*=\s*(?:true|True)\s*%}", - re.DOTALL, - ) - or ctx.has_pattern( - r"set\s+enable_thinking\s*=\s*enable_thinking\s+if\s+enable_thinking\s+is\s+defined\s+else\s+(?:true|True)" - ) - or ctx.has_pattern( - r"enable_thinking\s+is\s+defined\s+and\s+(?:enable_thinking\s+is\s+false|not\s+enable_thinking)" - ) - or ctx.has_pattern( - r"enable_thinking\s+is\s+not\s+defined\s+or\s+enable_thinking" - ) - or ctx.has_pattern(r"namespace\([^)]*enable_thinking\s*=\s*true") - or _has_toggle_default_assignment(ctx, "enable_thinking", True), + predicate=lambda ctx: ( + ctx.has_pattern( + r"{%\s*if\s+not\s+enable_thinking\s+is\s+defined\s*%}.*?" + r"{%\s*set\s+enable_thinking\s*=\s*(?:true|True)\s*%}", + re.DOTALL, + ) + or ctx.has_pattern( + r"set\s+enable_thinking\s*=\s*enable_thinking\s+if\s+enable_thinking\s+is\s+defined\s+else\s+(?:true|True)" + ) + or ctx.has_pattern( + r"enable_thinking\s+is\s+defined\s+and\s+(?:enable_thinking\s+is\s+false|not\s+enable_thinking)" + ) + or ctx.has_pattern( + r"enable_thinking\s+is\s+not\s+defined\s+or\s+enable_thinking" + ) + or ctx.has_pattern(r"namespace\([^)]*enable_thinking\s*=\s*true") + or _has_toggle_default_assignment(ctx, "enable_thinking", True) + ), ), DetectionRule( name="explicit_thinking_default_false", value=ReasoningToggleConfig(toggle_param="thinking", default_enabled=False), - predicate=lambda ctx: ctx.has_pattern( - r"{%\s*if\s+not\s+thinking\s+is\s+defined\s*%}.*?" - r"{%\s*set\s+thinking\s*=\s*(?:false|False)\s*%}", - re.DOTALL, - ) - or _has_toggle_default_assignment(ctx, "thinking", False), + predicate=lambda ctx: ( + ctx.has_pattern( + r"{%\s*if\s+not\s+thinking\s+is\s+defined\s*%}.*?" + r"{%\s*set\s+thinking\s*=\s*(?:false|False)\s*%}", + re.DOTALL, + ) + or _has_toggle_default_assignment(ctx, "thinking", False) + ), ), DetectionRule( name="thinking_default_true", value=ReasoningToggleConfig(toggle_param="thinking", default_enabled=True), - predicate=lambda ctx: ctx.has_pattern( - r"{%\s*if\s+not\s+thinking\s+is\s+defined\s*%}.*?" - r"{%\s*set\s+thinking\s*=\s*(?:true|True)\s*%}", - re.DOTALL, - ) - or ctx.has_pattern( - r"set\s+thinking\s*=\s*thinking\s+if\s+thinking\s+is\s+defined\s+else\s+(?:true|True)" - ) - or ctx.has_pattern( - r"thinking\s+is\s+defined\s+and\s+(?:thinking\s+is\s+false|not\s+thinking)" - ) - or ctx.has_pattern(r"thinking\s+is\s+not\s+defined\s+or\s+thinking") - or ctx.has_pattern(r"namespace\([^)]*thinking\s*=\s*true") - or _has_toggle_default_assignment(ctx, "thinking", True), + predicate=lambda ctx: ( + ctx.has_pattern( + r"{%\s*if\s+not\s+thinking\s+is\s+defined\s*%}.*?" + r"{%\s*set\s+thinking\s*=\s*(?:true|True)\s*%}", + re.DOTALL, + ) + or ctx.has_pattern( + r"set\s+thinking\s*=\s*thinking\s+if\s+thinking\s+is\s+defined\s+else\s+(?:true|True)" + ) + or ctx.has_pattern( + r"thinking\s+is\s+defined\s+and\s+(?:thinking\s+is\s+false|not\s+thinking)" + ) + or ctx.has_pattern(r"thinking\s+is\s+not\s+defined\s+or\s+thinking") + or ctx.has_pattern(r"namespace\([^)]*thinking\s*=\s*true") + or _has_toggle_default_assignment(ctx, "thinking", True) + ), ), ) diff --git a/python/sglang/srt/parser/template_manager.py b/python/sglang/srt/parser/template_manager.py index d69458f87..abb9daf9f 100644 --- a/python/sglang/srt/parser/template_manager.py +++ b/python/sglang/srt/parser/template_manager.py @@ -290,9 +290,9 @@ class TemplateManager: def _load_json_chat_template(self, template_path: str) -> None: """Load a JSON chat template file.""" - assert template_path.endswith( - ".json" - ), "unrecognized format of chat template file" + assert template_path.endswith(".json"), ( + "unrecognized format of chat template file" + ) with open(template_path, "r") as filep: template = json.load(filep) @@ -319,9 +319,9 @@ class TemplateManager: def _load_json_completion_template(self, template_path: str) -> None: """Load a JSON completion template file.""" - assert template_path.endswith( - ".json" - ), "unrecognized format of completion template file" + assert template_path.endswith(".json"), ( + "unrecognized format of completion template file" + ) with open(template_path, "r") as filep: template = json.load(filep) diff --git a/python/sglang/srt/plugins/hook_registry.py b/python/sglang/srt/plugins/hook_registry.py index c577b5232..4def317da 100644 --- a/python/sglang/srt/plugins/hook_registry.py +++ b/python/sglang/srt/plugins/hook_registry.py @@ -269,9 +269,7 @@ class HookRegistry: # Sort: REPLACE hooks first (stable sort preserves registration order # within the same type). This ensures AROUND/BEFORE/AFTER always wrap # the replaced function, regardless of registration order. - sorted_hooks = sorted( - hooks, key=lambda h: (0 if h[0] == HookType.REPLACE else 1) - ) + sorted_hooks = sorted(hooks, key=lambda h: 0 if h[0] == HookType.REPLACE else 1) wrapped = original for hook_type, hook, _src in sorted_hooks: if isinstance(hook, type) and hook_type == HookType.REPLACE: diff --git a/python/sglang/srt/rust_server/server.py b/python/sglang/srt/rust_server/server.py index e6aa021c0..1e4a9ddc0 100644 --- a/python/sglang/srt/rust_server/server.py +++ b/python/sglang/srt/rust_server/server.py @@ -230,9 +230,9 @@ class RustServer: # Invariant: control requests always carry a rust-minted rid; without # one the response is unroutable, so fail loudly rather than drop it. - assert ( - recv_req.rid is not None - ), f"control response without rid: {type(output).__name__}" + assert recv_req.rid is not None, ( + f"control response without rid: {type(output).__name__}" + ) # No local try/except: a failed push propagates to run_scheduler_process's # outer handler, which logs the full traceback (scheduler-fatal either way). payload = ( @@ -359,7 +359,9 @@ class RustServer: assert len(col) in ( 0, batch_size, - ), f"extras column {name}: {len(col)} entries for a batch of {batch_size}" + ), ( + f"extras column {name}: {len(col)} entries for a batch of {batch_size}" + ) populated |= len(col) > 0 if populated: active.append(extra) diff --git a/python/sglang/srt/sampling/custom_logit_processor.py b/python/sglang/srt/sampling/custom_logit_processor.py index a6073ae81..d66d872bd 100644 --- a/python/sglang/srt/sampling/custom_logit_processor.py +++ b/python/sglang/srt/sampling/custom_logit_processor.py @@ -51,9 +51,9 @@ class DisallowedTokensLogitsProcessor(CustomLogitProcessor): custom_param_list: Optional[List[Dict[str, Any]]] = None, ) -> torch.Tensor: disallowed_token_ids = custom_param_list[0]["token_ids"] - assert all( - disallowed_token_ids == c["token_ids"] for c in custom_param_list - ), f"{custom_param_list=}" + assert all(disallowed_token_ids == c["token_ids"] for c in custom_param_list), ( + f"{custom_param_list=}" + ) logits[..., disallowed_token_ids] = -float("inf") return logits diff --git a/python/sglang/srt/sampling/penaltylib/min_new_tokens.py b/python/sglang/srt/sampling/penaltylib/min_new_tokens.py index c77cb7851..bf47b92b4 100644 --- a/python/sglang/srt/sampling/penaltylib/min_new_tokens.py +++ b/python/sglang/srt/sampling/penaltylib/min_new_tokens.py @@ -60,9 +60,7 @@ class BatchedMinNewTokensPenalizer(_BatchedPenalizer): fill_value=float("-inf"), device=self.orchestrator.device, ), - )[ - :, : self.orchestrator.vocab_size - ] + )[:, : self.orchestrator.vocab_size] self.len_output_tokens = torch.zeros( size=(len(self.orchestrator.reqs()), 1), diff --git a/python/sglang/srt/sampling/sampling_params.py b/python/sglang/srt/sampling/sampling_params.py index da340e56c..fdc8d8377 100644 --- a/python/sglang/srt/sampling/sampling_params.py +++ b/python/sglang/srt/sampling/sampling_params.py @@ -171,12 +171,11 @@ class SamplingParams(msgspec.Struct, kw_only=True, array_like=True): ) if not -2.0 <= self.frequency_penalty <= 2.0: raise ValueError( - "frequency_penalty must be in [-2, 2], got " - f"{self.frequency_penalty}." + f"frequency_penalty must be in [-2, 2], got {self.frequency_penalty}." ) if not -2.0 <= self.presence_penalty <= 2.0: raise ValueError( - "presence_penalty must be in [-2, 2], got " f"{self.presence_penalty}." + f"presence_penalty must be in [-2, 2], got {self.presence_penalty}." ) if not 0.0 < self.repetition_penalty <= 2.0: raise ValueError( diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index fa4b80ed3..94392accd 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1286,9 +1286,7 @@ class ServerArgs: "Initial connection-level HTTP/2 receive window in bytes (1024 to " "2^31 - 1). Only applies with --enable-http2.", NS("serving"), - ] = ( - 1024 * 1024 - ) + ] = 1024 * 1024 # ------------------------------------------------------------------------- # SSL/TLS @@ -2334,7 +2332,7 @@ class ServerArgs: ] = 18 speculative_ngram_capacity: A[ int, "The cache capacity for ngram speculative decoding.", NS("spec") - ] = (10 * 1000 * 1000) + ] = 10 * 1000 * 1000 speculative_ngram_external_corpus_path: A[ Optional[str], "Path to an external JSONL corpus to pre-load into SAM at startup. Additional corpora can be added at runtime via POST /add_external_corpus.", diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index 4489b3f00..522790a53 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/python/sglang/srt/speculative/dflash_utils.py @@ -222,8 +222,7 @@ def apply_dflash_verify_logits_adjustments( return if next_token_logits.ndim != 2: raise ValueError( - "next_token_logits must be 2D, " - f"got shape={tuple(next_token_logits.shape)}." + f"next_token_logits must be 2D, got shape={tuple(next_token_logits.shape)}." ) if draft_token_num <= 0: raise ValueError(f"draft_token_num must be positive, got {draft_token_num}.") @@ -894,8 +893,7 @@ def compute_dflash_sampling_correct_drafts_and_bonus( raise ValueError(f"candidates must be 2D, got shape={tuple(candidates.shape)}") if next_token_logits.ndim != 2: raise ValueError( - "next_token_logits must be 2D, " - f"got shape={tuple(next_token_logits.shape)}." + f"next_token_logits must be 2D, got shape={tuple(next_token_logits.shape)}." ) bs, draft_token_num = candidates.shape @@ -915,10 +913,8 @@ def compute_dflash_sampling_correct_drafts_and_bonus( ) if threshold_single is None: - threshold_single = get_spec().speculative_accept_threshold_single if threshold_acc is None: - threshold_acc = get_spec().speculative_accept_threshold_acc threshold_single = float(threshold_single) threshold_acc = max(float(threshold_acc), 1e-9) diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index 1652e219e..7df89d5a8 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -1727,7 +1727,9 @@ class DFlashWorkerV2(BaseSpecWorker): ] self._accept_bonus_buffer_cap = new_cap - def _next_accept_bonus_buffers(self, bs: int) -> tuple[ + def _next_accept_bonus_buffers( + self, bs: int + ) -> tuple[ torch.Tensor, torch.Tensor, torch.Tensor, diff --git a/python/sglang/srt/speculative/dspark_components/dspark_observability.py b/python/sglang/srt/speculative/dspark_components/dspark_observability.py index e08d11d38..0a1c026f5 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_observability.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_observability.py @@ -495,7 +495,6 @@ def _format_float(value: float, digits: int = 4) -> str: class PerPositionConfidenceMetrics: - def __init__( self, *, @@ -655,7 +654,6 @@ class PerPositionConfidenceMetrics: class ConfidenceMetricsProbe: - def __init__( self, *, diff --git a/python/sglang/srt/speculative/dspark_components/dspark_planner.py b/python/sglang/srt/speculative/dspark_components/dspark_planner.py index 45e7c0544..e56026572 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_planner.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_planner.py @@ -260,9 +260,9 @@ class DSparkVerifyPlanner: return None compute_confidence_hook = getattr(self.draft_model, "compute_confidence", None) if compute_confidence_hook is not None: - assert ( - confidence_tap is not None - ), "dsv4 compute_confidence needs the compute_base_logits tap" + assert confidence_tap is not None, ( + "dsv4 compute_confidence needs the compute_base_logits tap" + ) with torch.inference_mode(): return compute_confidence_hook( anchor_tokens=anchor_tokens, @@ -1001,7 +1001,6 @@ def _additive_step_time_tensor( class HostConfidenceBudgetPlanner: - def __init__( self, *, diff --git a/python/sglang/srt/speculative/dspark_components/dspark_sps.py b/python/sglang/srt/speculative/dspark_components/dspark_sps.py index f19483f2d..bfefe1219 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_sps.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_sps.py @@ -60,7 +60,6 @@ def _interp_clamped(xs: list[int], ys: list[float], x: float) -> float: class SpsAdditiveCostTable(msgspec.Struct, frozen=True): - bias_seconds: float bs_probes: list[int] alpha_seconds: list[float] diff --git a/python/sglang/srt/speculative/dspark_components/dspark_verify.py b/python/sglang/srt/speculative/dspark_components/dspark_verify.py index cb4fc4e6d..f118c2d60 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_verify.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_verify.py @@ -475,7 +475,6 @@ class TargetVerifyExecutor: class CommitInjectCtx(msgspec.Struct): - draft_model: object block_pos_offsets: torch.Tensor resolve_pool: object @@ -492,7 +491,6 @@ class AcceptOuts(msgspec.Struct): class DsparkVerifyEpilogue: - def __init__( self, *, diff --git a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py index 4a5e68408..9b4a74a14 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -125,7 +125,6 @@ def _configure_target_hidden_projection( class DSparkWorkerV2(BaseSpecWorker): - def __init__( self, server_args: ServerArgs, diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 48e46f8a2..e89e61155 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -170,8 +170,12 @@ class EagleDraftWorker(EagleDraftWorkerBase): else: ctx = empty_context() with ( - ctx - ), draft_pp_context(), speculative_moe_backend_context(), speculative_moe_a2a_backend_context(), draft_model_build_scope(): + ctx, + draft_pp_context(), + speculative_moe_backend_context(), + speculative_moe_a2a_backend_context(), + draft_model_build_scope(), + ): self.draft_worker = TpModelWorker( server_args=server_args, gpu_id=gpu_id, @@ -999,9 +1003,7 @@ class EagleDraftWorker(EagleDraftWorkerBase): dsa_seed_topk_indices = None if self.seed_dsa_topk_from_draft_extend: if can_run_decode_cuda_graph: - dsa_extend_topk_capture = ( - self.cuda_graph_runner_for_draft_extend.buffers.dsa_seed_topk_capture - ) + dsa_extend_topk_capture = self.cuda_graph_runner_for_draft_extend.buffers.dsa_seed_topk_capture else: dsa_extend_topk_capture = forward_batch.spec_info.dsa_seed_topk_capture # Fancy indexing returns a fresh tensor (detached from the buffer). diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py index 056f73e7e..42a0d212a 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py @@ -134,8 +134,11 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker): self.hot_token_id = None with ( - empty_context() - ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context(), draft_model_build_scope(): + empty_context(), + speculative_moe_backend_context(), + speculative_moe_a2a_backend_context(), + draft_model_build_scope(), + ): # Both base classes own initialization, so initialize TpModelWorker # explicitly after EagleDraftWorkerBase above. TpModelWorker.__init__( @@ -720,9 +723,9 @@ class FrozenKVMTPWorkerV2(EAGLEWorkerV2): ) # Frozen MTP does not wire the adaptive controller yet. - assert ( - not get_spec().speculative_adaptive - ), "Frozen-KV MTP does not support adaptive speculative decoding yet." + assert not get_spec().speculative_adaptive, ( + "Frozen-KV MTP does not support adaptive speculative decoding yet." + ) self.adaptive_controller = None # Some dummy tensors (parity with EAGLEWorkerV2 init). diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index e5c21ceda..472537b7e 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -159,7 +159,11 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): ) # Load draft model weights only. - with empty_context(), speculative_moe_backend_context(), draft_model_build_scope(): + with ( + empty_context(), + speculative_moe_backend_context(), + draft_model_build_scope(), + ): self.draft_worker = TpModelWorker( server_args=server_args, gpu_id=gpu_id, @@ -376,9 +380,9 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): draft_backend_factory.create_draft_extend_backend() ) if self.draft_extend_attn_backend_list[-1] is not None: - self.draft_runner_list[step].attn_backend = ( - self.draft_extend_attn_backend_list[-1] - ) + self.draft_runner_list[ + step + ].attn_backend = self.draft_extend_attn_backend_list[-1] def _capture_cuda_graphs(self): self.cuda_graph_runner = None diff --git a/python/sglang/srt/speculative/ngram_worker.py b/python/sglang/srt/speculative/ngram_worker.py index 7d8c6c68b..826895b28 100644 --- a/python/sglang/srt/speculative/ngram_worker.py +++ b/python/sglang/srt/speculative/ngram_worker.py @@ -306,9 +306,9 @@ class NGRAMWorker(BaseSpecWorker): total_draft_token_num = len(req_drafts) # Check if speculative decoding is needed; here we always enforce it - assert ( - total_draft_token_num == bs * self.draft_token_num - ), f"{total_draft_token_num=}, {bs=}, {self.draft_token_num=}" + assert total_draft_token_num == bs * self.draft_token_num, ( + f"{total_draft_token_num=}, {bs=}, {self.draft_token_num=}" + ) return req_drafts, mask def _prepare_for_speculative_decoding(self, batch: ScheduleBatch): diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index f07f153eb..0cca5c007 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/python/sglang/srt/speculative/spec_info.py @@ -293,9 +293,9 @@ class SpeculativeAlgorithm(Enum): ) -> Optional[Union[Type[BaseSpecWorker], Type[TpModelWorker], Type[NGRAMWorker]]]: cfg = resolving_view(server_args) - assert ( - not self.is_none() - ), "Cannot create worker for NONE speculative algorithm." + assert not self.is_none(), ( + "Cannot create worker for NONE speculative algorithm." + ) if self.is_dflash(): # V2 worker drives both overlap and non-overlap (scheduler runs it diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 25db452c5..464ff501a 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -197,9 +197,9 @@ def draft_kv_indices_buffer_width( num_seqs * topk branches each attend up to max_context_len KV slots; the topk factor is mandatory -- dropping it under-allocates and overflows the row (#27338, #27460). """ - assert ( - num_seqs * topk * max_context_len < 2**31 - ), "kv_indices flat offset would overflow int32; reduce batch/topk/context" + assert num_seqs * topk * max_context_len < 2**31, ( + "kv_indices flat offset would overflow int32; reduce batch/topk/context" + ) return num_seqs * topk * max_context_len diff --git a/python/sglang/srt/speculative/standalone_worker_v2.py b/python/sglang/srt/speculative/standalone_worker_v2.py index ac37350b1..a270cb278 100644 --- a/python/sglang/srt/speculative/standalone_worker_v2.py +++ b/python/sglang/srt/speculative/standalone_worker_v2.py @@ -150,7 +150,6 @@ class StandaloneDraftWorker(EagleDraftWorker): class StandaloneWorkerV2(EAGLEWorkerV2): - def __init__( self, server_args: ServerArgs, diff --git a/python/sglang/srt/tokenizer/tiktoken_tokenizer.py b/python/sglang/srt/tokenizer/tiktoken_tokenizer.py index 866c1ad24..cd061c7af 100644 --- a/python/sglang/srt/tokenizer/tiktoken_tokenizer.py +++ b/python/sglang/srt/tokenizer/tiktoken_tokenizer.py @@ -76,9 +76,7 @@ class TiktokenTokenizer: self, text: str, *, - allowed_special: Union[ - Literal["all"], AbstractSet[str] - ] = set(), # noqa: B006 + allowed_special: Union[Literal["all"], AbstractSet[str]] = set(), # noqa: B006 disallowed_special: Union[Literal["all"], Collection[str]] = "all", ) -> List[int]: if isinstance(allowed_special, set): diff --git a/python/sglang/srt/utils/bench_utils.py b/python/sglang/srt/utils/bench_utils.py index ccb811482..a9d9ede3c 100644 --- a/python/sglang/srt/utils/bench_utils.py +++ b/python/sglang/srt/utils/bench_utils.py @@ -129,7 +129,9 @@ def bench_kineto( assert ( sum([int(re.search(name, line) is not None) for line in prof_lines]) == 1 - ), f"Errors of the kernel {name} in the profiling table (table: {prof_lines})" + ), ( + f"Errors of the kernel {name} in the profiling table (table: {prof_lines})" + ) # Save chrome traces if trace_path is not None: diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 87958d618..6100f9e02 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -1431,7 +1431,6 @@ def calculate_time(show=False, min_cost_ms=0.0): class LayerFn(Protocol): - def __call__(self, idx: int, prefix: str) -> torch.nn.Module: ... @@ -2460,7 +2459,9 @@ def broadcast_pyobj( device = torch.device( "cuda" if torch.cuda.is_available() and not force_cpu_device - else "musa" if is_musa() and not force_cpu_device else "cpu" + else "musa" + if is_musa() and not force_cpu_device + else "cpu" ) if rank == src: @@ -2735,9 +2736,9 @@ def init_custom_process_group( rendezvous, ) - assert (store is None) or ( - init_method is None - ), "Cannot specify both init_method and store." + assert (store is None) or (init_method is None), ( + "Cannot specify both init_method and store." + ) if store is not None: assert world_size > 0, "world_size must be positive if using store" @@ -3257,13 +3258,13 @@ class UvicornAccessLogFilter(logging.Filter): def set_uvicorn_logging_configs(server_args=None): from uvicorn.config import LOGGING_CONFIG - LOGGING_CONFIG["formatters"]["default"][ - "fmt" - ] = "[%(asctime)s] %(levelprefix)s %(message)s" + LOGGING_CONFIG["formatters"]["default"]["fmt"] = ( + "[%(asctime)s] %(levelprefix)s %(message)s" + ) LOGGING_CONFIG["formatters"]["default"]["datefmt"] = "%Y-%m-%d %H:%M:%S" - LOGGING_CONFIG["formatters"]["access"][ - "fmt" - ] = '[%(asctime)s] %(levelprefix)s %(client_addr)s - "%(request_line)s" %(status_code)s' + LOGGING_CONFIG["formatters"]["access"]["fmt"] = ( + '[%(asctime)s] %(levelprefix)s %(client_addr)s - "%(request_line)s" %(status_code)s' + ) LOGGING_CONFIG["formatters"]["access"]["datefmt"] = "%Y-%m-%d %H:%M:%S" _configure_uvicorn_access_log_filter(LOGGING_CONFIG, server_args) @@ -3964,9 +3965,9 @@ def _process_weight_after_loading(module, weight_names, transpose_dims=None) -> device = devices.pop() if transpose_dims: - assert len(weight_names) == len( - transpose_dims - ), "len(weight_names) should be equal to len(transpose_dims)" + assert len(weight_names) == len(transpose_dims), ( + "len(weight_names) should be equal to len(transpose_dims)" + ) for i, weight_name in enumerate(weight_names): weight_tensor = getattr(module, weight_name) @@ -4106,7 +4107,7 @@ def configure_gc_logger(): logger.info( f"GC end: Time {time.time()} | Generation {gen} | " f"Duration: {duration:.4f}s | Collected: {collected} | Uncollectable: {uncollectable} " - f'{"(LONG GC)" if duration > 0.1 else ""}' + f"{'(LONG GC)' if duration > 0.1 else ''}" ) gc.callbacks.append(gc_callback) @@ -4186,9 +4187,9 @@ def get_physical_cpus_by_numa(): for cpu, core, socket, node in cpu_info: key = (core, socket) if key not in physical_by_node[node]: - physical_by_node[node][ - key - ] = cpu # pick first CPU seen for that physical core + physical_by_node[node][key] = ( + cpu # pick first CPU seen for that physical core + ) # Retrieves CPUs that the current process is allowed to run on cpus_allowed_list = psutil.Process().cpu_affinity() @@ -4577,9 +4578,9 @@ class CachedKernel: # Check that no parameters have default values for name, param in self.signature.parameters.items(): - assert ( - param.default is inspect.Parameter.empty - ), f"Parameter '{name}' has a default value. Default parameters are not supported in cached kernels." + assert param.default is inspect.Parameter.empty, ( + f"Parameter '{name}' has a default value. Default parameters are not supported in cached kernels." + ) functools.update_wrapper(self, original_fn) self.kernel_cache = {} @@ -4592,9 +4593,9 @@ class CachedKernel: Index with grid to get a launcher function. Returns a launcher that will handle caching based on the key function. """ - assert ( - isinstance(grid, tuple) and len(grid) <= 3 - ), "Grid must be a tuple with at most 3 dimensions." + assert isinstance(grid, tuple) and len(grid) <= 3, ( + "Grid must be a tuple with at most 3 dimensions." + ) # Normalize grid once if len(grid) < 3: diff --git a/python/sglang/srt/utils/custom_op.py b/python/sglang/srt/utils/custom_op.py index 7895393b2..aadb857e5 100644 --- a/python/sglang/srt/utils/custom_op.py +++ b/python/sglang/srt/utils/custom_op.py @@ -103,15 +103,15 @@ def register_custom_op( """ extra_kwarg_keys = set(extra_kwargs.keys()) expected_kwarg_keys = set({"out_shape", "fake_impl"}) - assert ( - expected_kwarg_keys >= extra_kwarg_keys - ), f"Unexpected extra kwargs: {extra_kwarg_keys - expected_kwarg_keys}" + assert expected_kwarg_keys >= extra_kwarg_keys, ( + f"Unexpected extra kwargs: {extra_kwarg_keys - expected_kwarg_keys}" + ) has_out_shape = "out_shape" in extra_kwargs has_fake_impl = "fake_impl" in extra_kwargs - assert not ( - has_out_shape and has_fake_impl - ), "Only one of `out_shape` or `fake_impl` should be provided." + assert not (has_out_shape and has_fake_impl), ( + "Only one of `out_shape` or `fake_impl` should be provided." + ) # Assume inplace if neither out_shape nor fake_impl is provided if not (has_out_shape or has_fake_impl): extra_kwargs["out_shape"] = None @@ -261,9 +261,9 @@ def register_custom_op_from_extern( name = op_name or fn.__name__ computed_args = computed_args or {} - assert not ( - out_shape is not None and fake_impl is not None - ), "Only one of `out_shape` or `fake_impl` should be provided." + assert not (out_shape is not None and fake_impl is not None), ( + "Only one of `out_shape` or `fake_impl` should be provided." + ) # If computed_args specified, create a wrapper with a reduced signature # that computes the excluded args inside the op body. diff --git a/python/sglang/srt/utils/flatten.py b/python/sglang/srt/utils/flatten.py index 0ff6b4e12..ba5324be3 100644 --- a/python/sglang/srt/utils/flatten.py +++ b/python/sglang/srt/utils/flatten.py @@ -106,9 +106,9 @@ class FlatPairColumns: # decoder cannot catch it — the data buffer only grows, so the receiver's # bounds check still passes and it hands the client another column's bytes # reinterpreted as logprobs, with a 200. - assert len(ii) == len( - vv - ), f"{self.name}: request {j} has {len(ii)} idx entries but {len(vv)} vals" + assert len(ii) == len(vv), ( + f"{self.name}: request {j} has {len(ii)} idx entries but {len(vv)} vals" + ) if self.first_none_to_nan and vv and vv[0] is None: self.v.append(float("nan")) self.v.extend(vv[1:]) diff --git a/python/sglang/srt/utils/hf_transformers/mistral_utils.py b/python/sglang/srt/utils/hf_transformers/mistral_utils.py index 71cd8cc2a..c9b59d962 100644 --- a/python/sglang/srt/utils/hf_transformers/mistral_utils.py +++ b/python/sglang/srt/utils/hf_transformers/mistral_utils.py @@ -81,9 +81,9 @@ def adapt_config_dict( else: config_dict["architectures"] = ["MistralLarge3ForCausalLM"] - assert ( - "llama_4_scaling" in config_dict - ), "MistralLarge3 expect llama4 scaling config." + assert "llama_4_scaling" in config_dict, ( + "MistralLarge3 expect llama4 scaling config." + ) llama_4_scaling_config_keys = ["original_max_position_embeddings", "beta"] assert all( [ diff --git a/python/sglang/srt/utils/hf_transformers_patches.py b/python/sglang/srt/utils/hf_transformers_patches.py index 614fc95ba..83f40aa3a 100644 --- a/python/sglang/srt/utils/hf_transformers_patches.py +++ b/python/sglang/srt/utils/hf_transformers_patches.py @@ -228,8 +228,8 @@ def _patch_removed_symbols(): if not hasattr(_u, "is_flash_attn_greater_or_equal_2_10"): if hasattr(_u, "is_flash_attn_greater_or_equal"): - _u.is_flash_attn_greater_or_equal_2_10 = ( - lambda: _u.is_flash_attn_greater_or_equal("2.10.0") + _u.is_flash_attn_greater_or_equal_2_10 = lambda: ( + _u.is_flash_attn_greater_or_equal("2.10.0") ) else: _u.is_flash_attn_greater_or_equal_2_10 = lambda: False diff --git a/python/sglang/srt/utils/numa_utils.py b/python/sglang/srt/utils/numa_utils.py index af50e6d9b..9f497a84a 100644 --- a/python/sglang/srt/utils/numa_utils.py +++ b/python/sglang/srt/utils/numa_utils.py @@ -93,9 +93,9 @@ def _mp_set_executable(executable: str, debug_str: str): try: yield finally: - assert ( - os.fsdecode(multiprocessing.spawn.get_executable()) == executable - ), f"{multiprocessing.spawn.get_executable()=}" + assert os.fsdecode(multiprocessing.spawn.get_executable()) == executable, ( + f"{multiprocessing.spawn.get_executable()=}" + ) multiprocessing.spawn.set_executable(old_executable) logger.debug(f"mp.set_executable revert to {old_executable}") diff --git a/python/sglang/srt/utils/offloader.py b/python/sglang/srt/utils/offloader.py index 694fd92ca..2170ce9c5 100644 --- a/python/sglang/srt/utils/offloader.py +++ b/python/sglang/srt/utils/offloader.py @@ -71,9 +71,9 @@ def create_offloader(dp_rank: int): cpu_offload_max_bytes=int(get_exec().offload.cpu_offload_gb * 1024**3) ) if get_exec().offload.offload_group_size > 0: - assert ( - get_exec().offload.cpu_offload_gb == 0 - ), "V2 offload does not support cpu_offload_gb yet" + assert get_exec().offload.cpu_offload_gb == 0, ( + "V2 offload does not support cpu_offload_gb yet" + ) return OffloaderV2( group_size=get_exec().offload.offload_group_size, num_in_group=get_exec().offload.offload_num_in_group, @@ -173,7 +173,6 @@ class OffloaderV2(BaseOffloader): # Temporarily init inside Offloader, can move if other modules also need this if self.mode in {"sharded_gpu", "shm_cpu"}: - assert get_parallel().tp_size == 1, "not yet support tp_size!=1" set_naive_distributed( NaiveDistributed( @@ -288,17 +287,17 @@ class _ModuleOffloader(ABC): self.device = next(module.parameters()).device self.alt_stream = alt_stream - assert self.device != torch.device( - "cpu" - ), "not handled device=cpu case yet (should skip this tensor)" + assert self.device != torch.device("cpu"), ( + "not handled device=cpu case yet (should skip this tensor)" + ) self._device_tensors = None self._load_event = None param_dict = dict(self.module.named_parameters()) - assert all( - name in param_dict for name in whitelist_param_names - ), f"{whitelist_param_names=} {list(param_dict.keys())=}" + assert all(name in param_dict for name in whitelist_param_names), ( + f"{whitelist_param_names=} {list(param_dict.keys())=}" + ) self._param_offloaders = { name: _BaseParamOffloader.create(mode, module=module, param_name=name) @@ -391,9 +390,9 @@ class _ShmCpuParamOffloader(_BaseParamOffloader): self._world_size = get_naive_distributed().get_world_size() assert get_parallel().tp_size == 1, "not yet support tp_size!=1" - assert ( - self._param.data.is_contiguous() - ), f"not yet support non-contiguous tensor {self._param.shape=} {self._param.stride()=}" + assert self._param.data.is_contiguous(), ( + f"not yet support non-contiguous tensor {self._param.shape=} {self._param.stride()=}" + ) self.shm_cpu_data = get_host_shared_memory_manager().malloc( shape=self._param.shape, dtype=self._param.dtype @@ -408,9 +407,9 @@ class _ShmCpuParamOffloader(_BaseParamOffloader): def post_init(self): if self._rank == 0: - assert ( - self.shm_cpu_data.data_ptr() == self._param.data.data_ptr() - ), f"{self.shm_cpu_data.data_ptr()=} {self._param.data.data_ptr()=} {self.shm_cpu_data=} {self._param.data=}" + assert self.shm_cpu_data.data_ptr() == self._param.data.data_ptr(), ( + f"{self.shm_cpu_data.data_ptr()=} {self._param.data.data_ptr()=} {self.shm_cpu_data=} {self._param.data=}" + ) _move_param_to_meta(self._module, self._param_name) @@ -424,9 +423,9 @@ def update_param(param, new_tensor): if param.device == new_tensor.device: param.data = new_tensor else: - assert param.device == torch.device( - "cpu" - ), f"{param.device=} {new_tensor.device=}" + assert param.device == torch.device("cpu"), ( + f"{param.device=} {new_tensor.device=}" + ) param.data = _create_cpu_data(new_tensor, pin_memory=True) @@ -496,9 +495,9 @@ class _ShardedGpuParamOffloader(_BaseParamOffloader): self._world_size = get_naive_distributed().get_world_size() assert get_parallel().tp_size == 1, "not yet support tp_size!=1" - assert ( - self._param.data.is_contiguous() - ), f"not yet support non-contiguous tensor {self._param.shape=} {self._param.stride()=}" + assert self._param.data.is_contiguous(), ( + f"not yet support non-contiguous tensor {self._param.shape=} {self._param.stride()=}" + ) if self._rank == 0: _move_param_to_cpu(self._param, pin_memory=True) @@ -509,9 +508,9 @@ class _ShardedGpuParamOffloader(_BaseParamOffloader): def post_init(self): # check again since it may be changed - assert ( - self._param.data.is_contiguous() - ), f"not yet support non-contiguous tensor {self._param.shape=} {self._param.stride()=}" + assert self._param.data.is_contiguous(), ( + f"not yet support non-contiguous tensor {self._param.shape=} {self._param.stride()=}" + ) scatter_src = self._param.data diff --git a/python/sglang/srt/utils/patch_tokenizer.py b/python/sglang/srt/utils/patch_tokenizer.py index ad6407a4f..3da8c2e64 100644 --- a/python/sglang/srt/utils/patch_tokenizer.py +++ b/python/sglang/srt/utils/patch_tokenizer.py @@ -67,14 +67,14 @@ class _SpecialTokensCachePatcher: ) def patched_add_special_tokens(self, *args, **kwargs): - assert ( - False - ), "Cannot modify special tokens after patch. Call unpatch_tokenizer first." + assert False, ( + "Cannot modify special tokens after patch. Call unpatch_tokenizer first." + ) def patched_add_tokens(self, new_tokens, special_tokens=False): - assert ( - not special_tokens - ), "Cannot add special tokens after patch. Call unpatch_tokenizer first." + assert not special_tokens, ( + "Cannot add special tokens after patch. Call unpatch_tokenizer first." + ) return tokenizer_cls._original_add_tokens( self, new_tokens, special_tokens=False ) diff --git a/python/sglang/srt/utils/patch_torch.py b/python/sglang/srt/utils/patch_torch.py index 68cc94002..690b122e3 100644 --- a/python/sglang/srt/utils/patch_torch.py +++ b/python/sglang/srt/utils/patch_torch.py @@ -29,9 +29,9 @@ if _is_npu: return npu_reductions._rebuild_npu_tensor_original(*args) def npu_verl_to_sglang(device: int): - assert ( - SGLANG_TP_RANK is not None - ), "SGLANG_TP_RANK is not registered. Please call register_sgl_tp_rank() first." + assert SGLANG_TP_RANK is not None, ( + "SGLANG_TP_RANK is not registered. Please call register_sgl_tp_rank() first." + ) return SGLANG_TP_RANK diff --git a/python/sglang/srt/utils/phase_checker.py b/python/sglang/srt/utils/phase_checker.py index 84c1a9891..5546fa0fc 100644 --- a/python/sglang/srt/utils/phase_checker.py +++ b/python/sglang/srt/utils/phase_checker.py @@ -94,8 +94,7 @@ class SimplePhaseChecker: def _reset_to_idle(self) -> None: self._phase.fill_(self._initial_phase) _host_debug( - f"[SimplePhaseChecker._reset_to_idle] phase reset to " - f"{self._initial_phase}" + f"[SimplePhaseChecker._reset_to_idle] phase reset to {self._initial_phase}" ) def _resolve_caller_tag(self, caller_name: str) -> int: diff --git a/python/sglang/srt/utils/profile_utils.py b/python/sglang/srt/utils/profile_utils.py index 4d1a80c72..41f28a2e3 100644 --- a/python/sglang/srt/utils/profile_utils.py +++ b/python/sglang/srt/utils/profile_utils.py @@ -109,9 +109,9 @@ class ProfileManager: self.detailed_annotations = detailed_annotations # not supported yet assert start_step is None - assert ( - profile_by_stage - ), "only support profile_by_stage=true now" # `false` can be easily supported + assert profile_by_stage, ( + "only support profile_by_stage=true now" + ) # `false` can be easily supported assert not merge_profiles if output_dir is None: diff --git a/python/sglang/srt/utils/rank_consensus_checker.py b/python/sglang/srt/utils/rank_consensus_checker.py index 953782dfc..fb6bd0882 100644 --- a/python/sglang/srt/utils/rank_consensus_checker.py +++ b/python/sglang/srt/utils/rank_consensus_checker.py @@ -75,7 +75,7 @@ def rank_consensus(func=None, *, same_params=None, same_results=None, **kwargs): """ if kwargs: raise TypeError( - f"rank_consensus() got unexpected keyword argument(s): " f"{list(kwargs)}" + f"rank_consensus() got unexpected keyword argument(s): {list(kwargs)}" ) params_selector = _normalize_selector(same_params, "same_params") diff --git a/python/sglang/srt/utils/rpd_utils.py b/python/sglang/srt/utils/rpd_utils.py index bceb5f72e..e0d16c130 100644 --- a/python/sglang/srt/utils/rpd_utils.py +++ b/python/sglang/srt/utils/rpd_utils.py @@ -72,9 +72,9 @@ def rpd_to_chrome_trace( raise Exception("Trace file is empty.") print("Timestamps:") - print(f"\t first: \t{min_time/1000} us") - print(f"\t last: \t{max_time/1000} us") - print(f"\t duration: \t{(max_time-min_time) / 1000000000} seconds") + print(f"\t first: \t{min_time / 1000} us") + print(f"\t last: \t{max_time / 1000} us") + print(f"\t duration: \t{(max_time - min_time) / 1000000000} seconds") start_time = min_time / 1000 end_time = max_time / 1000 @@ -114,7 +114,7 @@ def rpd_to_chrome_trace( ) print("\nFilter: %s" % (rangeStringApi)) - print(f"Output duration: {(end_time-start_time)/1000000} seconds") + print(f"Output duration: {(end_time - start_time) / 1000000} seconds") # Output Ops diff --git a/python/sglang/srt/utils/watchdog.py b/python/sglang/srt/utils/watchdog.py index 7c774e758..70df32b5d 100644 --- a/python/sglang/srt/utils/watchdog.py +++ b/python/sglang/srt/utils/watchdog.py @@ -26,9 +26,9 @@ class Watchdog: test_stuck_time: float = 0, ) -> Watchdog: if watchdog_timeout is None: - assert ( - test_stuck_time == 0 - ), f"stuck tester can be enabled only if soft watchdog is enabled." + assert test_stuck_time == 0, ( + f"stuck tester can be enabled only if soft watchdog is enabled." + ) return _WatchdogNoop() return _WatchdogReal( debug_name=debug_name, diff --git a/python/sglang/srt/utils/weight_checker.py b/python/sglang/srt/utils/weight_checker.py index bb6c3377f..4ab69ddfe 100644 --- a/python/sglang/srt/utils/weight_checker.py +++ b/python/sglang/srt/utils/weight_checker.py @@ -89,9 +89,9 @@ class WeightChecker: (name, param.data.detach().cpu()) for name, param in self._model_state() ] self._snapshot_tensors = dict(named_tensors) - assert len(self._snapshot_tensors) == len( - named_tensors - ), f"should not have duplicated tensor name" + assert len(self._snapshot_tensors) == len(named_tensors), ( + f"should not have duplicated tensor name" + ) def _reset_tensors(self): for name, param in self._model_state(): @@ -198,9 +198,9 @@ def _check_tensors( # skip cos/sin cache which is deterministic from shape and dtype and may have different shapes due to different implementations. continue assert expect_name == actual_name, f"{expect_name=} {actual_name=}" - assert ( - should_compare == actual_should_compare - ), f"{should_compare=} {actual_should_compare=}" + assert should_compare == actual_should_compare, ( + f"{should_compare=} {actual_should_compare=}" + ) name = expect_name try: diff --git a/python/sglang/srt/utils/weight_checker_comparator.py b/python/sglang/srt/utils/weight_checker_comparator.py index 11d387e41..1f0471d91 100644 --- a/python/sglang/srt/utils/weight_checker_comparator.py +++ b/python/sglang/srt/utils/weight_checker_comparator.py @@ -132,9 +132,9 @@ def compare_weights( for (expect_dq, expect_tol), (actual_dq, actual_tol) in zip( expect.iter_chunks(), actual.iter_chunks(), strict=True ): - assert ( - expect_dq.shape == actual_dq.shape - ), f"{expect_dq.shape=} {actual_dq.shape=}" + assert expect_dq.shape == actual_dq.shape, ( + f"{expect_dq.shape=} {actual_dq.shape=}" + ) numel += expect_dq.numel() abs_diff = (actual_dq.float() - expect_dq.float()).abs() if torch.all(abs_diff == 0): diff --git a/python/sglang/srt/weight_cache/daemon.py b/python/sglang/srt/weight_cache/daemon.py index 077236cbe..28101c1eb 100644 --- a/python/sglang/srt/weight_cache/daemon.py +++ b/python/sglang/srt/weight_cache/daemon.py @@ -501,7 +501,7 @@ class WeightCacheDaemon: f.write(f"config={self.config.to_dict()}\n") logger.info( - f"[WeightCacheDaemon gpu={self.gpu_id}] " f"Listening on {self.socket_path}" + f"[WeightCacheDaemon gpu={self.gpu_id}] Listening on {self.socket_path}" ) self._running = True diff --git a/python/sglang/test/accuracy_test_runner.py b/python/sglang/test/accuracy_test_runner.py index 7da5d4289..fb7c5d6f0 100644 --- a/python/sglang/test/accuracy_test_runner.py +++ b/python/sglang/test/accuracy_test_runner.py @@ -189,11 +189,11 @@ def run_accuracy_test( """ base_url = base_url or DEFAULT_URL_FOR_TEST - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Running ACCURACY test for {model.model_path}") print(f" Dataset: {params.dataset}") print(f" Baseline: {params.baseline_accuracy}") - print(f"{'='*60}\n") + print(f"{'=' * 60}\n") success, error, metrics = _run_simple_eval( model=model, diff --git a/python/sglang/test/ascend/disaggregation_utils.py b/python/sglang/test/ascend/disaggregation_utils.py index cb46a9d3b..14e05785c 100644 --- a/python/sglang/test/ascend/disaggregation_utils.py +++ b/python/sglang/test/ascend/disaggregation_utils.py @@ -138,7 +138,7 @@ def get_rdma_devices_args(): if not (base_rdma_group <= gpu_idx < base_rdma_group + 4): warnings.warn( f"GPU index {gpu_idx} is outside expected group " - f"{base_rdma_group}-{base_rdma_group+3}" + f"{base_rdma_group}-{base_rdma_group + 3}" ) # 3. Generate RDMA device names diff --git a/python/sglang/test/ascend/e2e/gen_dataset_fixed_len.py b/python/sglang/test/ascend/e2e/gen_dataset_fixed_len.py index f3bddc594..109a14534 100644 --- a/python/sglang/test/ascend/e2e/gen_dataset_fixed_len.py +++ b/python/sglang/test/ascend/e2e/gen_dataset_fixed_len.py @@ -164,7 +164,7 @@ def generate_custom_dataset( for item in output_data ] print( - f"Token count stats: min={min(token_counts)}, max={max(token_counts)}, avg={sum(token_counts)/len(token_counts):.1f}" + f"Token count stats: min={min(token_counts)}, max={max(token_counts)}, avg={sum(token_counts) / len(token_counts):.1f}" ) return output_data @@ -459,8 +459,8 @@ def generate_random_dataset( } ) - print(f"#Input tokens: {np.sum(input_lens[:len(input_requests)])}") - print(f"#Output tokens: {np.sum(output_lens[:len(input_requests)])}") + print(f"#Input tokens: {np.sum(input_lens[: len(input_requests)])}") + print(f"#Output tokens: {np.sum(output_lens[: len(input_requests)])}") output_dir = os.path.dirname(output_file) if output_dir: diff --git a/python/sglang/test/ascend/e2e/test_npu_accuracy_utils.py b/python/sglang/test/ascend/e2e/test_npu_accuracy_utils.py index 98bf2d76f..5e5375555 100644 --- a/python/sglang/test/ascend/e2e/test_npu_accuracy_utils.py +++ b/python/sglang/test/ascend/e2e/test_npu_accuracy_utils.py @@ -196,8 +196,7 @@ def run_evalscope( process.wait() logger.info( - f"run_evalscope finished: pid={process.pid} " - f"returncode={process.returncode}" + f"run_evalscope finished: pid={process.pid} returncode={process.returncode}" ) kill_process_group(process) @@ -641,7 +640,9 @@ class TestNpuAccuracyMultiNodePdSepTestCaseBase(CustomTestCase): cls.role = ( "router" if "router" in cls.hostname - else "prefill" if "prefill" in cls.hostname else "decode" + else "prefill" + if "prefill" in cls.hostname + else "decode" ) logger.info(f"Init {cls.host} {cls.role=}!") diff --git a/python/sglang/test/ascend/e2e/test_npu_multi_node_utils.py b/python/sglang/test/ascend/e2e/test_npu_multi_node_utils.py index c27b4c581..5387aa133 100644 --- a/python/sglang/test/ascend/e2e/test_npu_multi_node_utils.py +++ b/python/sglang/test/ascend/e2e/test_npu_multi_node_utils.py @@ -970,7 +970,7 @@ class TestNpuMultiNodePdMixTestCaseBase(CustomTestCase): self.assertGreaterEqual( metrics["accuracy"], expect_accuracy, - f'Accuracy is {str(metrics["accuracy"])}, is lower than {expect_accuracy}', + f"Accuracy is {str(metrics['accuracy'])}, is lower than {expect_accuracy}", ) @@ -988,7 +988,9 @@ class TestNpuMultiNodePdSepTestCaseBase(CustomTestCase): cls.role = ( "router" if "router" in cls.hostname - else "prefill" if "prefill" in cls.hostname else "decode" + else "prefill" + if "prefill" in cls.hostname + else "decode" ) logger.info(f"Init {cls.host} {cls.role=}!") cls.sglang_thread = None @@ -1078,7 +1080,7 @@ class TestNpuMultiNodePdSepTestCaseBase(CustomTestCase): self.assertGreaterEqual( metrics["accuracy"], expect_accuracy, - f'Accuracy is {str(metrics["accuracy"])}, is lower than {expect_accuracy}', + f"Accuracy is {str(metrics['accuracy'])}, is lower than {expect_accuracy}", ) diff --git a/python/sglang/test/ascend/e2e/test_npu_performance_utils.py b/python/sglang/test/ascend/e2e/test_npu_performance_utils.py index 8d8e5a51e..86bb0d9bc 100644 --- a/python/sglang/test/ascend/e2e/test_npu_performance_utils.py +++ b/python/sglang/test/ascend/e2e/test_npu_performance_utils.py @@ -721,7 +721,7 @@ def run_aisbench( else: logger.warning("Could not extract mean_tpot from output") logger.error( - f"Simplified output snippet around TPOT: {simplified_output[simplified_output.find('TPOT')-20:simplified_output.find('TPOT')+50] if 'TPOT' in simplified_output else 'TPOT not found'}" + f"Simplified output snippet around TPOT: {simplified_output[simplified_output.find('TPOT') - 20 : simplified_output.find('TPOT') + 50] if 'TPOT' in simplified_output else 'TPOT not found'}" ) tps_matches = re.findall( @@ -751,7 +751,7 @@ def run_aisbench( else: logger.warning("Could not extract total_tps from output") logger.warning( - f"Simplified output snippet around Output Token Throughput: {simplified_output[simplified_output.find('Output')-20:simplified_output.find('Output')+100] if 'Output' in simplified_output else 'Output not found'}" + f"Simplified output snippet around Output Token Throughput: {simplified_output[simplified_output.find('Output') - 20 : simplified_output.find('Output') + 100] if 'Output' in simplified_output else 'Output not found'}" ) ttft_match = re.search(r"TTFT\s+total\s+([\d.]+)\s+ms", simplified_output) @@ -761,7 +761,7 @@ def run_aisbench( else: logger.warning("Could not extract mean_ttft from output") logger.warning( - f"Simplified output snippet around TTFT: {simplified_output[simplified_output.find('TTFT')-20:simplified_output.find('TTFT')+50] if 'TTFT' in simplified_output else 'TTFT not found'}" + f"Simplified output snippet around TTFT: {simplified_output[simplified_output.find('TTFT') - 20 : simplified_output.find('TTFT') + 50] if 'TTFT' in simplified_output else 'TTFT not found'}" ) e2el_match = re.search(r"E2EL\s+total\s+([\d.]+)\s+ms", simplified_output) @@ -771,7 +771,7 @@ def run_aisbench( else: logger.warning("Could not extract mean_e2e_latency from output") logger.warning( - f"Simplified output snippet around E2EL: {simplified_output[simplified_output.find('E2EL')-20:simplified_output.find('E2EL')+50] if 'E2EL' in simplified_output else 'E2EL not found'}" + f"Simplified output snippet around E2EL: {simplified_output[simplified_output.find('E2EL') - 20 : simplified_output.find('E2EL') + 50] if 'E2EL' in simplified_output else 'E2EL not found'}" ) concurrency_match = re.search( @@ -783,7 +783,7 @@ def run_aisbench( else: logger.warning("Could not extract concurrency from output") logger.warning( - f"Simplified output snippet around Concurrency: {simplified_output[simplified_output.find('Concurrency')-20:simplified_output.find('Concurrency')+50] if 'Concurrency' in simplified_output else 'Concurrency not found'}" + f"Simplified output snippet around Concurrency: {simplified_output[simplified_output.find('Concurrency') - 20 : simplified_output.find('Concurrency') + 50] if 'Concurrency' in simplified_output else 'Concurrency not found'}" ) max_concurrency_match = re.search( @@ -795,7 +795,7 @@ def run_aisbench( else: logger.warning("Could not extract max_concurrency from output") logger.warning( - f"Simplified output snippet around Max Concurrency: {simplified_output[simplified_output.find('Max Concurrency')-20:simplified_output.find('Max Concurrency')+50] if 'Max Concurrency' in simplified_output else 'Max Concurrency not found'}" + f"Simplified output snippet around Max Concurrency: {simplified_output[simplified_output.find('Max Concurrency') - 20 : simplified_output.find('Max Concurrency') + 50] if 'Max Concurrency' in simplified_output else 'Max Concurrency not found'}" ) req_throughput_match = re.search( @@ -810,7 +810,7 @@ def run_aisbench( else: logger.warning("Could not extract request_throughput from output") logger.warning( - f"Simplified output snippet around Request Throughput: {simplified_output[simplified_output.find('Request')-20:simplified_output.find('Request')+50] if 'Request' in simplified_output else 'Request not found'}" + f"Simplified output snippet around Request Throughput: {simplified_output[simplified_output.find('Request') - 20 : simplified_output.find('Request') + 50] if 'Request' in simplified_output else 'Request not found'}" ) total_requests_match = re.search( @@ -822,7 +822,7 @@ def run_aisbench( else: logger.warning("Could not extract total_requests from output") logger.warning( - f"Simplified output snippet around Total Requests: {simplified_output[simplified_output.find('Total Requests')-20:simplified_output.find('Total Requests')+50] if 'Total Requests' in simplified_output else 'Total Requests not found'}" + f"Simplified output snippet around Total Requests: {simplified_output[simplified_output.find('Total Requests') - 20 : simplified_output.find('Total Requests') + 50] if 'Total Requests' in simplified_output else 'Total Requests not found'}" ) failed_requests_match = re.search( @@ -834,7 +834,7 @@ def run_aisbench( else: logger.warning("Could not extract failed_requests from output") logger.warning( - f"Simplified output snippet around Failed Requests: {simplified_output[simplified_output.find('Failed Requests')-20:simplified_output.find('Failed Requests')+50] if 'Failed Requests' in simplified_output else 'Failed Requests not found'}" + f"Simplified output snippet around Failed Requests: {simplified_output[simplified_output.find('Failed Requests') - 20 : simplified_output.find('Failed Requests') + 50] if 'Failed Requests' in simplified_output else 'Failed Requests not found'}" ) logger.info(f"All extracted metrics: {metrics}") @@ -1365,7 +1365,9 @@ class TestNpuPerfMultiNodePdSepTestCaseBase(CustomTestCase): cls.role = ( "router" if "router" in cls.hostname - else "prefill" if "prefill" in cls.hostname else "decode" + else "prefill" + if "prefill" in cls.hostname + else "decode" ) logger.info(f"Init {cls.host} {cls.role=}!") diff --git a/python/sglang/test/ascend/gsm8k_ascend_mixin.py b/python/sglang/test/ascend/gsm8k_ascend_mixin.py index fcf012ea3..0f6dfef73 100644 --- a/python/sglang/test/ascend/gsm8k_ascend_mixin.py +++ b/python/sglang/test/ascend/gsm8k_ascend_mixin.py @@ -102,12 +102,12 @@ class GSM8KAscendMixin(ABC): self.assertGreaterEqual( metrics["score"], accuracy_threshold, - f'Accuracy of {self.model} is {str(metrics["score"])}, is lower than {accuracy_threshold}', + f"Accuracy of {self.model} is {str(metrics['score'])}, is lower than {accuracy_threshold}", ) self.assertGreaterEqual( metrics["output_throughput"], output_throughput_threshold, - f'Output throughput of {self.model} is {str(metrics["output_throughput"])}, is lower than {output_throughput_threshold}', + f"Output throughput of {self.model} is {str(metrics['output_throughput'])}, is lower than {output_throughput_threshold}", ) except Exception as e: model_metrics["error"] = e diff --git a/python/sglang/test/ascend/test_embedding_base.py b/python/sglang/test/ascend/test_embedding_base.py index 2857b8317..a751e94be 100644 --- a/python/sglang/test/ascend/test_embedding_base.py +++ b/python/sglang/test/ascend/test_embedding_base.py @@ -87,9 +87,9 @@ class BaseEmbeddingTest(ABC): print("similarity diff", abs(similarity - 1)) if len(prompts[i]) <= 1000: - assert torch.all( - abs(similarity - 1) < prefill_tolerance - ), "embeddings are not all close" + assert torch.all(abs(similarity - 1) < prefill_tolerance), ( + "embeddings are not all close" + ) def test_prefill_logits(self): """Main test method to run for all models and dtypes""" diff --git a/python/sglang/test/ascend/test_mmlu.py b/python/sglang/test/ascend/test_mmlu.py index 0b90f9c4f..84f8a25f5 100644 --- a/python/sglang/test/ascend/test_mmlu.py +++ b/python/sglang/test/ascend/test_mmlu.py @@ -6,7 +6,6 @@ from sglang.test.run_eval import run_eval class TestMMLU: - mmlu_num_examples = 128 def test_mmlu(self): diff --git a/python/sglang/test/cache_consistency_jitter.py b/python/sglang/test/cache_consistency_jitter.py index ee5d1e1fc..4289c93c4 100644 --- a/python/sglang/test/cache_consistency_jitter.py +++ b/python/sglang/test/cache_consistency_jitter.py @@ -199,9 +199,7 @@ def get_jitter_engine( port = portpicker.pick_unused_port_range( # pyright: ignore[reportAttributeAccessIssue] DP_ATTENTION_HANDSHAKE_PORT_DELTA + 1 - )[ - 0 - ] + )[0] engine_kwargs["dist_init_addr"] = f"127.0.0.1:{port}" scheduler_process = ( @@ -307,9 +305,9 @@ def _record_response( out_ids = [int(t) for t in out["output_ids"]] out_lps = meta["output_token_logprobs"] out_top = meta["output_top_logprobs"] - assert ( - len(out_ids) == req.max_new_tokens - ), f"{req.label}: got {len(out_ids)} output tokens, expected {req.max_new_tokens}" + assert len(out_ids) == req.max_new_tokens, ( + f"{req.label}: got {len(out_ids)} output tokens, expected {req.max_new_tokens}" + ) expected = out_ids if baseline is None else baseline[req.cut :] for m, (tid, lp_entry, top) in enumerate( zip(out_ids, out_lps, out_top, strict=True) diff --git a/python/sglang/test/ci/ci_stress_utils.py b/python/sglang/test/ci/ci_stress_utils.py index 942719da5..a693b6921 100644 --- a/python/sglang/test/ci/ci_stress_utils.py +++ b/python/sglang/test/ci/ci_stress_utils.py @@ -149,12 +149,12 @@ class StressTestRunner: Returns: True if successful, False otherwise """ - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Starting stress test for: {model_path}") print(f"Input length: {random_input_len}") print(f"Output length: {random_output_len}") print(f"Num prompts: {self.num_prompts}") - print(f"{'='*60}\n") + print(f"{'=' * 60}\n") # Launch server process = popen_launch_server( diff --git a/python/sglang/test/ci/ci_utils.py b/python/sglang/test/ci/ci_utils.py index 8b87d8b84..787a8e3b2 100644 --- a/python/sglang/test/ci/ci_utils.py +++ b/python/sglang/test/ci/ci_utils.py @@ -434,11 +434,11 @@ def run_unittest_files( logger.info(f"Fail. Time elapsed: {elapsed_total:.2f}s") # Print summary - logger.info(f"\n{'='*60}") + logger.info(f"\n{'=' * 60}") logger.info(f"Test Summary: {len(passed_tests)}/{len(files)} passed") if enable_retry and retried_tests: logger.info(f"Retries: {len(retried_tests)} test(s) were retried") - logger.info(f"{'='*60}") + logger.info(f"{'=' * 60}") if passed_tests: logger.info("✓ PASSED:") for test in passed_tests: @@ -451,7 +451,7 @@ def run_unittest_files( logger.info("\n↻ RETRIED:") for test, attempts, result in retried_tests: logger.info(f" {test} ({attempts} attempts, {result})") - logger.info(f"{'='*60}\n") + logger.info(f"{'=' * 60}\n") # Machine-readable timings block for downstream scrapers/dashboards. # One JSON object per executed file (post-retry: only the latest diff --git a/python/sglang/test/cpu_test_utils.py b/python/sglang/test/cpu_test_utils.py index f75279687..816f2ee25 100644 --- a/python/sglang/test/cpu_test_utils.py +++ b/python/sglang/test/cpu_test_utils.py @@ -441,9 +441,9 @@ class MXFP4QuantizeUtil: new_data = ( right_side.clone() << 4 ) # Put odd indices (higher addresses) in high bits - new_data[ - ..., : left_side.shape[-1] - ] += left_side # Put even indices in low bits + new_data[..., : left_side.shape[-1]] += ( + left_side # Put even indices in low bits + ) return new_data original_shape = input.shape diff --git a/python/sglang/test/kernels/kv_canary/_canary_helpers.py b/python/sglang/test/kernels/kv_canary/_canary_helpers.py index 99bacfb51..6d14f7cee 100644 --- a/python/sglang/test/kernels/kv_canary/_canary_helpers.py +++ b/python/sglang/test/kernels/kv_canary/_canary_helpers.py @@ -324,9 +324,9 @@ def assert_canary_state_equal( *, log_a: FakeViolationLog, log_b: FakeViolationLog ) -> None: for name in ("ring", "write_index", "slot_run_counter", "kernel_run_counter"): - assert torch.equal( - getattr(log_a, name), getattr(log_b, name) - ), f"{name} diverged (CUDA vs ref)" + assert torch.equal(getattr(log_a, name), getattr(log_b, name)), ( + f"{name} diverged (CUDA vs ref)" + ) def assert_canary_buf_equal(*, buf_a: torch.Tensor, buf_b: torch.Tensor) -> None: @@ -334,11 +334,9 @@ def assert_canary_buf_equal(*, buf_a: torch.Tensor, buf_b: torch.Tensor) -> None def assert_only_bits_set(fail_bits: int, expected_bits: int) -> None: - assert ( - fail_bits & expected_bits - ) == expected_bits, ( + assert (fail_bits & expected_bits) == expected_bits, ( f"missing expected bits: expected {expected_bits:#b} got {fail_bits:#b}" ) - assert ( - fail_bits & ~expected_bits - ) == 0, f"unexpected extra bits: got {fail_bits:#b} extras {fail_bits & ~expected_bits:#b}" + assert (fail_bits & ~expected_bits) == 0, ( + f"unexpected extra bits: got {fail_bits:#b} extras {fail_bits & ~expected_bits:#b}" + ) diff --git a/python/sglang/test/kernels/kv_canary/_differential.py b/python/sglang/test/kernels/kv_canary/_differential.py index b122a2e12..dac5cb7b7 100644 --- a/python/sglang/test/kernels/kv_canary/_differential.py +++ b/python/sglang/test/kernels/kv_canary/_differential.py @@ -130,18 +130,18 @@ def _assert_plans_byte_equal( else int(triton_verify.verify_num_valid[0].item()) ) n_verify_ref = int(ref_verify.verify_num_valid[0].item()) - assert ( - n_verify == n_verify_ref - ), f"verify_num_valid diverged: triton={n_verify} ref={n_verify_ref}" + assert n_verify == n_verify_ref, ( + f"verify_num_valid diverged: triton={n_verify} ref={n_verify_ref}" + ) # When total_verify > VERIFY_CAPACITY the offsets kernel clears verify_enable and # plan_entries skips its scatter — leaving verify_slot_indices/positions/prev_slot_indices # as whatever the (torch.empty) allocation contained. Skip the byte-equal probe in that # case; verify_num_valid being clamped + verify_enable=0 is the contract here. triton_enable = int(triton_verify.enable[0].item()) ref_enable = int(ref_verify.enable[0].item()) - assert ( - triton_enable == ref_enable - ), f"verify_enable diverged: triton={triton_enable} ref={ref_enable}" + assert triton_enable == ref_enable, ( + f"verify_enable diverged: triton={triton_enable} ref={ref_enable}" + ) if n_verify > 0 and triton_enable != 0: assert torch.equal( triton_verify.verify_slot_indices[:n_verify], @@ -166,9 +166,9 @@ def _assert_plans_byte_equal( else int(triton_write.write_num_valid_reqs[0].item()) ) n_write_ref = int(ref_write.write_num_valid_reqs[0].item()) - assert ( - n_write == n_write_ref - ), f"write_num_valid_reqs diverged: triton={n_write} ref={n_write_ref}" + assert n_write == n_write_ref, ( + f"write_num_valid_reqs diverged: triton={n_write} ref={n_write_ref}" + ) assert torch.equal( triton_write.write_offsets[: n_write + 1], ref_write.write_offsets[: n_write + 1], @@ -350,7 +350,8 @@ def _yield_simpler(inputs: Any) -> Iterator[tuple[str, Any]]: Plan / Verify / Write fuzz failures uniformly. """ fields = { - f: getattr(inputs, f) for f in inputs.__dataclass_fields__ # type: ignore[attr-defined] + f: getattr(inputs, f) + for f in inputs.__dataclass_fields__ # type: ignore[attr-defined] } def emit(label: str, **overrides: Any) -> Iterator[tuple[str, Any]]: diff --git a/python/sglang/test/kernels/kv_canary/_invariants.py b/python/sglang/test/kernels/kv_canary/_invariants.py index 588a242ba..3a619cce6 100644 --- a/python/sglang/test/kernels/kv_canary/_invariants.py +++ b/python/sglang/test/kernels/kv_canary/_invariants.py @@ -76,9 +76,9 @@ class PlanInvariants: raise AssertionError(f"write_num_valid_reqs negative: {n_active}") offsets = write_plan.write_offsets[: n_active + 1].detach().cpu().tolist() for i in range(len(offsets) - 1): - assert ( - offsets[i] <= offsets[i + 1] - ), f"write_offsets non-monotone at {i}: {offsets[i]} > {offsets[i + 1]}" + assert offsets[i] <= offsets[i + 1], ( + f"write_offsets non-monotone at {i}: {offsets[i]} > {offsets[i + 1]}" + ) @staticmethod def _assert_write_offsets_total_matches_active_extend_sum( @@ -92,9 +92,9 @@ class PlanInvariants: rpi_cpu = req_pool_indices.detach().cpu().tolist() ext_cpu = extend_seq_lens.detach().cpu().tolist() expected_total = sum(ext for rpi, ext in zip(rpi_cpu, ext_cpu) if rpi != 0) - assert ( - total == expected_total - ), f"write_offsets total {total} != active extend sum {expected_total}" + assert total == expected_total, ( + f"write_offsets total {total} != active extend sum {expected_total}" + ) @staticmethod def _assert_extras_land_at_tail( @@ -111,9 +111,9 @@ class PlanInvariants: tail_start = derived_verify_count tail_end = derived_verify_count + extras_count n_valid = int(verify_plan.verify_num_valid[0].item()) - assert ( - tail_end <= n_valid - ), f"extras tail {tail_end} exceeds verify_num_valid {n_valid}" + assert tail_end <= n_valid, ( + f"extras tail {tail_end} exceeds verify_num_valid {n_valid}" + ) plan_slots = verify_plan.verify_slot_indices[tail_start:tail_end] plan_positions = verify_plan.verify_expected_positions[tail_start:tail_end] plan_prevs = verify_plan.verify_prev_slot_indices[tail_start:tail_end] @@ -136,9 +136,9 @@ class PlanInvariants: ) for r in range(min(n_active, len(rpi_cpu))): if rpi_cpu[r] == 0: - assert ( - seeds_cpu[r] == -1 - ), f"padding row {r} has seed {seeds_cpu[r]} != -1" + assert seeds_cpu[r] == -1, ( + f"padding row {r} has seed {seeds_cpu[r]} != -1" + ) @staticmethod def _assert_prev_slot_minus_one_iff_chain_head( @@ -163,14 +163,14 @@ class PlanInvariants: ) for i, (pos, prev) in enumerate(zip(positions_cpu, prevs_cpu)): if pos == 0: - assert ( - prev == -1 - ), f"entry {i} at position 0 must have prev=-1, got {prev}" + assert prev == -1, ( + f"entry {i} at position 0 must have prev=-1, got {prev}" + ) else: if swa_window_size == 0: - assert ( - prev != -1 - ), f"FULL entry {i} at position {pos} must have prev != -1, got {prev}" + assert prev != -1, ( + f"FULL entry {i} at position {pos} must have prev != -1, got {prev}" + ) @staticmethod def _assert_verify_num_valid_equals_derived_plus_extras( @@ -251,9 +251,9 @@ class VerifyInvariants: canary_buf_before: torch.Tensor, canary_buf_after: torch.Tensor, ) -> None: - assert torch.equal( - canary_buf_before, canary_buf_after - ), "verify kernel mutated canary_buf (must be read-only)" + assert torch.equal(canary_buf_before, canary_buf_after), ( + "verify kernel mutated canary_buf (must be read-only)" + ) @staticmethod def _assert_violation_count_le_active_entries( @@ -266,9 +266,9 @@ class VerifyInvariants: log_before.write_index[0].item() ) n_active = int(plan.verify_num_valid[0].item()) - assert ( - 0 <= delta <= n_active - ), f"violation_write_index delta {delta} out of [0, {n_active}]" + assert 0 <= delta <= n_active, ( + f"violation_write_index delta {delta} out of [0, {n_active}]" + ) @staticmethod def _assert_violation_rows_have_valid_slot_and_kernel_kind( @@ -292,13 +292,13 @@ class VerifyInvariants: rows = log_after.ring[visible_start:visible_end].detach().cpu() for i in range(rows.shape[0]): kind = int(rows[i, consts.VIOLATION_FIELD_KERNEL_KIND].item()) - assert kind == int( - kernel_kind - ), f"row {visible_start + i} kernel_kind {kind} != expected {int(kernel_kind)}" + assert kind == int(kernel_kind), ( + f"row {visible_start + i} kernel_kind {kind} != expected {int(kernel_kind)}" + ) slot = int(rows[i, consts.VIOLATION_FIELD_SLOT_IDX].item()) - assert ( - slot in plan_slots - ), f"row {visible_start + i} slot {slot} not in plan_slots" + assert slot in plan_slots, ( + f"row {visible_start + i} slot {slot} not in plan_slots" + ) @staticmethod def _assert_slot_run_counter_incremented_by_active_entries( @@ -311,9 +311,9 @@ class VerifyInvariants: delta = int(log_after.slot_run_counter[0].item()) - int( log_before.slot_run_counter[0].item() ) - assert ( - delta == n_active - ), f"slot_run_counter delta {delta} != active entries {n_active}" + assert delta == n_active, ( + f"slot_run_counter delta {delta} != active entries {n_active}" + ) @staticmethod def _assert_kernel_run_counter_incremented_by_one( @@ -401,12 +401,12 @@ class WriteInvariants: continue stored_token = int(view[slot, 0].item()) stored_position = int(view[slot, 1].item()) - assert ( - stored_token == tokens_cpu[i] - ), f"slot {slot}: stored token {stored_token} != input {tokens_cpu[i]}" - assert ( - stored_position == pos_cpu[i] - ), f"slot {slot}: stored position {stored_position} != input {pos_cpu[i]}" + assert stored_token == tokens_cpu[i], ( + f"slot {slot}: stored token {stored_token} != input {tokens_cpu[i]}" + ) + assert stored_position == pos_cpu[i], ( + f"slot {slot}: stored position {stored_position} != input {pos_cpu[i]}" + ) @staticmethod def _assert_slot_minus_one_skipped( @@ -428,9 +428,9 @@ class WriteInvariants: for slot in range(num_slots): if slot in written_slots: continue - assert torch.equal( - view_before[slot], view_after[slot] - ), f"slot {slot} not in out_cache_loc but canary_buf changed" + assert torch.equal(view_before[slot], view_after[slot]), ( + f"slot {slot} not in out_cache_loc but canary_buf changed" + ) @staticmethod def _assert_pseudo_violation_only_on_mismatch( @@ -449,9 +449,9 @@ class WriteInvariants: log_before.write_index[0].item() ) if not enable_write_verify_inputs: - assert ( - delta == 0 - ), f"enable_write_verify_inputs=OFF must produce no violations, got {delta}" + assert delta == 0, ( + f"enable_write_verify_inputs=OFF must produce no violations, got {delta}" + ) return if expected_input_tokens is None or expected_input_positions is None: return @@ -472,13 +472,13 @@ class WriteInvariants: ) no_mismatch = mismatch_entries == 0 if no_mismatch: - assert ( - delta == 0 - ), f"enable_write_verify_inputs=ON with no mismatch produced {delta} violations" + assert delta == 0, ( + f"enable_write_verify_inputs=ON with no mismatch produced {delta} violations" + ) else: - assert ( - delta == mismatch_entries - ), f"write input mismatch count {mismatch_entries} produced {delta} violations" + assert delta == mismatch_entries, ( + f"write input mismatch count {mismatch_entries} produced {delta} violations" + ) @staticmethod def _assert_write_slot_run_counter_incremented( diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py index 7cc4c340a..057e8a3db 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py @@ -1456,8 +1456,7 @@ def run_dsa_sparse_cuda_graph_decode_impl_variant_case( ) if not case.forward_mode.is_decode(): raise ValueError( - "run_dsa_sparse_cuda_graph_decode_impl_variant_case expects a " - "DECODE case." + "run_dsa_sparse_cuda_graph_decode_impl_variant_case expects a DECODE case." ) from ..runner_modes.cuda_graph_decode_runner import ( run_dsa_sparse_cuda_graph_decode_case, diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py index 16e8714d0..8202ec398 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py @@ -1015,8 +1015,7 @@ def make_dsv4_padded_replay_inputs( pad_token_count = case.num_input_tokens - base_inputs["input_hidden"].shape[0] if pad_token_count < 0: raise ValueError( - f"replay input shrink not supported: {pad_token_count=}; " - f"case={case.name}" + f"replay input shrink not supported: {pad_token_count=}; case={case.name}" ) if pad_token_count == 0: padded_input_hidden = base_inputs["input_hidden"] @@ -1446,12 +1445,12 @@ def _seed_c4_sparse_prefill_indices( max_len = int(lens.max().item()) pool = fixture.runner.token_to_kv_pool c4_page_size = pool.get_extra_key_page_size(layer_id=0) - assert max_len <= min( - num_entries, c4_page_size - ), f"case attends {max_len} c4 entries; only {min(num_entries, c4_page_size)} populated" - assert ( - md.page_table[:, 0] == 0 - ).all(), "sparse seeding requires the raw==physical identity (first page 0)" + assert max_len <= min(num_entries, c4_page_size), ( + f"case attends {max_len} c4 entries; only {min(num_entries, c4_page_size)} populated" + ) + assert (md.page_table[:, 0] == 0).all(), ( + "sparse seeding requires the raw==physical identity (first page 0)" + ) seq = ( torch.arange(width, dtype=raw_indices.dtype, device=raw_indices.device) .unsqueeze(0) @@ -1487,9 +1486,9 @@ def run_dsv4_target_verify_attention_case( "DSV4 target_verify is chain-only — `deepseek_v4_backend.py:369` " "asserts `self.topk in [0, 1]`. Pass topk=1." ) - assert ( - case.forward_mode.is_target_verify() - ), f"run_dsv4_target_verify_attention_case requires TARGET_VERIFY case; got {case.forward_mode}" + assert case.forward_mode.is_target_verify(), ( + f"run_dsv4_target_verify_attention_case requires TARGET_VERIFY case; got {case.forward_mode}" + ) # Lazy import to avoid cycles (runner_modes imports attention_methods). from sglang.test.kits.attention_unittest.runner_modes.speculative_target_verify_runner import ( _make_eagle_verify_input, @@ -1560,9 +1559,9 @@ def run_dsv4_draft_extend_attention_case( "`deepseek_v4_backend.py:636-663` and the 'Production-Unsupported' " "section in dsv4/README.md." ) - assert ( - case.forward_mode.is_draft_extend_v2() - ), f"run_dsv4_draft_extend_attention_case requires DRAFT_EXTEND; got {case.forward_mode}" + assert case.forward_mode.is_draft_extend_v2(), ( + f"run_dsv4_draft_extend_attention_case requires DRAFT_EXTEND; got {case.forward_mode}" + ) from sglang.test.kits.attention_unittest.runner_modes.speculative_draft_extend_runner import ( _make_eagle_draft_extend_input, ) @@ -1631,11 +1630,13 @@ def run_dsv4_compress_attention_case( assert case.compress_ratio in ( 4, 128, - ), f"DSV4 compact runner requires compress_ratio in (4, 128); got {case.compress_ratio}" + ), ( + f"DSV4 compact runner requires compress_ratio in (4, 128); got {case.compress_ratio}" + ) if sparse_prefill: - assert ( - case.forward_mode.is_extend_without_speculative() - ), f"sparse prefill only serves extend; got {case.forward_mode}" + assert case.forward_mode.is_extend_without_speculative(), ( + f"sparse prefill only serves extend; got {case.forward_mode}" + ) fixture = build_dsv4_attention_fixture( testcase, case, diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py index d37693b0c..bc4f6003a 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py +++ b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py @@ -567,8 +567,9 @@ def _run_eagle_draft_extend_eager( model_runner = ( worker.model_runner if hasattr(worker, "model_runner") else worker.draft_runner ) - with torch.no_grad(), forward_context( - ForwardContext(attn_backend=worker.draft_extend_attn_backend) + with ( + torch.no_grad(), + forward_context(ForwardContext(attn_backend=worker.draft_extend_attn_backend)), ): worker.draft_extend_attn_backend.init_forward_metadata(batch) ret = model_runner.model.forward( diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py index cbebd85c0..79ba07b03 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py +++ b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py @@ -606,9 +606,9 @@ class _DenseEagleDraftForward: ) def __call__(self, forward_batch: ForwardBatch): - assert ( - forward_batch.forward_metadata_ready - ), "draft-loop forward reached the runner without a pre-planned batch" + assert forward_batch.forward_metadata_ready, ( + "draft-loop forward reached the runner without a pre-planned batch" + ) spec_info = forward_batch.spec_info hidden_states = spec_info.hidden_states if hidden_states is None: @@ -645,9 +645,9 @@ class _FrozenKVMTPDenseDraftForward: ) def __call__(self, forward_batch: ForwardBatch): - assert ( - forward_batch.forward_metadata_ready - ), "draft-loop forward reached the runner without a pre-planned batch" + assert forward_batch.forward_metadata_ready, ( + "draft-loop forward reached the runner without a pre-planned batch" + ) spec_info = forward_batch.spec_info hidden_states = spec_info.hidden_states if hidden_states is None: @@ -1034,9 +1034,9 @@ class _MLAEagleDraftForward: ) def __call__(self, forward_batch: ForwardBatch): - assert ( - forward_batch.forward_metadata_ready - ), "draft-loop forward reached the runner without a pre-planned batch" + assert forward_batch.forward_metadata_ready, ( + "draft-loop forward reached the runner without a pre-planned batch" + ) spec_info = forward_batch.spec_info hidden_states = spec_info.hidden_states if hidden_states is None: @@ -1292,9 +1292,9 @@ class _DSV4EagleDraftForward: ) def __call__(self, forward_batch: ForwardBatch): - assert ( - forward_batch.forward_metadata_ready - ), "draft-loop forward reached the runner without a pre-planned batch" + assert forward_batch.forward_metadata_ready, ( + "draft-loop forward reached the runner without a pre-planned batch" + ) spec_info = forward_batch.spec_info hidden_states = spec_info.hidden_states if hidden_states is None: @@ -1603,9 +1603,9 @@ class _DSAEagleDraftForward: ) def __call__(self, forward_batch: ForwardBatch): - assert ( - forward_batch.forward_metadata_ready - ), "draft-loop forward reached the runner without a pre-planned batch" + assert forward_batch.forward_metadata_ready, ( + "draft-loop forward reached the runner without a pre-planned batch" + ) spec_info = forward_batch.spec_info hidden_states = spec_info.hidden_states if hidden_states is None: diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_target_verify_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_target_verify_runner.py index 6f45af8a3..b7d7d9569 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_target_verify_runner.py +++ b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_target_verify_runner.py @@ -493,8 +493,8 @@ def _run_spec_verify_cuda_graph_case( make_capture_case=lambda base, name, capture_prefix_len, bs: ( make_case_with_prefix_lens(base, name, (capture_prefix_len,) * bs) ), - make_replay_case=lambda base, name, pad_prefix_lens: ( - make_case_with_prefix_lens(base, name, base.prefix_lens + pad_prefix_lens) + make_replay_case=lambda base, name, pad_prefix_lens: make_case_with_prefix_lens( + base, name, base.prefix_lens + pad_prefix_lens ), make_forward_batch=make_forward_batch, fixture_inputs=fixture_inputs, @@ -740,8 +740,8 @@ def run_gdn_eagle_verify_cuda_graph_case( make_forward_batch=_make_gdn_forward_batch, fixture_inputs=gdn_fixture_inputs, make_capture_inputs=make_gdn_random_inputs, - make_replay_inputs=lambda _case, fixture, *_args, **_kwargs: ( - gdn_fixture_inputs(fixture) + make_replay_inputs=lambda _case, fixture, *_args, **_kwargs: gdn_fixture_inputs( + fixture ), prepare_batch=lambda spec_case, batch: _prepare_gdn_verify_batch( spec_case, @@ -928,9 +928,9 @@ def run_dsv4_eagle_verify_cuda_graph_case( ) num_draft_tokens = case.extend_lens[0] if case.extend_lens else 0 - assert ( - num_draft_tokens > 0 - ), "DSV4 verify cases must set `extend_lens=(num_draft, ...)`." + assert num_draft_tokens > 0, ( + "DSV4 verify cases must set `extend_lens=(num_draft, ...)`." + ) def _prepare_dsv4_verify_batch(spec_case, batch): _prepare_target_verify_batch(batch, spec_case, device) @@ -1080,8 +1080,8 @@ def run_kda_eagle_verify_cuda_graph_case( make_forward_batch=_make_kda_forward_batch, fixture_inputs=kda_fixture_inputs, make_capture_inputs=make_kda_random_inputs, - make_replay_inputs=lambda _case, fixture, *_args, **_kwargs: ( - kda_fixture_inputs(fixture) + make_replay_inputs=lambda _case, fixture, *_args, **_kwargs: kda_fixture_inputs( + fixture ), prepare_batch=lambda spec_case, batch: _prepare_kda_verify_batch( spec_case, diff --git a/python/sglang/test/kits/ebnf_constrained_kit.py b/python/sglang/test/kits/ebnf_constrained_kit.py index 292faf937..e4ecf21c2 100644 --- a/python/sglang/test/kits/ebnf_constrained_kit.py +++ b/python/sglang/test/kits/ebnf_constrained_kit.py @@ -4,7 +4,6 @@ import requests class EBNFConstrainedMixin: - ebnf_grammar = 'root ::= "test"' # Default grammar def _run_decode_ebnf( diff --git a/python/sglang/test/kits/eval_accuracy_kit.py b/python/sglang/test/kits/eval_accuracy_kit.py index 58b167200..034927578 100644 --- a/python/sglang/test/kits/eval_accuracy_kit.py +++ b/python/sglang/test/kits/eval_accuracy_kit.py @@ -60,9 +60,9 @@ def _run_accuracy_eval( ``None``, so the common case stays identical to ``run_eval``'s defaults. Returns the metrics dict. """ - assert ( - score_threshold == score_threshold - ), f"{type(test_case).__name__} must set the {eval_name} score threshold" + assert score_threshold == score_threshold, ( + f"{type(test_case).__name__} must set the {eval_name} score threshold" + ) model = eval_overrides.pop("model", getattr(test_case, "model", None)) kwargs = dict( @@ -113,9 +113,9 @@ def _run_sgl_eval( True}`` so the server separates reasoning from the final answer. Skips the test if sgl-eval is not installed. Returns the RunResult. """ - assert ( - score_threshold == score_threshold - ), f"{type(test_case).__name__} must set the {eval_name} score threshold" + assert score_threshold == score_threshold, ( + f"{type(test_case).__name__} must set the {eval_name} score threshold" + ) try: from sgl_eval.registry import get as get_eval_spec @@ -299,7 +299,7 @@ class MMMUProMixin: def test_mmmu_pro(self): assert self.mmmu_pro_load_preset_from_model_id, ( - f"{type(self).__name__} must set " "mmmu_pro_load_preset_from_model_id" + f"{type(self).__name__} must set mmmu_pro_load_preset_from_model_id" ) _run_accuracy_eval( self, diff --git a/python/sglang/test/kits/fwd_occupancy_kit.py b/python/sglang/test/kits/fwd_occupancy_kit.py index e18659a2f..3c5acd6ad 100644 --- a/python/sglang/test/kits/fwd_occupancy_kit.py +++ b/python/sglang/test/kits/fwd_occupancy_kit.py @@ -51,9 +51,7 @@ class FwdOccupancyMixin: # Measurement: one long single-batch request -- max_new_tokens must # span several decode_log_interval windows for enough samples. fwd_occupancy_max_new_tokens: int = 2048 - fwd_occupancy_prompt: str = ( - "Human: Give me a fully functional FastAPI server. Show the python code.\n\nAssistant:" - ) + fwd_occupancy_prompt: str = "Human: Give me a fully functional FastAPI server. Show the python code.\n\nAssistant:" def _scrape_fwd_occupancy(self): """Max non-NaN gauge value across exposed labels (e.g. dp ranks); diff --git a/python/sglang/test/kits/json_constrained_kit.py b/python/sglang/test/kits/json_constrained_kit.py index 6576e61aa..cdc61438d 100644 --- a/python/sglang/test/kits/json_constrained_kit.py +++ b/python/sglang/test/kits/json_constrained_kit.py @@ -6,7 +6,6 @@ import requests class JSONConstrainedMixin: - json_schema = json.dumps( { "type": "object", diff --git a/python/sglang/test/kits/matched_stop_kit.py b/python/sglang/test/kits/matched_stop_kit.py index afccc6779..5693809ac 100644 --- a/python/sglang/test/kits/matched_stop_kit.py +++ b/python/sglang/test/kits/matched_stop_kit.py @@ -45,12 +45,12 @@ class MatchedStopMixin: if not isinstance(matched_stop, list): matched_stop = [matched_stop] - assert ( - res["choices"][0]["finish_reason"] == finish_reason - ), f"Expected finish_reason: {finish_reason}, but got: {res['choices'][0]['finish_reason']}" - assert ( - res["choices"][0]["matched_stop"] in matched_stop - ), f"Expected matched_stop: {matched_stop}, but got: {res['choices'][0]['matched_stop']}" + assert res["choices"][0]["finish_reason"] == finish_reason, ( + f"Expected finish_reason: {finish_reason}, but got: {res['choices'][0]['finish_reason']}" + ) + assert res["choices"][0]["matched_stop"] in matched_stop, ( + f"Expected matched_stop: {matched_stop}, but got: {res['choices'][0]['matched_stop']}" + ) def _run_chat_completions_generation( self, @@ -89,12 +89,12 @@ class MatchedStopMixin: if not isinstance(matched_stop, list): matched_stop = [matched_stop] - assert ( - res["choices"][0]["finish_reason"] == finish_reason - ), f"Expected finish_reason: {finish_reason}, but got: {res['choices'][0]['finish_reason']}" - assert ( - res["choices"][0]["matched_stop"] in matched_stop - ), f"Expected matched_stop: {matched_stop}, but got: {res['choices'][0]['matched_stop']}" + assert res["choices"][0]["finish_reason"] == finish_reason, ( + f"Expected finish_reason: {finish_reason}, but got: {res['choices'][0]['finish_reason']}" + ) + assert res["choices"][0]["matched_stop"] in matched_stop, ( + f"Expected matched_stop: {matched_stop}, but got: {res['choices'][0]['matched_stop']}" + ) def test_finish_stop_str(self): self._run_completions_generation( diff --git a/python/sglang/test/kits/mmmu_vlm_kit.py b/python/sglang/test/kits/mmmu_vlm_kit.py index 8c587037a..c9c5fda65 100644 --- a/python/sglang/test/kits/mmmu_vlm_kit.py +++ b/python/sglang/test/kits/mmmu_vlm_kit.py @@ -161,7 +161,7 @@ class MMMUMixin: os.makedirs(output_path, exist_ok=True) # -------- compose --model_args -------- - model_args = f'model_version="{model_version}",' f"tp={tp}" + model_args = f'model_version="{model_version}",tp={tp}' # -------- build command list -------- cmd = [ @@ -293,7 +293,7 @@ class MMMUMultiModelTestBase(CustomTestCase): os.makedirs(output_path, exist_ok=True) # -------- compose --model_args -------- - model_args = f'model_version="{model_version}",' f"tp={tp}" + model_args = f'model_version="{model_version}",tp={tp}' # -------- build command list -------- cmd = [ diff --git a/python/sglang/test/kits/prefix_cache_branching_kit.py b/python/sglang/test/kits/prefix_cache_branching_kit.py index 780e423a3..16ec2379e 100644 --- a/python/sglang/test/kits/prefix_cache_branching_kit.py +++ b/python/sglang/test/kits/prefix_cache_branching_kit.py @@ -41,10 +41,10 @@ class PrefixCacheBranchingMixin: expected_cached_tokens = ( branching_pos // cls.cache_chunk_size * cls.cache_chunk_size ) - assert ( - cached_tokens == expected_cached_tokens - ), f"{i=}, {cache_hit=}, {cached_tokens=} is not equal to {expected_cached_tokens=}, {branching_pos=}" + assert cached_tokens == expected_cached_tokens, ( + f"{i=}, {cache_hit=}, {cached_tokens=} is not equal to {expected_cached_tokens=}, {branching_pos=}" + ) else: - assert ( - cached_tokens == 0 - ), f"{i=}, {cache_hit=}, {cached_tokens=} is not 0" + assert cached_tokens == 0, ( + f"{i=}, {cache_hit=}, {cached_tokens=} is not 0" + ) diff --git a/python/sglang/test/kits/regex_constrained_kit.py b/python/sglang/test/kits/regex_constrained_kit.py index f4c3e31d9..c37b97a46 100644 --- a/python/sglang/test/kits/regex_constrained_kit.py +++ b/python/sglang/test/kits/regex_constrained_kit.py @@ -4,7 +4,6 @@ import requests class RegexConstrainedMixin: - def _run_decode_regex( self, regex, diff --git a/python/sglang/test/kits/streaming_session_kit.py b/python/sglang/test/kits/streaming_session_kit.py index 4d02263e6..356fb7a1c 100644 --- a/python/sglang/test/kits/streaming_session_kit.py +++ b/python/sglang/test/kits/streaming_session_kit.py @@ -111,9 +111,9 @@ class StreamingSessionKitMixin: # No logprob asyncio.run(_concurrent_logprob_run(self.base_url, self.tokenizer)) time.sleep(3) - assert ( - requests.get(self.base_url + "/health").status_code == 200 - ), "Server unhealthy after concurrent logprob sessions." + assert requests.get(self.base_url + "/health").status_code == 200, ( + "Server unhealthy after concurrent logprob sessions." + ) def test_stress_concurrent_sessions(self) -> None: """High concurrency streaming + non-streaming with retract pressure; diff --git a/python/sglang/test/kl_multiturn_utils.py b/python/sglang/test/kl_multiturn_utils.py index 97219b6d3..a24f18bb8 100644 --- a/python/sglang/test/kl_multiturn_utils.py +++ b/python/sglang/test/kl_multiturn_utils.py @@ -44,9 +44,9 @@ __all__ = [ def default_prefill_cache_assert(result: dict, prefix_len: int, label: str): """Standard radix cache: cached_tokens == prefix_len.""" actual = result["meta_info"]["cached_tokens"] - assert ( - actual == prefix_len - ), f"{label}: expected cached_tokens={prefix_len}, got {actual}" + assert actual == prefix_len, ( + f"{label}: expected cached_tokens={prefix_len}, got {actual}" + ) def default_decode_cache_assert( @@ -55,9 +55,9 @@ def default_decode_cache_assert( """Standard radix cache: cached_tokens == history_len + output_len.""" expected = history_len + output_len actual = result["meta_info"]["cached_tokens"] - assert ( - actual == expected - ), f"{label}: expected cached_tokens={expected}, got {actual}" + assert actual == expected, ( + f"{label}: expected cached_tokens={expected}, got {actual}" + ) def make_mamba_prefill_assert(chunk_size: int = 64) -> Callable: @@ -67,9 +67,9 @@ def make_mamba_prefill_assert(chunk_size: int = 64) -> Callable: actual = result["meta_info"]["cached_tokens"] upper = (prefix_len // chunk_size) * chunk_size lower = max(0, upper - chunk_size) - assert ( - lower <= actual <= upper - ), f"{label}: expected cached_tokens in [{lower}, {upper}], got {actual}" + assert lower <= actual <= upper, ( + f"{label}: expected cached_tokens in [{lower}, {upper}], got {actual}" + ) return _check @@ -85,9 +85,9 @@ def make_mamba_decode_assert(track_interval: int = 16) -> Callable: expected = ( (history_len + output_len - 1) // track_interval ) * track_interval - assert ( - actual >= expected - ), f"{label}: expected cached_tokens={expected}, got {actual}" + assert actual >= expected, ( + f"{label}: expected cached_tokens={expected}, got {actual}" + ) return _check @@ -458,9 +458,9 @@ def test_input_output_logprobs_match_decode_cache_hit_helper( different suffixes per branch. Use branches_per_group for interleaved submission to stress the radix tree. """ - assert ( - len(turn_suffixes) >= 1 - ), "turn_suffixes must have at least 1 entry (for turn 2)" + assert len(turn_suffixes) >= 1, ( + "turn_suffixes must have at least 1 entry (for turn 2)" + ) if assert_decode_cached_tokens is None: assert_decode_cached_tokens = default_decode_cache_assert diff --git a/python/sglang/test/kl_test_utils.py b/python/sglang/test/kl_test_utils.py index dd5ae682c..28efcbfa9 100644 --- a/python/sglang/test/kl_test_utils.py +++ b/python/sglang/test/kl_test_utils.py @@ -285,9 +285,9 @@ def test_input_output_logprobs_match_prefill_cache_hit_helper( output_logprobs.append(_extract_output_logprobs(result)) if not os.environ.get("SGLANG_TEST_SKIP_CACHE_HIT_ASSERT"): - assert len(new_input_ids) > 0.5 * len( - input_ids - ), f"Too few prefill cache hits: {len(new_input_ids)}/{len(input_ids)}" + assert len(new_input_ids) > 0.5 * len(input_ids), ( + f"Too few prefill cache hits: {len(new_input_ids)}/{len(input_ids)}" + ) print("Flush Cache and run prefill to get input logprobs ...") input_logprobs = _get_input_logprobs(base_url, new_input_ids, output_logprobs) @@ -367,9 +367,9 @@ def test_input_output_logprobs_match_decode_cache_hit_helper( # Page-aligned SWA retention decides which prompts hit at all, so the default # only screens out a vacuous run. A caller whose checkpoint interval makes # every prompt hit raises this to pin that down. - assert len(new_input_ids) > min_cache_hit_ratio * len( - second_turn_input_ids - ), f"Too few decode cache hits: {len(new_input_ids)}/{len(second_turn_input_ids)}" + assert len(new_input_ids) > min_cache_hit_ratio * len(second_turn_input_ids), ( + f"Too few decode cache hits: {len(new_input_ids)}/{len(second_turn_input_ids)}" + ) print("Flush Cache and run prefill to get input logprobs ...") input_logprobs = _get_input_logprobs(base_url, new_input_ids, output_logprobs) diff --git a/python/sglang/test/kv_canary/pp_fixture.py b/python/sglang/test/kv_canary/pp_fixture.py index e9f73eea2..9db8695f2 100644 --- a/python/sglang/test/kv_canary/pp_fixture.py +++ b/python/sglang/test/kv_canary/pp_fixture.py @@ -9,7 +9,6 @@ PP_SIZE: int = 2 class CanaryPPFixture(CanaryE2EBase): - model_mode: ClassVar[str] = "swa" workload_n_batches: ClassVar[int] = 2 diff --git a/python/sglang/test/lora_utils.py b/python/sglang/test/lora_utils.py index 0bfd7fa07..f5c65b18e 100644 --- a/python/sglang/test/lora_utils.py +++ b/python/sglang/test/lora_utils.py @@ -616,7 +616,6 @@ def run_lora_test_by_batch( ) for i in range(len(prompts)): - srt_output_str = srt_outputs.output_strs[i].strip() hf_output_str = hf_outputs.output_strs[i].strip() rouge_score = calculate_rouge_l([srt_output_str], [hf_output_str])[0] @@ -785,7 +784,7 @@ def run_lora_multiple_batch_on_model_cases( with srt_runner, hf_runner: for i, (prompts, lora_paths) in enumerate(batches): print( - f"\n--- Running Batch {i+1} --- prompts: {prompts}, lora_paths: {lora_paths}" + f"\n--- Running Batch {i + 1} --- prompts: {prompts}, lora_paths: {lora_paths}" ) srt_outputs = srt_runner.batch_forward( @@ -816,7 +815,7 @@ def run_lora_multiple_batch_on_model_cases( f"for base '{base_path}', adaptor '{lora_paths}', prompt: '{prompts}...'" ) - print(f"--- Batch {i+1} Comparison Passed --- ") + print(f"--- Batch {i + 1} Comparison Passed --- ") def run_lora_batch_splitting_equivalence_test( @@ -851,9 +850,9 @@ def run_lora_batch_splitting_equivalence_test( def _run_test(model_case: LoRAModelCase, torch_dtype: torch.dtype): lora_adapter_paths = [a.name for a in model_case.adaptors] - assert ( - len(lora_adapter_paths) >= max_loras_per_batch - ), f"Need at least {max_loras_per_batch} adapters for this test" + assert len(lora_adapter_paths) >= max_loras_per_batch, ( + f"Need at least {max_loras_per_batch} adapters for this test" + ) max_new_tokens = 64 base_path = model_case.base diff --git a/python/sglang/test/performance_test_runner.py b/python/sglang/test/performance_test_runner.py index 5f1f2b055..a73c95287 100644 --- a/python/sglang/test/performance_test_runner.py +++ b/python/sglang/test/performance_test_runner.py @@ -55,7 +55,7 @@ def run_performance_test( if batch_sizes is None: batch_sizes = [1, 8, 16, 64] - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Running PERFORMANCE test for {model.model_path}") print(f" Variant: {model.variant}") print(f" Batch sizes: {batch_sizes}") @@ -63,7 +63,7 @@ def run_performance_test( print(f" Output lens: {output_lens}") if spec_accept_length_threshold is not None: print(f" Spec accept length threshold: {spec_accept_length_threshold}") - print(f"{'='*60}\n") + print(f"{'=' * 60}\n") # Build extra args for benchmarks extra_bench_args = ["--trust-remote-code"] diff --git a/python/sglang/test/runners.py b/python/sglang/test/runners.py index 03a33e20f..a2f4c1c1c 100644 --- a/python/sglang/test/runners.py +++ b/python/sglang/test/runners.py @@ -935,9 +935,9 @@ def check_close_model_outputs( print(f"{srt_outputs.output_strs=}") rouge_l_scores = calculate_rouge_l(hf_outputs.output_strs, srt_outputs.output_strs) print(f"{rouge_l_scores=}") - assert all( - score >= rouge_l_tolerance for score in rouge_l_scores - ), f"Not all ROUGE-L scores are greater than rouge_l_tolerance={rouge_l_tolerance}" + assert all(score >= rouge_l_tolerance for score in rouge_l_scores), ( + f"Not all ROUGE-L scores are greater than rouge_l_tolerance={rouge_l_tolerance}" + ) if check_logprobs: for i in range(len(hf_outputs.output_strs)): diff --git a/python/sglang/test/scripted_runtime/background_http_poster.py b/python/sglang/test/scripted_runtime/background_http_poster.py index 671d22e1a..201e7eff8 100644 --- a/python/sglang/test/scripted_runtime/background_http_poster.py +++ b/python/sglang/test/scripted_runtime/background_http_poster.py @@ -14,7 +14,6 @@ JOIN_TIMEOUT_S: float = 10.0 class BackgroundHttpPoster: - def __init__(self) -> None: self._session: Optional[aiohttp.ClientSession] = None self._loop = asyncio.new_event_loop() diff --git a/python/sglang/test/scripted_runtime/context/api.py b/python/sglang/test/scripted_runtime/context/api.py index bb3a8fbd3..8576b2cfb 100644 --- a/python/sglang/test/scripted_runtime/context/api.py +++ b/python/sglang/test/scripted_runtime/context/api.py @@ -30,7 +30,6 @@ logger = logging.getLogger(__name__) class ScriptedContext: - def __init__( self, *, @@ -38,9 +37,9 @@ class ScriptedContext: tokenizer_recv_proxy: Optional[ScriptedTokenizerRecvProxy], http_poster: BackgroundHttpPoster, ) -> None: - assert ( - scheduler_hook._is_driver - ), "ScriptedContext only exists on the driver rank" + assert scheduler_hook._is_driver, ( + "ScriptedContext only exists on the driver rank" + ) self.scheduler = scheduler_hook.scheduler self._scheduler_hook = scheduler_hook self._tokenizer_recv_proxy = tokenizer_recv_proxy @@ -100,9 +99,9 @@ class ScriptedContext: return lifecycle.flush_cache(self) def evict_radix(self, *, prefix_tokens: Optional[List[int]]) -> None: - assert ( - prefix_tokens is None - ), "evict_radix currently supports only full eviction (prefix_tokens=None)" + assert prefix_tokens is None, ( + "evict_radix currently supports only full eviction (prefix_tokens=None)" + ) return lifecycle.flush_cache(self) def exhaust_kv(self, *, leave_pages: int) -> None: diff --git a/python/sglang/test/scripted_runtime/context/kv_pool_exhauster.py b/python/sglang/test/scripted_runtime/context/kv_pool_exhauster.py index 3004dffea..ae6e2bbad 100644 --- a/python/sglang/test/scripted_runtime/context/kv_pool_exhauster.py +++ b/python/sglang/test/scripted_runtime/context/kv_pool_exhauster.py @@ -9,7 +9,6 @@ if TYPE_CHECKING: class ScriptedKvPoolExhauster: - def __init__(self, scheduler: Scheduler) -> None: self.scheduler = scheduler self._held: List[torch.Tensor] = [] @@ -23,9 +22,9 @@ class ScriptedKvPoolExhauster: return held = allocator.alloc(need) - assert ( - held is not None - ), f"exhaust_kv: allocator could not grab {need} tokens to create pressure" + assert held is not None, ( + f"exhaust_kv: allocator could not grab {need} tokens to create pressure" + ) self._held.append(held) def release(self) -> None: diff --git a/python/sglang/test/scripted_runtime/context/lock_ref_exhauster.py b/python/sglang/test/scripted_runtime/context/lock_ref_exhauster.py index 2530348fd..cf3e7bd8c 100644 --- a/python/sglang/test/scripted_runtime/context/lock_ref_exhauster.py +++ b/python/sglang/test/scripted_runtime/context/lock_ref_exhauster.py @@ -9,7 +9,6 @@ if TYPE_CHECKING: class ScriptedLockRefExhauster: - def __init__(self, scheduler: Scheduler) -> None: self.scheduler = scheduler self._locked: List[Any] = [] diff --git a/python/sglang/test/scripted_runtime/http_server.py b/python/sglang/test/scripted_runtime/http_server.py index 6d86a8652..21e99e97c 100644 --- a/python/sglang/test/scripted_runtime/http_server.py +++ b/python/sglang/test/scripted_runtime/http_server.py @@ -34,7 +34,6 @@ SERVER_HOST: str = "127.0.0.1" class ScriptedHttpServer: - def __init__( self, *, diff --git a/python/sglang/test/scripted_runtime/io_struct.py b/python/sglang/test/scripted_runtime/io_struct.py index c026329a1..80cadc202 100644 --- a/python/sglang/test/scripted_runtime/io_struct.py +++ b/python/sglang/test/scripted_runtime/io_struct.py @@ -8,7 +8,6 @@ from typing import Any, Tuple, Union @dataclass(frozen=True, slots=True) class RunScript: - fn_path: str args: Tuple[Any, ...] = () @@ -30,13 +29,11 @@ class ScriptSucceeded: @dataclass(frozen=True, slots=True) class ScriptFailed: - traceback: str @dataclass(frozen=True, slots=True) class OutOfBandError: - traceback: str def to_json(self) -> str: diff --git a/python/sglang/test/scripted_runtime/scheduler_hook.py b/python/sglang/test/scripted_runtime/scheduler_hook.py index 51462b5d4..b2a4ca68b 100644 --- a/python/sglang/test/scripted_runtime/scheduler_hook.py +++ b/python/sglang/test/scripted_runtime/scheduler_hook.py @@ -117,7 +117,6 @@ def _reset_engine_state(ctx: ScriptedContext) -> Generator: class ScriptedSchedulerHook: - def __init__( self, *, diff --git a/python/sglang/test/scripted_runtime/test_case.py b/python/sglang/test/scripted_runtime/test_case.py index 491ab5e09..51ff9f0e0 100644 --- a/python/sglang/test/scripted_runtime/test_case.py +++ b/python/sglang/test/scripted_runtime/test_case.py @@ -7,7 +7,6 @@ from sglang.test.test_utils import CustomTestCase class ScriptedTestCase(CustomTestCase): - ENGINE_KWARGS: ClassVar[Dict[str, Any]] = {} server: ClassVar[ScriptedHttpServer] diff --git a/python/sglang/test/scripted_runtime/tokenizer_recv_proxy.py b/python/sglang/test/scripted_runtime/tokenizer_recv_proxy.py index aee5795dd..a65187342 100644 --- a/python/sglang/test/scripted_runtime/tokenizer_recv_proxy.py +++ b/python/sglang/test/scripted_runtime/tokenizer_recv_proxy.py @@ -23,7 +23,6 @@ _WORK_REQ_TYPES = ( class ScriptedTokenizerRecvProxy: - def __init__(self, *, underlying: zmq.Socket) -> None: self._underlying = underlying self._buffer: deque = deque() diff --git a/python/sglang/test/scripted_runtime_chunked_helpers.py b/python/sglang/test/scripted_runtime_chunked_helpers.py index 393d064ff..2324313cd 100644 --- a/python/sglang/test/scripted_runtime_chunked_helpers.py +++ b/python/sglang/test/scripted_runtime_chunked_helpers.py @@ -57,9 +57,9 @@ def run_until_all_finished(handles: List[Any], *, max_steps: int = DEFAULT_MAX_S def warmup_radix(t, prompt_tokens: List[int], *, max_steps: int = DEFAULT_MAX_STEPS): assert prompt_tokens, "warmup_radix needs a non-empty prompt" token = prompt_tokens[0] - assert all( - x == token for x in prompt_tokens - ), "warmup_radix supports only uniform prompts" + assert all(x == token for x in prompt_tokens), ( + "warmup_radix supports only uniform prompts" + ) handle = t.start_req( prompt_len=len(prompt_tokens), max_new_tokens=1, prompt_token=token ) @@ -126,9 +126,9 @@ def advance_to_decode_step( r, target_output_len: int, *, max_steps: int = DEFAULT_MAX_STEPS ): for _ in range(max_steps): - assert ( - not r.finished - ), f"req finished before reaching decode step {target_output_len}" + assert not r.finished, ( + f"req finished before reaching decode step {target_output_len}" + ) req = r.req if req is not None and len(req.output_ids) >= target_output_len: return diff --git a/python/sglang/test/send_one.py b/python/sglang/test/send_one.py index e6cd7d992..678d7444c 100644 --- a/python/sglang/test/send_one.py +++ b/python/sglang/test/send_one.py @@ -39,9 +39,7 @@ class BenchArgs: presence_penalty: float = 0.0 json: bool = False return_logprob: bool = False - prompt: str = ( - "Human: Give me a fully functional FastAPI server. Show the python code.\n\nAssistant:" - ) + prompt: str = "Human: Give me a fully functional FastAPI server. Show the python code.\n\nAssistant:" image: bool = False many_images: bool = False stop: Optional[list] = None @@ -144,7 +142,7 @@ def send_one_prompt( else: if args.different_prompts: prompt = [ - f"Test case {i+1}: " + args.prompt for i in range(args.batch_size) + f"Test case {i + 1}: " + args.prompt for i in range(args.batch_size) ] else: prompt = [args.prompt] * args.batch_size diff --git a/python/sglang/test/server_fixtures/disaggregation_fixture.py b/python/sglang/test/server_fixtures/disaggregation_fixture.py index 0aba39f7e..2169ec3e6 100644 --- a/python/sglang/test/server_fixtures/disaggregation_fixture.py +++ b/python/sglang/test/server_fixtures/disaggregation_fixture.py @@ -378,7 +378,7 @@ def get_rdma_devices_args(): if not (base_rdma_group <= gpu_idx < base_rdma_group + 4): warnings.warn( f"GPU index {gpu_idx} is outside expected group " - f"{base_rdma_group}-{base_rdma_group+3}" + f"{base_rdma_group}-{base_rdma_group + 3}" ) # 3. Generate RDMA device names diff --git a/python/sglang/test/server_fixtures/hybrid_attn_backend_fixture.py b/python/sglang/test/server_fixtures/hybrid_attn_backend_fixture.py index 9aba6ee48..c1685edce 100644 --- a/python/sglang/test/server_fixtures/hybrid_attn_backend_fixture.py +++ b/python/sglang/test/server_fixtures/hybrid_attn_backend_fixture.py @@ -40,7 +40,6 @@ DEFAULT_HYBRID_ATTN_SERVER_ARGS = [ @unittest.skipIf(get_device_sm() < 90, "Test requires CUDA SM 90 or higher") class TestHybridAttnBackendBase(CustomTestCase): - model = DEFAULT_MODEL_NAME_FOR_TEST base_url = DEFAULT_URL_FOR_TEST accuracy_threshold = 0.65 # derived tests need to override this diff --git a/python/sglang/test/server_fixtures/pcg_spec_fixture.py b/python/sglang/test/server_fixtures/pcg_spec_fixture.py index 3bcfd2f61..77b88b2db 100644 --- a/python/sglang/test/server_fixtures/pcg_spec_fixture.py +++ b/python/sglang/test/server_fixtures/pcg_spec_fixture.py @@ -40,9 +40,9 @@ class PCGSpecBase: @classmethod def setUpClass(cls): - assert ( - cls.model and cls.server_args - ), f"{cls.__name__} must set `model` and `server_args`" + assert cls.model and cls.server_args, ( + f"{cls.__name__} must set `model` and `server_args`" + ) cls.base_url = DEFAULT_URL_FOR_TEST kwargs = dict( timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * cls.timeout_mult, diff --git a/python/sglang/test/server_fixtures/streaming_session_fixture.py b/python/sglang/test/server_fixtures/streaming_session_fixture.py index f23babbe2..0996c7457 100644 --- a/python/sglang/test/server_fixtures/streaming_session_fixture.py +++ b/python/sglang/test/server_fixtures/streaming_session_fixture.py @@ -314,8 +314,7 @@ async def _concurrent_logprob_run(base_url: str, tokenizer: Any, **gen_kwargs) - tasks = [] for s in range(CONCURRENT_LOGPROB_SESSIONS): text = ( - f"S{s} T{turn}: " - f"{LOGPROB_PROMPTS[turn % len(LOGPROB_PROMPTS)]}" + f"S{s} T{turn}: {LOGPROB_PROMPTS[turn % len(LOGPROB_PROMPTS)]}" ) ids = tokenizer.encode(text) tasks.append( @@ -360,9 +359,7 @@ async def _stress_run_all(base_url: str, tokenizer: Any) -> None: # Streaming requests — long prompts to trigger chunked prefill. for s in range(STRESS_NUM_SESSIONS): offset = (s * STRESS_NUM_TURNS + turn) * 200 - text = ( - f"Session {s} turn {turn}: " f"{LEAK_FILLER[offset : offset + 800]}" - ) + text = f"Session {s} turn {turn}: {LEAK_FILLER[offset : offset + 800]}" ids = tokenizer.encode(text) tasks.append( _async_generate( diff --git a/python/sglang/test/simple_eval_mixed_prefix_gsm8k.py b/python/sglang/test/simple_eval_mixed_prefix_gsm8k.py index 34a0db672..61b997952 100644 --- a/python/sglang/test/simple_eval_mixed_prefix_gsm8k.py +++ b/python/sglang/test/simple_eval_mixed_prefix_gsm8k.py @@ -112,7 +112,6 @@ class GSM8KEval(Eval): class MixedPrefixGSM8KEval(GSM8KEval): - def __init__( self, num_examples: Optional[int], diff --git a/python/sglang/test/test_deepep_utils.py b/python/sglang/test/test_deepep_utils.py index aa15b5a0b..d11c40313 100644 --- a/python/sglang/test/test_deepep_utils.py +++ b/python/sglang/test/test_deepep_utils.py @@ -192,9 +192,9 @@ def bench_kineto( kernel_names = (kernel_names,) if isinstance(kernel_names, str) else kernel_names assert all([isinstance(name, str) for name in kernel_names]) for name in kernel_names: - assert ( - sum([name in line for line in prof_lines]) == 1 - ), f"Errors of the kernel {name} in the profiling table" + assert sum([name in line for line in prof_lines]) == 1, ( + f"Errors of the kernel {name} in the profiling table" + ) # Save chrome traces if trace_path is not None: diff --git a/python/sglang/test/test_deterministic.py b/python/sglang/test/test_deterministic.py index 17f5495ab..8d02a674a 100644 --- a/python/sglang/test/test_deterministic.py +++ b/python/sglang/test/test_deterministic.py @@ -510,7 +510,7 @@ def test_deterministic(args): # If logprobs are enabled, compare them across different batch sizes if args.return_logprob: - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print("Logprobs Comparison Across Batch Sizes") print("=" * 60) @@ -536,7 +536,7 @@ def test_deterministic(args): match, msg = compare_logprobs(ref_logprobs, resp_logprobs) if not match: - print(f" ✗ Sample {j+1}: {msg}") + print(f" ✗ Sample {j + 1}: {msg}") mismatches.append((j + 1, msg)) all_match = False @@ -549,7 +549,7 @@ def test_deterministic(args): ) logprob_results.append(0) - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") if all(r == 1 for r in logprob_results): print("✓✓✓ Logprobs are identical across all batch sizes! ✓✓✓") else: @@ -653,7 +653,7 @@ def test_deterministic(args): print(f" Logprob: {uncached_logprob:.10f}") # Step 6: Compare results - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print("Comparison 1: Decode (Request 1) vs Prefill with Cache (Request 2)") print("=" * 60) @@ -679,7 +679,7 @@ def test_deterministic(args): print(f" Logprob difference: {diff:.10e}") print(f" Note: We expect these to be DIFFERENT (decode vs prefill kernels)") - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print( "Comparison 2: Cached Prefill (Request 2) vs Uncached Prefill (Request 3)" ) @@ -708,7 +708,7 @@ def test_deterministic(args): print(f" Difference: {diff:.10e}") print(f" Note: We expect these to be IDENTICAL (both prefill kernels)") - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") if token_match and logprob_match: print("✓✓✓ TEST PASSED - Radix cache is consistent! ✓✓✓") return [1] diff --git a/python/sglang/test/test_marlin_utils.py b/python/sglang/test/test_marlin_utils.py index bd53c8332..6bf4693a0 100644 --- a/python/sglang/test/test_marlin_utils.py +++ b/python/sglang/test/test_marlin_utils.py @@ -27,12 +27,11 @@ from sglang.srt.layers.quantization.utils import ( class MarlinWorkspace: - def __init__(self, out_features, min_thread_n, max_parallel): - assert ( - out_features % min_thread_n == 0 - ), "out_features = {} is undivisible by min_thread_n = {}".format( - out_features, min_thread_n + assert out_features % min_thread_n == 0, ( + "out_features = {} is undivisible by min_thread_n = {}".format( + out_features, min_thread_n + ) ) max_workspace_size = (out_features // min_thread_n) * max_parallel diff --git a/python/sglang/test/test_programs.py b/python/sglang/test/test_programs.py index 6779814a0..01755192b 100644 --- a/python/sglang/test/test_programs.py +++ b/python/sglang/test/test_programs.py @@ -262,16 +262,16 @@ def test_parallel_decoding(): # Generate detailed tips forks = s.fork(fork_size) for i in range(fork_size): - forks[ - i - ] += f"Now, I expand tip {i+1} into a detailed paragraph:\nTip {i+1}:" + forks[i] += ( + f"Now, I expand tip {i + 1} into a detailed paragraph:\nTip {i + 1}:" + ) forks[i] += sgl.gen("detailed_tip", max_tokens, stop=["\n\n"]) forks.join() # Concatenate tips and summarize s += "Here are these tips with detailed explanation:\n" for i in range(fork_size): - s += f"Tip {i+1}:" + forks[i]["detailed_tip"] + "\n" + s += f"Tip {i + 1}:" + forks[i]["detailed_tip"] + "\n" s += "\nIn summary," + sgl.gen("summary", max_tokens=512) @@ -294,7 +294,7 @@ def test_parallel_encoding(check_answer=True): forks += lambda i: f"Statement {i}: " + contexts[i] + "\n" forks.join(mode="concate_and_append") - s += "Now, please answer the following question. " "Do not list options." + s += "Now, please answer the following question. Do not list options." s += "\nQuestion: " + question + "\n" s += "ASSISTANT:" + sgl.gen("answer", max_tokens=max_tokens) @@ -474,9 +474,9 @@ def test_completion_speculative(): gen_character_no_spec().sync() usage_with_no_spec = token_usage.prompt_tokens - assert ( - usage_with_spec < usage_with_no_spec - ), f"{usage_with_spec} vs {usage_with_no_spec}" + assert usage_with_spec < usage_with_no_spec, ( + f"{usage_with_spec} vs {usage_with_no_spec}" + ) def test_chat_completion_speculative(): @@ -612,9 +612,9 @@ def test_gen_min_new_tokens(): def assert_min_tokens(tokenizer, text): token_ids = tokenizer.encode(text) - assert ( - len(token_ids) >= MIN_TOKENS - ), f"Generated {len(token_ids)} tokens, min required: {MIN_TOKENS}. Text: {text}" + assert len(token_ids) >= MIN_TOKENS, ( + f"Generated {len(token_ids)} tokens, min required: {MIN_TOKENS}. Text: {text}" + ) tokenizer = get_tokenizer(model_path) diff --git a/python/sglang/test/test_utils.py b/python/sglang/test/test_utils.py index 143e9f641..8c5ae9eae 100644 --- a/python/sglang/test/test_utils.py +++ b/python/sglang/test/test_utils.py @@ -2144,7 +2144,6 @@ def server_args_variant(server_args, **fields): class CustomTestCase(unittest.TestCase): - def __init_subclass__(cls, **kwargs): super().__init_subclass__(**kwargs) diff --git a/python/sglang/test/tool_call_test_runner.py b/python/sglang/test/tool_call_test_runner.py index a030d5605..b642260cb 100644 --- a/python/sglang/test/tool_call_test_runner.py +++ b/python/sglang/test/tool_call_test_runner.py @@ -263,14 +263,14 @@ def _test_reasoning_usage(client, model): ) usage = response.usage assert usage is not None, "usage should not be None" - assert ( - usage.reasoning_tokens and usage.reasoning_tokens > 0 - ), f"expected reasoning_tokens > 0, got {usage.reasoning_tokens}" + assert usage.reasoning_tokens and usage.reasoning_tokens > 0, ( + f"expected reasoning_tokens > 0, got {usage.reasoning_tokens}" + ) if usage.completion_tokens_details: detail_reasoning = usage.completion_tokens_details.get("reasoning_tokens", 0) - assert ( - detail_reasoning > 0 - ), f"expected completion_tokens_details.reasoning_tokens > 0, got {detail_reasoning}" + assert detail_reasoning > 0, ( + f"expected completion_tokens_details.reasoning_tokens > 0, got {detail_reasoning}" + ) def _test_parallel(client, model): diff --git a/python/sglang/test/vlm_utils.py b/python/sglang/test/vlm_utils.py index 4dfa8417f..8446a276d 100644 --- a/python/sglang/test/vlm_utils.py +++ b/python/sglang/test/vlm_utils.py @@ -95,9 +95,9 @@ class AudioOpenAITestMixin(TestOpenAIMLLMServerBase): "art", ] for check_word in check_list: - assert ( - check_word in text.lower() - ), f"audio_response: |{text}| should contain |{check_word}|" + assert check_word in text.lower(), ( + f"audio_response: |{text}| should contain |{check_word}|" + ) def prepare_audio_messages(self, prompt, audio_file_name): messages = [ @@ -285,9 +285,9 @@ class ImageOpenAITestMixin(TestOpenAIMLLMServerBase): assert isinstance(text, str) # `driver` is for gemma-3-it - assert any( - keyword in text for keyword in ("man", "person", "driver") - ), f"text: {text}, should contain man, person or driver" + assert any(keyword in text for keyword in ("man", "person", "driver")), ( + f"text: {text}, should contain man, person or driver" + ) assert ( "cab" in text or "taxi" in text @@ -377,9 +377,9 @@ class ImageOpenAITestMixin(TestOpenAIMLLMServerBase): assert response.choices[0].message.role == "assistant" text = response.choices[0].message.content assert isinstance(text, str) - assert ( - "man" in text or "cab" in text - ), f"text: {text}, should contain man or cab" + assert "man" in text or "cab" in text, ( + f"text: {text}, should contain man or cab" + ) assert response.id assert response.created assert response.usage.prompt_tokens > 0 @@ -429,9 +429,9 @@ class ImageOpenAITestMixin(TestOpenAIMLLMServerBase): or "taxi" in text or "car" in text ), f"text: {text}, should contain man, cab, SUV, taxi or car" - assert ( - "logo" in text or '"S"' in text or "SG" in text or "graphic" in text - ), f"text: {text}, should contain logo, S or SG or graphic" + assert "logo" in text or '"S"' in text or "SG" in text or "graphic" in text, ( + f"text: {text}, should contain logo, S or SG or graphic" + ) assert response.id assert response.created assert response.usage.prompt_tokens > 0 @@ -594,16 +594,20 @@ class VideoOpenAITestMixin(TestOpenAIMLLMServerBase): or "speaker" in video_response or "presenter" in video_response or "hand" in video_response - ), f"video_response: {video_response}, should either have 'man' in video_response, or 'person' in video_response, or 'individual' in video_response or 'speaker' in video_response or 'presenter' or 'hand' in video_response" + ), ( + f"video_response: {video_response}, should either have 'man' in video_response, or 'person' in video_response, or 'individual' in video_response or 'speaker' in video_response or 'presenter' or 'hand' in video_response" + ) assert ( "present" in video_response or "examine" in video_response or "display" in video_response or "hold" in video_response - ), f"video_response: {video_response}, should contain 'present', 'examine', 'display', or 'hold'" - assert ( - "black" in video_response or "dark" in video_response - ), f"video_response: {video_response}, should contain 'black' or 'dark'" + ), ( + f"video_response: {video_response}, should contain 'present', 'examine', 'display', or 'hold'" + ) + assert "black" in video_response or "dark" in video_response, ( + f"video_response: {video_response}, should contain 'black' or 'dark'" + ) self.assertIsNotNone(video_response) self.assertGreater(len(video_response), 0) diff --git a/python/sglang/test/xpu/simple_eval_gsm8k_xpu_mixin.py b/python/sglang/test/xpu/simple_eval_gsm8k_xpu_mixin.py index 6ed2e284d..edd10be7d 100644 --- a/python/sglang/test/xpu/simple_eval_gsm8k_xpu_mixin.py +++ b/python/sglang/test/xpu/simple_eval_gsm8k_xpu_mixin.py @@ -105,7 +105,7 @@ class SimpleEvalGSM8KXPUMixin(ABC): self.assertGreaterEqual( metrics["score"], accuracy_threshold, - f'Accuracy of {self.model} is {metrics["score"]}, ' + f"Accuracy of {self.model} is {metrics['score']}, " f"is lower than {accuracy_threshold}", ) if "output_throughput" in metrics: @@ -113,7 +113,7 @@ class SimpleEvalGSM8KXPUMixin(ABC): metrics["output_throughput"], output_throughput_threshold, f"Output throughput of {self.model} is " - f'{metrics["output_throughput"]}, is lower than ' + f"{metrics['output_throughput']}, is lower than " f"{output_throughput_threshold}", ) except Exception as e: diff --git a/rust/sglang-mm/bench/bench_parity.py b/rust/sglang-mm/bench/bench_parity.py index 584245992..487ae1a49 100644 --- a/rust/sglang-mm/bench/bench_parity.py +++ b/rust/sglang-mm/bench/bench_parity.py @@ -92,7 +92,7 @@ def parity_b(): if torch.equal( ref.contiguous().view(torch.uint16), got2.contiguous().view(torch.uint16) ): - print(f" {fmt}: bit-exact=True ({len(data)/1e6:.2f}MB)") + print(f" {fmt}: bit-exact=True ({len(data) / 1e6:.2f}MB)") else: d = (ref.float() - got2.float()).abs() print( diff --git a/rust/sglang-mm/tests/test_resize_parity.py b/rust/sglang-mm/tests/test_resize_parity.py index 727fef31e..d922e9871 100644 --- a/rust/sglang-mm/tests/test_resize_parity.py +++ b/rust/sglang-mm/tests/test_resize_parity.py @@ -156,5 +156,5 @@ def test_resize_bench(): rs_resize(arr, tw, th) t_rs = (time.perf_counter() - t0) / 10 * 1e3 print( - f"\nresize 1920x1080->{tw}x{th}: PIL {t_pil:.1f}ms rust {t_rs:.1f}ms ({t_pil/t_rs:.1f}x)" + f"\nresize 1920x1080->{tw}x{th}: PIL {t_pil:.1f}ms rust {t_rs:.1f}ms ({t_pil / t_rs:.1f}x)" ) diff --git a/scripts/ci/test_list_stage_models.py b/scripts/ci/test_list_stage_models.py index b4287800a..661d15a42 100644 --- a/scripts/ci/test_list_stage_models.py +++ b/scripts/ci/test_list_stage_models.py @@ -590,8 +590,9 @@ class ResolveCommit(unittest.TestCase): self.assertEqual(lsm.resolve_commit(None, "/nonexistent"), "deadbeef") def test_unknown_when_no_git(self): - with tempfile.TemporaryDirectory() as tmp, mock.patch.dict( - os.environ, {}, clear=True + with ( + tempfile.TemporaryDirectory() as tmp, + mock.patch.dict(os.environ, {}, clear=True), ): self.assertEqual(lsm.resolve_commit(None, tmp), "unknown") diff --git a/scripts/ci/update_est_time.py b/scripts/ci/update_est_time.py index cce4b6353..52cc4c0f6 100755 --- a/scripts/ci/update_est_time.py +++ b/scripts/ci/update_est_time.py @@ -95,7 +95,7 @@ def update_files(model, dry_run=False): new_content = pattern.sub(rf"\g<1>{p90}\3", new_content) changes.append((relpath, suite, old_val, p90)) print( - f" {relpath}: suite={suite!r} " f"est_time {old_val} -> {p90}", + f" {relpath}: suite={suite!r} est_time {old_val} -> {p90}", file=sys.stderr, ) break # one (file, suite) -> at most one register call @@ -121,8 +121,7 @@ def write_summary(changes, summary_file): lines = [] if sig: lines.append( - f"### Significant est_time changes " - f"({len(sig)} of {len(changes)} updates)" + f"### Significant est_time changes ({len(sig)} of {len(changes)} updates)" ) lines.append("") lines.append("| File | Suite | Old (s) | New (s) | Δ |") diff --git a/scripts/ci/utils/ci_coverage_report.py b/scripts/ci/utils/ci_coverage_report.py index 1a405cef9..3829e0465 100755 --- a/scripts/ci/utils/ci_coverage_report.py +++ b/scripts/ci/utils/ci_coverage_report.py @@ -34,9 +34,9 @@ from ci_register import CIRegistry, HWBackend, ut_parse_one_file # (highest test volume historically), then accelerators that have been # wired into the registry more recently (XPU, MUSA, MLX). BACKEND_DISPLAY_ORDER = ("CUDA", "AMD", "NPU", "CPU", "XPU", "MUSA", "MLX") -assert set(BACKEND_DISPLAY_ORDER) == { - b.name for b in HWBackend -}, "BACKEND_DISPLAY_ORDER is out of sync with HWBackend" +assert set(BACKEND_DISPLAY_ORDER) == {b.name for b in HWBackend}, ( + "BACKEND_DISPLAY_ORDER is out of sync with HWBackend" +) # --------------------------------------------------------------------------- # # multimodal_gen test coverage diff --git a/scripts/ci/utils/diffusion/compute_diffusion_partitions.py b/scripts/ci/utils/diffusion/compute_diffusion_partitions.py index bdfe2ac03..f5f920dff 100755 --- a/scripts/ci/utils/diffusion/compute_diffusion_partitions.py +++ b/scripts/ci/utils/diffusion/compute_diffusion_partitions.py @@ -198,7 +198,7 @@ def print_suite_summary( ) for standalone_file in suite_info.missing_standalone_estimates: print(f" - {standalone_file}") - print(f" Total estimated time: {total_time:.1f}s ({total_time/60:.1f} min)") + print(f" Total estimated time: {total_time:.1f}s ({total_time / 60:.1f} min)") print(f" Selected partitions: {len(partitions)}") print() @@ -207,7 +207,7 @@ def print_suite_summary( partition_time = sum(item.est_time for item in partition) print(f" Partition {idx}:") print( - f" Estimated time: {partition_time:.1f}s ({partition_time/60:.1f} min)" + f" Estimated time: {partition_time:.1f}s ({partition_time / 60:.1f} min)" ) for item in partition: fallback_suffix = ( @@ -304,9 +304,11 @@ def main(): sys.exit(1) print("=== Diffusion Partition Computation ===") - print(f"Min partition time: {args.min_time}s ({args.min_time/60:.1f} min)") - print(f"Target partition time: {args.target_time}s ({args.target_time/60:.1f} min)") - print(f"Max partition time: {args.max_time}s ({args.max_time/60:.1f} min)") + print(f"Min partition time: {args.min_time}s ({args.min_time / 60:.1f} min)") + print( + f"Target partition time: {args.target_time}s ({args.target_time / 60:.1f} min)" + ) + print(f"Max partition time: {args.max_time}s ({args.max_time / 60:.1f} min)") print() for suite_name, suite_info in suites.items(): diff --git a/scripts/ci/utils/diffusion/generate_diffusion_dashboard.py b/scripts/ci/utils/diffusion/generate_diffusion_dashboard.py index a069b2463..9ca9b41bc 100644 --- a/scripts/ci/utils/diffusion/generate_diffusion_dashboard.py +++ b/scripts/ci/utils/diffusion/generate_diffusion_dashboard.py @@ -273,7 +273,7 @@ def generate_dashboard( if pct > REGRESSION_THRESHOLD: regressions.append( f"**{cid}** ({fw}): {prev:.2f}s -> {cur:.2f}s " - f"(+{pct*100:.1f}%)" + f"(+{pct * 100:.1f}%)" ) if regressions: diff --git a/scripts/ci/utils/diffusion/run_comparison.py b/scripts/ci/utils/diffusion/run_comparison.py index 86d765130..74cc651a5 100644 --- a/scripts/ci/utils/diffusion/run_comparison.py +++ b/scripts/ci/utils/diffusion/run_comparison.py @@ -855,9 +855,9 @@ def _install_framework(fw_name: str, dry_run: bool = False) -> bool: if dry_run: print(f" [DRY-RUN] Would install: bash {INSTALL_SCRIPT} {fw_name}") return True - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Installing framework: {fw_name}") - print(f"{'='*60}") + print(f"{'=' * 60}") ret = subprocess.run( ["bash", str(INSTALL_SCRIPT), fw_name], timeout=600, @@ -943,9 +943,9 @@ def run_comparison( installed_fws.add(fw_name) for case, fw_cfg in pairs: - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Case: {case['id']} | Model: {case['model']} | Framework: {fw_name}") - print(f"{'='*60}") + print(f"{'=' * 60}") if dry_run: cmd = build_server_cmd(fw_name, case, fw_cfg, port) @@ -982,9 +982,9 @@ def run_comparison( print(f"\nResults written to {output}") # Print summary table - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print("SUMMARY") - print(f"{'='*60}") + print(f"{'=' * 60}") for r in results: lat = f"{r['latency_s']:.2f}s" if r["latency_s"] else r.get("error", "N/A") print(f" {r['case_id']:30s} | {r['framework']:12s} | {lat}") diff --git a/scripts/ci/utils/diffusion/verify_diffusion_coverage.py b/scripts/ci/utils/diffusion/verify_diffusion_coverage.py index b327e38bb..925b99c26 100755 --- a/scripts/ci/utils/diffusion/verify_diffusion_coverage.py +++ b/scripts/ci/utils/diffusion/verify_diffusion_coverage.py @@ -155,8 +155,7 @@ def print_missing_standalone_estimates_summary( print("\n" + "=" * 60) print( - "Add standalone estimate(s) to " - "python/sglang/multimodal_gen/test/run_suite.py" + "Add standalone estimate(s) to python/sglang/multimodal_gen/test/run_suite.py" ) print("=" * 60) print("The following standalone file(s) used fallback estimate 300.0s.") diff --git a/scripts/ci/utils/publish_traces.py b/scripts/ci/utils/publish_traces.py index a10f7ef10..bb9e6a6ba 100644 --- a/scripts/ci/utils/publish_traces.py +++ b/scripts/ci/utils/publish_traces.py @@ -311,7 +311,6 @@ def copy_trace_files(source_dir, target_base_path): for root, dirs, files in os.walk(source_dir): for file in files: if file.endswith(".json.gz"): - # Only upload TP rank 0 traces to avoid duplicates across tensor parallel ranks if "TP-" in file and "TP-0" not in file: continue diff --git a/scripts/ci/utils/query_job_status.py b/scripts/ci/utils/query_job_status.py index e4ddafa13..097c3e5b8 100755 --- a/scripts/ci/utils/query_job_status.py +++ b/scripts/ci/utils/query_job_status.py @@ -321,7 +321,7 @@ def fetch_all_jobs_snapshot( for i, run in enumerate(unique_runs): if (i + 1) % 20 == 0: - print(f"Processing run {i+1}/{total_runs}...", file=sys.stderr) + print(f"Processing run {i + 1}/{total_runs}...", file=sys.stderr) workflow_name = run.get("_workflow", "-") try: diff --git a/scripts/ci/utils/slash_command_handler.py b/scripts/ci/utils/slash_command_handler.py index a4fe6e42c..5e288c367 100644 --- a/scripts/ci/utils/slash_command_handler.py +++ b/scripts/ci/utils/slash_command_handler.py @@ -675,7 +675,7 @@ def resolve_test_file(file_part): full_path = ( file_part if file_part.startswith("python/") - else f"python/sglang/multimodal_gen/test/{file_part[len(prefix):]}" + else f"python/sglang/multimodal_gen/test/{file_part[len(prefix) :]}" ) if not os.path.isfile(full_path): return None, False, f"File not found: `{full_path}`" @@ -1207,9 +1207,8 @@ def handle_rerun_test( ): return False - if ( - refresh_precision_baseline - and not _check_precision_baseline_refresh_permissions(gh_repo, pr, comment) + if refresh_precision_baseline and not _check_precision_baseline_refresh_permissions( + gh_repo, pr, comment ): return False diff --git a/scripts/ci/utils/xpu_job_monitor.py b/scripts/ci/utils/xpu_job_monitor.py index 0053673dd..fec0bd487 100755 --- a/scripts/ci/utils/xpu_job_monitor.py +++ b/scripts/ci/utils/xpu_job_monitor.py @@ -323,7 +323,7 @@ def fetch_all_jobs_snapshot( for i, run in enumerate(unique_runs): if (i + 1) % 20 == 0: - print(f"Processing run {i+1}/{total_runs}...", file=sys.stderr) + print(f"Processing run {i + 1}/{total_runs}...", file=sys.stderr) workflow_name = run.get("_workflow", "-") try: diff --git a/scripts/ci_monitor/ci_failures_analysis.py b/scripts/ci_monitor/ci_failures_analysis.py index d06532801..9f40103f6 100644 --- a/scripts/ci_monitor/ci_failures_analysis.py +++ b/scripts/ci_monitor/ci_failures_analysis.py @@ -709,9 +709,9 @@ class SGLangFailuresAnalyzer: runner_instance_key = f"{runner_labels_str}_{runner_id}" runner_instance_stats[runner_instance_key]["total_jobs"] += 1 # Store runner name for reference - runner_instance_stats[runner_instance_key][ - "runner_name" - ] = runner_name + runner_instance_stats[runner_instance_key]["runner_name"] = ( + runner_name + ) # Calculate queue time (time from created to started) per instance created_at = job.get("created_at") diff --git a/scripts/code_sync/copy_to_oss.py b/scripts/code_sync/copy_to_oss.py index 96bc0af25..22579e85e 100644 --- a/scripts/code_sync/copy_to_oss.py +++ b/scripts/code_sync/copy_to_oss.py @@ -157,9 +157,7 @@ def get_oss_repo(dry_run): """ gh_token = os.getenv("GH_TOKEN") if not gh_token: - print( - "⚠️ Warning: GH_TOKEN environment variable not set. Skipping PR creation." - ) + print("⚠️ Warning: GH_TOKEN environment variable not set. Skipping PR creation.") if not dry_run: return @@ -361,9 +359,7 @@ def create_pull_request(oss_root, branch_name, title, body, dry_run): """Create a pull request in the OSS repo using the GitHub CLI.""" gh_token = os.getenv("GH_TOKEN") if not gh_token: - print( - "⚠️ Warning: GH_TOKEN environment variable not set. Skipping PR creation." - ) + print("⚠️ Warning: GH_TOKEN environment variable not set. Skipping PR creation.") if not dry_run: return @@ -569,8 +565,7 @@ def main(): "Please check the CI logs for the full patch and conflict details." ) pr_body_parts.append( - f"\n\n---\n\n" - f"*This is an automated PR created by scripts/copy_to_oss.py.*" + f"\n\n---\n\n*This is an automated PR created by scripts/copy_to_oss.py.*" ) pr_body = "\n".join(pr_body_parts) diff --git a/scripts/code_sync/utils.py b/scripts/code_sync/utils.py index 31b3cc067..6485a6d7f 100644 --- a/scripts/code_sync/utils.py +++ b/scripts/code_sync/utils.py @@ -120,14 +120,11 @@ def find_latest_oss_sync_commit(repo_root: Optional[str] = None) -> Optional[str if m: oss_commit = m.group(1) print( - f"✅ Latest OSS sync commit found: {oss_commit} " - f"(from: {subject})" + f"✅ Latest OSS sync commit found: {oss_commit} (from: {subject})" ) return oss_commit - print( - "⚠️ No '[Automated PR] Copy OSS code from commit ...' " "found in history." - ) + print("⚠️ No '[Automated PR] Copy OSS code from commit ...' found in history.") return None except subprocess.CalledProcessError as e: diff --git a/scripts/convert_otel_2_perfetto.py b/scripts/convert_otel_2_perfetto.py index 39ad3f78e..455fbb6fd 100644 --- a/scripts/convert_otel_2_perfetto.py +++ b/scripts/convert_otel_2_perfetto.py @@ -237,7 +237,7 @@ def generate_perfetto_span(engine_root_spans, smg_otel_spans, thread_meta_data): for thread_span in root_span["child"]: pid = int(thread_span["attributes"]["pid"]) host_id = thread_span["attributes"]["host_id"] - thread_name = f'{thread_span["attributes"]["host_id"][:8]}:{thread_span["attributes"]["thread_label"]}' + thread_name = f"{thread_span['attributes']['host_id'][:8]}:{thread_span['attributes']['thread_label']}" if "pp_rank" in thread_span["attributes"]: thread_name += f"-PP{thread_span['attributes']['pp_rank']}" if "dp_rank" in thread_span["attributes"]: diff --git a/scripts/lint/check_registered_tests.py b/scripts/lint/check_registered_tests.py index 1a33a977e..32828e802 100755 --- a/scripts/lint/check_registered_tests.py +++ b/scripts/lint/check_registered_tests.py @@ -156,11 +156,7 @@ def main() -> int: "test silently never runs. Switch to the modern form:\n" ) for f, suite in non_dispatchable: - print( - f" {f}\n" - f' suite="{suite}"' - f' -> stage="...", runner_config="..."' - ) + print(f' {f}\n suite="{suite}" -> stage="...", runner_config="..."') print() exit_code = 1 if dead_tests: diff --git a/scripts/playground/router/test_tree.py b/scripts/playground/router/test_tree.py index af41c738e..b2432860a 100644 --- a/scripts/playground/router/test_tree.py +++ b/scripts/playground/router/test_tree.py @@ -189,7 +189,7 @@ class TestMultiTenantRadixTree(unittest.TestCase): for i in range(100): for j, prefix in enumerate(prefixes): random_suffix = "".join(random.choices(string.ascii_letters, k=10)) - self.tree.insert(prefix + random_suffix, f"tenant{j+1}") + self.tree.insert(prefix + random_suffix, f"tenant{j + 1}") sizes_before = self.tree.get_used_size_per_tenant() print(sizes_before) diff --git a/scripts/release/update_others_whl_index.py b/scripts/release/update_others_whl_index.py index 9599697c6..c9059d5b4 100755 --- a/scripts/release/update_others_whl_index.py +++ b/scripts/release/update_others_whl_index.py @@ -84,7 +84,7 @@ def update_others_index( return False index_dir.mkdir(parents=True, exist_ok=True) - updated = f"{OTHERS_HEADER}{entry}{content[len(OTHERS_HEADER):]}" + updated = f"{OTHERS_HEADER}{entry}{content[len(OTHERS_HEADER) :]}" index_path.write_text(updated, encoding="utf-8") return True diff --git a/sgl-model-gateway/bindings/python/src/sglang_router/mini_lb.py b/sgl-model-gateway/bindings/python/src/sglang_router/mini_lb.py index 2156a2dfe..cb8feefad 100644 --- a/sgl-model-gateway/bindings/python/src/sglang_router/mini_lb.py +++ b/sgl-model-gateway/bindings/python/src/sglang_router/mini_lb.py @@ -131,7 +131,6 @@ class MiniLoadBalancer: total=self.timeout ) # Add timeout for request reliability ) as session: - tasks = [ session.post(f"{prefill_server}/{endpoint}", json=prefill_req), session.post(f"{decode_server}/{endpoint}", json=decode_req), @@ -141,7 +140,6 @@ class MiniLoadBalancer: prefill_response, decode_response = await asyncio.gather(*tasks) if "return_logprob" in modified_request: - prefill_json = await prefill_response.json() ret_json = await decode_response.json() diff --git a/sgl-model-gateway/bindings/python/src/sglang_router/router_args.py b/sgl-model-gateway/bindings/python/src/sglang_router/router_args.py index d02a75992..a2ea4f2ea 100644 --- a/sgl-model-gateway/bindings/python/src/sglang_router/router_args.py +++ b/sgl-model-gateway/bindings/python/src/sglang_router/router_args.py @@ -1069,7 +1069,6 @@ class RouterArgs: prefill_urls = [] for prefill_args in prefill_list: - url = prefill_args[0] # Handle optional bootstrap port diff --git a/sgl-model-gateway/bindings/python/tests/test_startup_sequence.py b/sgl-model-gateway/bindings/python/tests/test_startup_sequence.py index 3f6086bf1..e2b9567d2 100644 --- a/sgl-model-gateway/bindings/python/tests/test_startup_sequence.py +++ b/sgl-model-gateway/bindings/python/tests/test_startup_sequence.py @@ -788,7 +788,7 @@ def test_find_available_ports_and_wait_health(monkeypatch): monkeypatch.setattr( ls.time, "perf_counter", - lambda: (base.__setitem__("t", base["t"] + 0.1) or base["t"]), + lambda: base.__setitem__("t", base["t"] + 0.1) or base["t"], ) assert ls.wait_for_server_health("127.0.0.1", 12345, timeout=1) @@ -922,7 +922,6 @@ def test_launch_server_process_declares_on_a_resolved_record(monkeypatch): ) with patch("sglang_router.launch_router.logger") as mock_logger: - with pytest.raises( ValueError, match="PD disaggregation mode requires --prefill" ): diff --git a/sgl-model-gateway/e2e_test/chat_completions/test_enable_thinking.py b/sgl-model-gateway/e2e_test/chat_completions/test_enable_thinking.py index 79523268b..7b1dcee24 100644 --- a/sgl-model-gateway/e2e_test/chat_completions/test_enable_thinking.py +++ b/sgl-model-gateway/e2e_test/chat_completions/test_enable_thinking.py @@ -120,9 +120,9 @@ class TestEnableThinking: if "content" in delta and delta["content"]: has_content = True - assert ( - has_reasoning - ), "The reasoning content is not included in the stream response" + assert has_reasoning, ( + "The reasoning content is not included in the stream response" + ) assert has_content, "The stream response does not contain normal content" def test_stream_chat_completion_without_reasoning(self, setup_backend): @@ -162,7 +162,7 @@ class TestEnableThinking: if "content" in delta and delta["content"]: has_content = True - assert ( - not has_reasoning - ), "The reasoning content should not be included in the stream response" + assert not has_reasoning, ( + "The reasoning content should not be included in the stream response" + ) assert has_content, "The stream response does not contain normal content" diff --git a/sgl-model-gateway/e2e_test/chat_completions/test_function_calling.py b/sgl-model-gateway/e2e_test/chat_completions/test_function_calling.py index 94d145e11..73f198d2b 100644 --- a/sgl-model-gateway/e2e_test/chat_completions/test_function_calling.py +++ b/sgl-model-gateway/e2e_test/chat_completions/test_function_calling.py @@ -153,9 +153,9 @@ class TestOpenAIServerFunctionCalling: tool_calls = response.choices[0].message.tool_calls - assert ( - isinstance(tool_calls, list) and len(tool_calls) > 0 - ), "tool_calls should be a non-empty list" + assert isinstance(tool_calls, list) and len(tool_calls) > 0, ( + "tool_calls should be a non-empty list" + ) function_name = tool_calls[0].function.name assert function_name == "add", "Function name should be 'add'" @@ -221,20 +221,20 @@ class TestOpenAIServerFunctionCalling: if choice.delta.tool_calls: tool_call = choice.delta.tool_calls[0] if tool_call.function.name: - assert ( - tool_call.function.name == "get_current_weather" - ), "Function name should be 'get_current_weather'" + assert tool_call.function.name == "get_current_weather", ( + "Function name should be 'get_current_weather'" + ) found_function_name = True break - assert ( - found_function_name - ), "Target function name 'get_current_weather' was not found in the streaming chunks" + assert found_function_name, ( + "Target function name 'get_current_weather' was not found in the streaming chunks" + ) finish_reason = chunks[-1].choices[0].finish_reason - assert ( - finish_reason == "tool_calls" - ), "Final response of function calling should have finish_reason 'tool_calls'" + assert finish_reason == "tool_calls", ( + "Final response of function calling should have finish_reason 'tool_calls'" + ) def test_function_calling_streaming_args_parsing(self, setup_backend): """Test: Whether the function call arguments returned in streaming mode can be correctly concatenated into valid JSON. @@ -299,14 +299,14 @@ class TestOpenAIServerFunctionCalling: assert function_name == "add", "Function name should be 'add'" joined_args = "".join(argument_fragments) - assert ( - len(joined_args) > 0 - ), "No parameter fragments were returned in the function call" + assert len(joined_args) > 0, ( + "No parameter fragments were returned in the function call" + ) finish_reason = chunks[-1].choices[0].finish_reason - assert ( - finish_reason == "tool_calls" - ), "Final response of function calling should have finish_reason 'tool_calls'" + assert finish_reason == "tool_calls", ( + "Final response of function calling should have finish_reason 'tool_calls'" + ) # Check whether the concatenated JSON is valid try: @@ -445,21 +445,21 @@ class TestOpenAIServerFunctionCalling: arguments = tool_calls[0].function.arguments args_obj = json.loads(arguments) - assert ( - function_name == "get_weather" - ), f"Function name should be 'get_weather', got: {function_name}" - assert ( - "city" in args_obj - ), f"Function arguments should have 'city', got: {args_obj}" + assert function_name == "get_weather", ( + f"Function name should be 'get_weather', got: {function_name}" + ) + assert "city" in args_obj, ( + f"Function arguments should have 'city', got: {args_obj}" + ) # Make the test more robust by checking type and accepting valid responses city_value = args_obj["city"] - assert isinstance( - city_value, str - ), f"Parameter city should be a string, got: {type(city_value)}" - assert ( - "Paris" in city_value or "France" in city_value - ), f"Parameter city should contain either 'Paris' or 'France', got: {city_value}" + assert isinstance(city_value, str), ( + f"Parameter city should be a string, got: {type(city_value)}" + ) + assert "Paris" in city_value or "France" in city_value, ( + f"Parameter city should contain either 'Paris' or 'France', got: {city_value}" + ) def test_function_call_specific(self, setup_backend): """Test: Whether tool_choice: ToolChoice works as expected. @@ -592,9 +592,9 @@ class TestOpenAIServerFunctionCalling: finish_reason_chunks[index].append(choice.finish_reason) # Verify we got finish_reason chunks for both indices - assert ( - len(finish_reason_chunks) == 2 - ), f"Expected finish_reason chunks for 2 indices, got {len(finish_reason_chunks)}" + assert len(finish_reason_chunks) == 2, ( + f"Expected finish_reason chunks for 2 indices, got {len(finish_reason_chunks)}" + ) # Verify both index 0 and 1 have finish_reason assert 0 in finish_reason_chunks, "Missing finish_reason chunk for index 0" @@ -602,9 +602,9 @@ class TestOpenAIServerFunctionCalling: # Verify the finish_reason is "tool_calls" since we forced tool calls for index, reasons in finish_reason_chunks.items(): - assert ( - reasons[-1] == "tool_calls" - ), f"Expected finish_reason 'tool_calls' for index {index}, got {reasons[-1]}" + assert reasons[-1] == "tool_calls", ( + f"Expected finish_reason 'tool_calls' for index {index}, got {reasons[-1]}" + ) def test_function_calling_streaming_no_tool_call(self, setup_backend): """Test: Whether the finish_reason is stop in streaming mode when no tool call is given. @@ -663,14 +663,14 @@ class TestOpenAIServerFunctionCalling: found_tool_call = True break - assert ( - not found_tool_call - ), "Shouldn't have any tool_call in the streaming chunks" + assert not found_tool_call, ( + "Shouldn't have any tool_call in the streaming chunks" + ) finish_reason = chunks[-1].choices[0].finish_reason - assert ( - finish_reason == "stop" - ), "Final response of no function calling should have finish_reason 'stop'" + assert finish_reason == "stop", ( + "Final response of no function calling should have finish_reason 'stop'" + ) def test_streaming_multiple_choices_without_tools(self, setup_backend): """Test: Verify that each choice gets its own finish_reason chunk without tool calls. @@ -705,9 +705,9 @@ class TestOpenAIServerFunctionCalling: finish_reason_chunks[index].append(choice.finish_reason) # Verify we got finish_reason chunks for both indices - assert ( - len(finish_reason_chunks) == 2 - ), f"Expected finish_reason chunks for 2 indices, got {len(finish_reason_chunks)}" + assert len(finish_reason_chunks) == 2, ( + f"Expected finish_reason chunks for 2 indices, got {len(finish_reason_chunks)}" + ) # Verify both index 0 and 1 have finish_reason assert 0 in finish_reason_chunks, "Missing finish_reason chunk for index 0" @@ -718,7 +718,9 @@ class TestOpenAIServerFunctionCalling: assert reasons[-1] in [ "stop", "length", - ], f"Expected finish_reason 'stop' or 'length' for index {index}, got {reasons[-1]}" + ], ( + f"Expected finish_reason 'stop' or 'length' for index {index}, got {reasons[-1]}" + ) # ============================================================================= @@ -749,9 +751,9 @@ class TestOpenAIPythonicFunctionCalling: assert isinstance(tool_calls, list), "No tool_calls found" assert len(tool_calls) >= 1 names = [tc.function.name for tc in tool_calls] - assert ( - "get_weather" in names or "get_tourist_attractions" in names - ), f"Function name '{names}' should contain either 'get_weather' or 'get_tourist_attractions'" + assert "get_weather" in names or "get_tourist_attractions" in names, ( + f"Function name '{names}' should contain either 'get_weather' or 'get_tourist_attractions'" + ) def test_pythonic_tool_call_streaming(self, setup_backend): """Test: Streaming pythonic tool call format; assert tool_call index is present.""" @@ -782,7 +784,9 @@ class TestOpenAIPythonicFunctionCalling: assert found_index, "No index field found in any streamed tool_call" assert ( "get_weather" in found_names or "get_tourist_attractions" in found_names - ), f"Function name '{found_names}' should contain either 'get_weather' or 'get_tourist_attractions'" + ), ( + f"Function name '{found_names}' should contain either 'get_weather' or 'get_tourist_attractions'" + ) # ============================================================================= @@ -1132,13 +1136,13 @@ class _TestToolChoiceBase: tool_call["id"] = tool_call_delta.id if tool_call_delta.function: if tool_call_delta.function.name: - tool_call["function"][ - "name" - ] = tool_call_delta.function.name + tool_call["function"]["name"] = ( + tool_call_delta.function.name + ) if tool_call_delta.function.arguments: - tool_call["function"][ - "arguments" - ] += tool_call_delta.function.arguments + tool_call["function"]["arguments"] += ( + tool_call_delta.function.arguments + ) assert len(tool_calls_by_index) > 0 @@ -1271,9 +1275,9 @@ class _TestToolChoiceBase: assert len(tool_calls) == 2, f"Expected 2 tool calls, got {len(tool_calls)}" called_functions = {call.function.name for call in tool_calls} - assert ( - called_functions == expected_functions - ), f"Expected functions {expected_functions}, got {called_functions}" + assert called_functions == expected_functions, ( + f"Expected functions {expected_functions}, got {called_functions}" + ) def test_multi_tool_scenario_required(self, setup_backend): """Test multi-tool scenario with tool_choice='required'.""" @@ -1307,9 +1311,9 @@ class _TestToolChoiceBase: if self._is_flaky_test("test_multi_tool_scenario_required"): # For flaky tests, just ensure basic functionality works - assert ( - len(tool_calls) > 0 - ), f"Expected at least 1 tool call, got {len(tool_calls)}" + assert len(tool_calls) > 0, ( + f"Expected at least 1 tool call, got {len(tool_calls)}" + ) for call in tool_calls: assert call.function.name in available_names else: @@ -1317,9 +1321,9 @@ class _TestToolChoiceBase: assert len(tool_calls) == 2, f"Expected 2 tool calls, got {len(tool_calls)}" called_functions = {call.function.name for call in tool_calls} - assert ( - called_functions == expected_functions - ), f"Expected functions {expected_functions}, got {called_functions}" + assert called_functions == expected_functions, ( + f"Expected functions {expected_functions}, got {called_functions}" + ) def test_error_handling_invalid_tool_choice(self, setup_backend): """Test error handling for invalid tool_choice.""" diff --git a/sgl-model-gateway/e2e_test/chat_completions/test_openai_server.py b/sgl-model-gateway/e2e_test/chat_completions/test_openai_server.py index 676f9d98a..6a44500f0 100644 --- a/sgl-model-gateway/e2e_test/chat_completions/test_openai_server.py +++ b/sgl-model-gateway/e2e_test/chat_completions/test_openai_server.py @@ -172,9 +172,9 @@ The SmartHome Mini is a compact smart home assistant available in black or white ret_num_top_logprobs = len( response.choices[0].logprobs.content[0].top_logprobs ) - assert ( - ret_num_top_logprobs == logprobs - ), f"{ret_num_top_logprobs} vs {logprobs}" + assert ret_num_top_logprobs == logprobs, ( + f"{ret_num_top_logprobs} vs {logprobs}" + ) assert len(response.choices) == parallel_sample_num assert response.choices[0].message.role == "assistant" @@ -223,9 +223,9 @@ The SmartHome Mini is a compact smart home assistant available in black or white data = response.choices[0].delta if is_firsts.get(index, True): - assert ( - data.role == "assistant" - ), "data.role was not 'assistant' for first chunk" + assert data.role == "assistant", ( + "data.role was not 'assistant' for first chunk" + ) is_firsts[index] = False continue @@ -240,9 +240,9 @@ The SmartHome Mini is a compact smart home assistant available in black or white ret_num_top_logprobs = len( response.choices[0].logprobs.content[0].top_logprobs ) - assert ( - ret_num_top_logprobs == logprobs - ), f"{ret_num_top_logprobs} vs {logprobs}" + assert ret_num_top_logprobs == logprobs, ( + f"{ret_num_top_logprobs} vs {logprobs}" + ) assert ( isinstance(data.content, str) @@ -254,14 +254,14 @@ The SmartHome Mini is a compact smart home assistant available in black or white assert response.created for index in range(parallel_sample_num): - assert not is_firsts.get( - index, True - ), f"index {index} is not found in the response" + assert not is_firsts.get(index, True), ( + f"index {index} is not found in the response" + ) for index in range(parallel_sample_num): - assert ( - index in finish_reason_counts - ), f"No finish_reason found for index {index}" + assert index in finish_reason_counts, ( + f"No finish_reason found for index {index}" + ) assert finish_reason_counts[index] == 1, ( f"Expected 1 finish_reason chunk for index {index}, " f"got {finish_reason_counts[index]}" diff --git a/sgl-model-gateway/e2e_test/chat_completions/test_validation.py b/sgl-model-gateway/e2e_test/chat_completions/test_validation.py index c6412ab0c..ed6355159 100644 --- a/sgl-model-gateway/e2e_test/chat_completions/test_validation.py +++ b/sgl-model-gateway/e2e_test/chat_completions/test_validation.py @@ -91,9 +91,9 @@ class TestIgnoreEOS: # The ignore_eos response should either: # 1. Have more tokens than the default response (if default stopped at EOS before max_tokens) # 2. Have exactly max_tokens (if it reached the max_tokens limit) - assert ( - ignore_eos_tokens > default_tokens or ignore_eos_tokens >= max_tokens - ), f"ignore_eos did not generate more tokens: {ignore_eos_tokens} vs {default_tokens}" + assert ignore_eos_tokens > default_tokens or ignore_eos_tokens >= max_tokens, ( + f"ignore_eos did not generate more tokens: {ignore_eos_tokens} vs {default_tokens}" + ) assert response_ignore_eos.choices[0].finish_reason == "length", ( f"Expected finish_reason='length' for ignore_eos=True, " @@ -158,9 +158,9 @@ class TestLargeMaxNewTokens: # Verify all requests completed successfully assert len(responses) == num_requests for i, response in enumerate(responses): - assert response.choices[ - 0 - ].message.content, f"Request {i} returned empty content" + assert response.choices[0].message.content, ( + f"Request {i} returned empty content" + ) assert response.choices[0].finish_reason in ("stop", "length"), ( f"Request {i} had unexpected finish_reason: " f"{response.choices[0].finish_reason}" diff --git a/sgl-model-gateway/e2e_test/embeddings/test_correctness.py b/sgl-model-gateway/e2e_test/embeddings/test_correctness.py index 0a7162cfd..05e072313 100644 --- a/sgl-model-gateway/e2e_test/embeddings/test_correctness.py +++ b/sgl-model-gateway/e2e_test/embeddings/test_correctness.py @@ -239,9 +239,9 @@ class TestEmbeddingCorrectness: # Verify all similarities are close to 1.0 for j, sim in enumerate(similarities): - assert ( - abs(sim - 1.0) < tolerance - ), f"Set {i+1}, text {j+1}: similarity {sim:.4f} not close to 1.0" + assert abs(sim - 1.0) < tolerance, ( + f"Set {i + 1}, text {j + 1}: similarity {sim:.4f} not close to 1.0" + ) logger.info("Semantic similarity test set %d passed", i + 1) @@ -273,8 +273,8 @@ class TestEmbeddingCorrectness: logger.info("Gateway relevance scores: %s", scores_gateway) logger.info("HF relevance scores: %s", scores_hf) - assert np.allclose( - scores_gateway, scores_hf, atol=tolerance - ), f"Scores differ beyond tolerance:\nGateway: {scores_gateway}\nHF: {scores_hf}" + assert np.allclose(scores_gateway, scores_hf, atol=tolerance), ( + f"Scores differ beyond tolerance:\nGateway: {scores_gateway}\nHF: {scores_hf}" + ) logger.info("Relevance scores comparison passed") diff --git a/sgl-model-gateway/e2e_test/infra/gpu_monitor.py b/sgl-model-gateway/e2e_test/infra/gpu_monitor.py index 870eec2b0..822e9a2f4 100644 --- a/sgl-model-gateway/e2e_test/infra/gpu_monitor.py +++ b/sgl-model-gateway/e2e_test/infra/gpu_monitor.py @@ -300,17 +300,17 @@ class GPUMonitor: mean_threshold = thresholds.get("gpu_util_mean_min") if mean_threshold is not None: mean_value = overall.get("mean", 0.0) - assert ( - mean_value >= mean_threshold - ), f"GPU utilization mean below threshold: {mean_value:.2f}% < {mean_threshold}%" + assert mean_value >= mean_threshold, ( + f"GPU utilization mean below threshold: {mean_value:.2f}% < {mean_threshold}%" + ) p50_threshold = thresholds.get("gpu_util_p50_min") if p50_threshold is not None: p50_value = overall.get("p50") if p50_value is not None: - assert ( - p50_value >= p50_threshold - ), f"GPU utilization p50 below threshold: {p50_value:.2f}% < {p50_threshold}%" + assert p50_value >= p50_threshold, ( + f"GPU utilization p50 below threshold: {p50_value:.2f}% < {p50_threshold}%" + ) def should_monitor(thresholds: dict[str, Any] | None) -> bool: diff --git a/sgl-model-gateway/e2e_test/k8s_integration/test_cross_namespace.py b/sgl-model-gateway/e2e_test/k8s_integration/test_cross_namespace.py index 504a7afdd..a596433a5 100644 --- a/sgl-model-gateway/e2e_test/k8s_integration/test_cross_namespace.py +++ b/sgl-model-gateway/e2e_test/k8s_integration/test_cross_namespace.py @@ -255,12 +255,12 @@ class TestClusterWideDiscovery: # A regression that quietly hardcoded a namespace filter would # still produce total=2 if labels happened to match elsewhere, # but only one of these IPs would surface. - assert any( - ip_a in u for u in urls - ), f"worker_a IP {ip_a} (ns {NAMESPACE}) not in {urls}" - assert any( - ip_b in u for u in urls - ), f"worker_b IP {ip_b} (ns {EXTRA_NAMESPACE}) not in {urls}" + assert any(ip_a in u for u in urls), ( + f"worker_a IP {ip_a} (ns {NAMESPACE}) not in {urls}" + ) + assert any(ip_b in u for u in urls), ( + f"worker_b IP {ip_b} (ns {EXTRA_NAMESPACE}) not in {urls}" + ) finally: _safe_delete_pod(worker_a, NAMESPACE) _safe_delete_pod(worker_b, EXTRA_NAMESPACE) diff --git a/sgl-model-gateway/e2e_test/k8s_integration/test_lifecycle.py b/sgl-model-gateway/e2e_test/k8s_integration/test_lifecycle.py index 0ab7022ba..8b63dccec 100644 --- a/sgl-model-gateway/e2e_test/k8s_integration/test_lifecycle.py +++ b/sgl-model-gateway/e2e_test/k8s_integration/test_lifecycle.py @@ -256,9 +256,9 @@ class TestGatewayRestart: "-l", "app=smg-gateway-restart", ) - assert res.get( - "items" - ), "No pods found for selector app=smg-gateway-restart" + assert res.get("items"), ( + "No pods found for selector app=smg-gateway-restart" + ) old_pod = res["items"][0]["metadata"]["name"] _kubectl( "delete", @@ -304,9 +304,9 @@ class TestGatewayRestart: logger.info("Workers after restart: %s", urls_after) # No duplicates: each pod should appear exactly once. - assert len(urls_after) == len( - set(urls_after) - ), f"Duplicate worker registrations after gateway restart: {urls_after}" + assert len(urls_after) == len(set(urls_after)), ( + f"Duplicate worker registrations after gateway restart: {urls_after}" + ) # Set equality: the same workers come back, neither dropped # nor duplicated. assert set(urls_after) == set(urls_before), ( @@ -364,9 +364,9 @@ class TestPodIpChange: ip_before = _get_pod_ip(pod_name) urls_before = {w["url"] for w in _get_workers(gateway_url)["workers"]} - assert any( - ip_before in url for url in urls_before - ), f"Expected initial worker URL containing {ip_before}, got {urls_before}" + assert any(ip_before in url for url in urls_before), ( + f"Expected initial worker URL containing {ip_before}, got {urls_before}" + ) logger.info("Pod IP before: %s, urls: %s", ip_before, urls_before) # Force-delete and wait until the registry no longer references @@ -384,9 +384,11 @@ class TestPodIpChange: ) _wait_for_pod_gone(pod_name) _poll_until( - lambda: not any( - ip_before in w["url"] - for w in _get_workers(gateway_url).get("workers", []) + lambda: ( + not any( + ip_before in w["url"] + for w in _get_workers(gateway_url).get("workers", []) + ) ), f"stale worker for IP {ip_before} removed", timeout=RECONCILIATION_WAIT_SECS, @@ -431,9 +433,9 @@ class TestPodIpChange: f"Expected exactly one worker URL containing current IP " f"{ip_after}, got {matching_after} (all urls: {urls_after})" ) - assert not any( - ip_before in u for u in urls_after - ), f"Stale URL with old IP {ip_before} still in registry: {urls_after}" + assert not any(ip_before in u for u in urls_after), ( + f"Stale URL with old IP {ip_before} still in registry: {urls_after}" + ) finally: _safe_force_delete(pod_name) @@ -496,9 +498,11 @@ class TestGracefulDrain: # without affecting whether *this* pod's IP got removed. The # meaningful timing guarantee (`elapsed < grace_secs`) is below. _poll_until( - lambda: not any( - pod_ip in w["url"] - for w in _get_workers(gateway_url).get("workers", []) + lambda: ( + not any( + pod_ip in w["url"] + for w in _get_workers(gateway_url).get("workers", []) + ) ), f"worker for ip {pod_ip} deregistered after graceful delete", timeout=RECONCILIATION_WAIT_SECS, diff --git a/sgl-model-gateway/e2e_test/k8s_integration/test_multi_model.py b/sgl-model-gateway/e2e_test/k8s_integration/test_multi_model.py index f0cce529e..269570b2f 100644 --- a/sgl-model-gateway/e2e_test/k8s_integration/test_multi_model.py +++ b/sgl-model-gateway/e2e_test/k8s_integration/test_multi_model.py @@ -196,9 +196,9 @@ class TestMultiModelSelectorIsolation: # No URL should appear in both views — that would mean a # selector mismatch leaked a worker into the wrong gateway. cross_talk = set(llama_urls) & set(qwen_urls) - assert ( - not cross_talk - ), f"Workers leaked across model selectors: {cross_talk}" + assert not cross_talk, ( + f"Workers leaked across model selectors: {cross_talk}" + ) finally: for name in llama_workers + qwen_workers: _safe_force_delete(name) diff --git a/sgl-model-gateway/e2e_test/k8s_integration/test_pd_type_change.py b/sgl-model-gateway/e2e_test/k8s_integration/test_pd_type_change.py index 4a82bef20..6a45fa142 100644 --- a/sgl-model-gateway/e2e_test/k8s_integration/test_pd_type_change.py +++ b/sgl-model-gateway/e2e_test/k8s_integration/test_pd_type_change.py @@ -181,9 +181,9 @@ class TestPDRolloutTypeChange: by_type = _get_workers_by_type(pd_gateway) logger.info("Workers by type: %s", json.dumps(by_type, indent=2)) - assert ( - "prefill" in by_type - ), f"Expected prefill, got: {list(by_type.keys())}" + assert "prefill" in by_type, ( + f"Expected prefill, got: {list(by_type.keys())}" + ) finally: _safe_delete_pod(pod_name) @@ -257,12 +257,12 @@ class TestPDRolloutTypeChange: by_type = _get_workers_by_type(pd_gateway) logger.info("After rollout: %s", json.dumps(by_type, indent=2)) - assert ( - "decode" in by_type - ), f"Expected decode worker after rollout, got: {list(by_type.keys())}" - assert ( - "prefill" not in by_type - ), "Stale prefill worker persists after rollout" + assert "decode" in by_type, ( + f"Expected decode worker after rollout, got: {list(by_type.keys())}" + ) + assert "prefill" not in by_type, ( + "Stale prefill worker persists after rollout" + ) finally: _safe_delete_pod(pod_name) diff --git a/sgl-model-gateway/e2e_test/k8s_integration/test_reconciliation.py b/sgl-model-gateway/e2e_test/k8s_integration/test_reconciliation.py index f3bba5267..385002b03 100644 --- a/sgl-model-gateway/e2e_test/k8s_integration/test_reconciliation.py +++ b/sgl-model-gateway/e2e_test/k8s_integration/test_reconciliation.py @@ -377,9 +377,9 @@ class TestReconciliationMetrics: {"source": "kubernetes", "result": "success"}, ) logger.info("Registration success metric: %s", reg_value) - assert ( - reg_value is not None and reg_value >= 1 - ), f"Expected at least 1 registration, got {reg_value}" + assert reg_value is not None and reg_value >= 1, ( + f"Expected at least 1 registration, got {reg_value}" + ) gauge_value = _parse_metric_value( metrics_text, @@ -387,9 +387,9 @@ class TestReconciliationMetrics: {"source": "kubernetes"}, ) logger.info("Workers discovered gauge: %s", gauge_value) - assert ( - gauge_value is not None and gauge_value >= 1 - ), f"Expected workers_discovered >= 1, got {gauge_value}" + assert gauge_value is not None and gauge_value >= 1, ( + f"Expected workers_discovered >= 1, got {gauge_value}" + ) finally: _safe_delete_worker_pod(pod_name) @@ -497,9 +497,9 @@ class TestReconciliationConsistency: logger.info("Worker count samples over time: %s", samples) - assert all( - s == stable_count for s in samples - ), f"Worker count fluctuated: {samples} (expected stable at {stable_count})" + assert all(s == stable_count for s in samples), ( + f"Worker count fluctuated: {samples} (expected stable at {stable_count})" + ) finally: for name in pod_names: diff --git a/sgl-model-gateway/e2e_test/responses/test_streaming_events.py b/sgl-model-gateway/e2e_test/responses/test_streaming_events.py index 4aaaf0b0a..a0dce5e39 100644 --- a/sgl-model-gateway/e2e_test/responses/test_streaming_events.py +++ b/sgl-model-gateway/e2e_test/responses/test_streaming_events.py @@ -58,15 +58,15 @@ class TestStreamingEventsLocal: first_item_event = output_item_added_events[0] assert first_item_event.item is not None assert first_item_event.output_index is not None - assert ( - first_item_event.output_index == 0 - ), "First output item must have output_index: 0 (zero-based indexing)" + assert first_item_event.output_index == 0, ( + "First output item must have output_index: 0 (zero-based indexing)" + ) # Verify subsequent items increment correctly for i, event in enumerate(output_item_added_events): - assert ( - event.output_index == i - ), f"Output item {i} should have output_index: {i}" + assert event.output_index == i, ( + f"Output item {i} should have output_index: {i}" + ) # Verify output_item.done event exists output_item_done_events = [ @@ -101,9 +101,9 @@ class TestStreamingEventsLocal: output_item_added_events = [ event for event in events if event.type == "response.output_item.added" ] - assert len(output_item_added_events) == len( - output_array - ), "Number of output_item.added events should match output array length" + assert len(output_item_added_events) == len(output_array), ( + "Number of output_item.added events should match output array length" + ) # ============================================================================= @@ -147,15 +147,15 @@ class TestStreamingEventsHarmony: first_item_event = output_item_added_events[0] assert first_item_event.item is not None assert first_item_event.output_index is not None - assert ( - first_item_event.output_index == 0 - ), "First output item must have output_index: 0 (zero-based indexing)" + assert first_item_event.output_index == 0, ( + "First output item must have output_index: 0 (zero-based indexing)" + ) # Verify subsequent items increment correctly for i, event in enumerate(output_item_added_events): - assert ( - event.output_index == i - ), f"Output item {i} should have output_index: {i}" + assert event.output_index == i, ( + f"Output item {i} should have output_index: {i}" + ) # Verify output_item.done event exists output_item_done_events = [ @@ -190,9 +190,9 @@ class TestStreamingEventsHarmony: output_item_added_events = [ event for event in events if event.type == "response.output_item.added" ] - assert len(output_item_added_events) == len( - output_array - ), "Number of output_item.added events should match output array length" + assert len(output_item_added_events) == len(output_array), ( + "Number of output_item.added events should match output array length" + ) def test_reasoning_content(self, setup_backend): """Test that reasoning content has correct zero-based output_index. @@ -228,16 +228,16 @@ class TestStreamingEventsHarmony: # If reasoning is present, verify it has output_index: 0 if reasoning_items: reasoning_item = reasoning_items[0] - assert ( - reasoning_item.output_index == 0 - ), "Reasoning item should have output_index: 0" + assert reasoning_item.output_index == 0, ( + "Reasoning item should have output_index: 0" + ) # If message is present after reasoning, verify it has output_index: 1 if reasoning_items and message_items: message_item = message_items[0] - assert ( - message_item.output_index == 1 - ), "Message item after reasoning should have output_index: 1" + assert message_item.output_index == 1, ( + "Message item after reasoning should have output_index: 1" + ) # Find response.completed event completed_events = [ diff --git a/sgl-model-gateway/e2e_test/responses/test_tools_call.py b/sgl-model-gateway/e2e_test/responses/test_tools_call.py index d025fb52a..16a358694 100644 --- a/sgl-model-gateway/e2e_test/responses/test_tools_call.py +++ b/sgl-model-gateway/e2e_test/responses/test_tools_call.py @@ -168,9 +168,9 @@ class TestToolCallingCloud: # Check for function_call in output function_calls = [item for item in output if item.type == "function_call"] - assert ( - len(function_calls) > 0 - ), "Response should contain at least one function_call" + assert len(function_calls) > 0, ( + "Response should contain at least one function_call" + ) # Verify function_call structure function_call = function_calls[0] @@ -285,30 +285,30 @@ class TestToolCallingCloud: event_types = [event.type for event in events] assert "response.created" in event_types, "Should have response.created event" - assert ( - "response.completed" in event_types - ), "Should have response.completed event" - assert ( - "response.output_item.added" in event_types - ), "Should have output_item.added events" - assert ( - "response.mcp_list_tools.in_progress" in event_types - ), "Should have mcp_list_tools.in_progress event" - assert ( - "response.mcp_list_tools.completed" in event_types - ), "Should have mcp_list_tools.completed event" - assert ( - "response.mcp_call.in_progress" in event_types - ), "Should have mcp_call.in_progress event" - assert ( - "response.mcp_call_arguments.delta" in event_types - ), "Should have mcp_call_arguments.delta event" - assert ( - "response.mcp_call_arguments.done" in event_types - ), "Should have mcp_call_arguments.done event" - assert ( - "response.mcp_call.completed" in event_types - ), "Should have mcp_call.completed event" + assert "response.completed" in event_types, ( + "Should have response.completed event" + ) + assert "response.output_item.added" in event_types, ( + "Should have output_item.added events" + ) + assert "response.mcp_list_tools.in_progress" in event_types, ( + "Should have mcp_list_tools.in_progress event" + ) + assert "response.mcp_list_tools.completed" in event_types, ( + "Should have mcp_list_tools.completed event" + ) + assert "response.mcp_call.in_progress" in event_types, ( + "Should have mcp_call.in_progress event" + ) + assert "response.mcp_call_arguments.delta" in event_types, ( + "Should have mcp_call_arguments.delta event" + ) + assert "response.mcp_call_arguments.done" in event_types, ( + "Should have mcp_call_arguments.done event" + ) + assert "response.mcp_call.completed" in event_types, ( + "Should have mcp_call.completed event" + ) completed_events = [e for e in events if e.type == "response.completed"] assert len(completed_events) == 1 @@ -336,18 +336,18 @@ class TestToolCallingCloud: assert mcp_call.output is not None # Strict validation for cloud backends - check for text output events - assert ( - "response.content_part.added" in event_types - ), "Should have content_part.added event" - assert ( - "response.output_text.delta" in event_types - ), "Should have output_text.delta events" - assert ( - "response.output_text.done" in event_types - ), "Should have output_text.done event" - assert ( - "response.content_part.done" in event_types - ), "Should have content_part.done event" + assert "response.content_part.added" in event_types, ( + "Should have content_part.added event" + ) + assert "response.output_text.delta" in event_types, ( + "Should have output_text.delta events" + ) + assert "response.output_text.done" in event_types, ( + "Should have output_text.done event" + ) + assert "response.content_part.done" in event_types, ( + "Should have content_part.done event" + ) assert "message" in final_output_types @@ -400,9 +400,9 @@ class TestToolChoiceHarmony: assert len(output) > 0 function_calls = [item for item in output if item.type == "function_call"] - assert ( - len(function_calls) > 0 - ), "Model should choose to call function with tool_choice='auto'" + assert len(function_calls) > 0, ( + "Model should choose to call function with tool_choice='auto'" + ) def test_tool_choice_required(self, setup_backend): """Test tool_choice="required" forces the model to call at least one tool.""" @@ -423,9 +423,9 @@ class TestToolChoiceHarmony: output = resp.output function_calls = [item for item in output if item.type == "function_call"] - assert ( - len(function_calls) > 0 - ), "tool_choice='required' must force at least one function call" + assert len(function_calls) > 0, ( + "tool_choice='required' must force at least one function call" + ) def test_tool_choice_specific_function(self, setup_backend): """Test tool_choice with specific function name forces that function to be called.""" @@ -447,9 +447,9 @@ class TestToolChoiceHarmony: output = resp.output function_calls = [item for item in output if item.type == "function_call"] assert len(function_calls) > 0, "Must call the specified function" - assert ( - function_calls[0].name == "search_web" - ), "Must call the function specified in tool_choice" + assert function_calls[0].name == "search_web", ( + "Must call the function specified in tool_choice" + ) def test_tool_choice_streaming(self, setup_backend): """Test tool_choice parameter works correctly with streaming.""" diff --git a/sgl-model-gateway/e2e_test/router/test_mmlu.py b/sgl-model-gateway/e2e_test/router/test_mmlu.py index 770211125..dbbe9cf8d 100644 --- a/sgl-model-gateway/e2e_test/router/test_mmlu.py +++ b/sgl-model-gateway/e2e_test/router/test_mmlu.py @@ -47,9 +47,9 @@ class TestMMLU: ) metrics = run_eval(args) - assert ( - metrics["score"] >= 0.65 - ), f"MMLU score {metrics['score']:.2f} below threshold 0.65" + assert metrics["score"] >= 0.65, ( + f"MMLU score {metrics['score']:.2f} below threshold 0.65" + ) logger.info("MMLU score: %.2f (threshold: 0.65)", metrics["score"]) def test_mmlu_extended(self, setup_backend): @@ -70,7 +70,7 @@ class TestMMLU: ) metrics = run_eval(args) - assert ( - metrics["score"] >= 0.65 - ), f"MMLU score {metrics['score']:.2f} below threshold 0.65" + assert metrics["score"] >= 0.65, ( + f"MMLU score {metrics['score']:.2f} below threshold 0.65" + ) logger.info("MMLU extended score: %.2f (threshold: 0.65)", metrics["score"]) diff --git a/sgl-model-gateway/e2e_test/router/test_pd_mmlu.py b/sgl-model-gateway/e2e_test/router/test_pd_mmlu.py index bd2a4ae8f..0c2e0f368 100644 --- a/sgl-model-gateway/e2e_test/router/test_pd_mmlu.py +++ b/sgl-model-gateway/e2e_test/router/test_pd_mmlu.py @@ -55,7 +55,7 @@ class TestPDMMLU: ) metrics = run_eval(args) - assert ( - metrics["score"] >= 0.65 - ), f"PD MMLU score {metrics['score']:.2f} below threshold 0.65" + assert metrics["score"] >= 0.65, ( + f"PD MMLU score {metrics['score']:.2f} below threshold 0.65" + ) logger.info("PD MMLU score: %.2f (threshold: 0.65)", metrics["score"]) diff --git a/sgl-model-gateway/e2e_test/router/test_worker_api.py b/sgl-model-gateway/e2e_test/router/test_worker_api.py index 5fa187cfd..056302032 100644 --- a/sgl-model-gateway/e2e_test/router/test_worker_api.py +++ b/sgl-model-gateway/e2e_test/router/test_worker_api.py @@ -201,9 +201,9 @@ class TestDisableHealthCheck: worker.metadata.get("disable_health_check"), ) # Worker should be healthy immediately - assert ( - worker.status == "healthy" - ), "Worker should be healthy when health checks disabled" + assert worker.status == "healthy", ( + "Worker should be healthy when health checks disabled" + ) finally: gateway.shutdown() http_instance.release() diff --git a/test/manual/8-gpu-models/test_deepseek_v3_basic.py b/test/manual/8-gpu-models/test_deepseek_v3_basic.py index 1ae60098c..57c34c3a8 100644 --- a/test/manual/8-gpu-models/test_deepseek_v3_basic.py +++ b/test/manual/8-gpu-models/test_deepseek_v3_basic.py @@ -58,7 +58,7 @@ class TestDeepseekV3Basic(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k (deepseek-v3)\n" f'{metrics["score"]=:.3f}\n' + f'### test_gsm8k (deepseek-v3)\n{metrics["score"]=:.3f}\n' ) self.assertGreater(metrics["score"], 0.935) @@ -70,7 +70,7 @@ class TestDeepseekV3Basic(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_bs_1_speed (deepseek-v3)\n" f"{speed=:.2f} token/s\n" + f"### test_bs_1_speed (deepseek-v3)\n{speed=:.2f} token/s\n" ) if is_in_amd_ci(): self.assertGreater(speed, 12) diff --git a/test/manual/8-gpu-models/test_dsa_models_basic.py b/test/manual/8-gpu-models/test_dsa_models_basic.py index b8c90cab7..14dfde0c9 100644 --- a/test/manual/8-gpu-models/test_dsa_models_basic.py +++ b/test/manual/8-gpu-models/test_dsa_models_basic.py @@ -61,7 +61,7 @@ class TestDeepseekV32DP(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k (deepseek-v32)\n" f'{metrics["score"]=:.3f}\n' + f'### test_gsm8k (deepseek-v32)\n{metrics["score"]=:.3f}\n' ) self.assertGreater(metrics["score"], 0.935) @@ -73,7 +73,7 @@ class TestDeepseekV32DP(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_bs_1_speed (deepseek-v32)\n" f"{speed=:.2f} token/s\n" + f"### test_bs_1_speed (deepseek-v32)\n{speed=:.2f} token/s\n" ) self.assertGreater(speed, 50) @@ -119,7 +119,7 @@ class TestDeepseekV32TP(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k (deepseek-v32)\n" f'{metrics["score"]=:.3f}\n' + f'### test_gsm8k (deepseek-v32)\n{metrics["score"]=:.3f}\n' ) self.assertGreater(metrics["score"], 0.935) @@ -131,7 +131,7 @@ class TestDeepseekV32TP(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_bs_1_speed (deepseek-v32)\n" f"{speed=:.2f} token/s\n" + f"### test_bs_1_speed (deepseek-v32)\n{speed=:.2f} token/s\n" ) self.assertGreater(speed, 80) @@ -180,7 +180,7 @@ class TestGLM5DP(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k (glm-5)\n" f'{metrics["score"]=:.3f}\n' + f'### test_gsm8k (glm-5)\n{metrics["score"]=:.3f}\n' ) self.assertGreater(metrics["score"], 0.935) @@ -192,7 +192,7 @@ class TestGLM5DP(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_bs_1_speed (glm-5)\n" f"{speed=:.2f} token/s\n" + f"### test_bs_1_speed (glm-5)\n{speed=:.2f} token/s\n" ) self.assertGreater(speed, 40) @@ -238,7 +238,7 @@ class TestGLM5TP(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k (glm-5)\n" f'{metrics["score"]=:.3f}\n' + f'### test_gsm8k (glm-5)\n{metrics["score"]=:.3f}\n' ) self.assertGreater(metrics["score"], 0.935) @@ -250,7 +250,7 @@ class TestGLM5TP(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_bs_1_speed (glm-5)\n" f"{speed=:.2f} token/s\n" + f"### test_bs_1_speed (glm-5)\n{speed=:.2f} token/s\n" ) self.assertGreater(speed, 60) diff --git a/test/manual/ascend/disaggregation_utils.py b/test/manual/ascend/disaggregation_utils.py index cb46a9d3b..14e05785c 100644 --- a/test/manual/ascend/disaggregation_utils.py +++ b/test/manual/ascend/disaggregation_utils.py @@ -138,7 +138,7 @@ def get_rdma_devices_args(): if not (base_rdma_group <= gpu_idx < base_rdma_group + 4): warnings.warn( f"GPU index {gpu_idx} is outside expected group " - f"{base_rdma_group}-{base_rdma_group+3}" + f"{base_rdma_group}-{base_rdma_group + 3}" ) # 3. Generate RDMA device names diff --git a/test/manual/ascend/test_ascend_deepseek_mtp.py b/test/manual/ascend/test_ascend_deepseek_mtp.py index acc78fa5b..6e45a149e 100644 --- a/test/manual/ascend/test_ascend_deepseek_mtp.py +++ b/test/manual/ascend/test_ascend_deepseek_mtp.py @@ -21,7 +21,6 @@ TEST_MODEL_MATRIX = { class TestAscendDeepSeekMTP(CustomTestCase): - @classmethod def setUpClass(cls): cls.models = TEST_MODEL_MATRIX.keys() diff --git a/test/manual/attention/test_trtllm_mla_backend.py b/test/manual/attention/test_trtllm_mla_backend.py index 344e5bcce..193ba8919 100755 --- a/test/manual/attention/test_trtllm_mla_backend.py +++ b/test/manual/attention/test_trtllm_mla_backend.py @@ -274,12 +274,12 @@ def compare_outputs(trtllm_out, reference_out, tolerance=1e-2): """Compare outputs with detailed analysis.""" # Basic checks - assert ( - trtllm_out.shape == reference_out.shape - ), f"Shape mismatch: {trtllm_out.shape} vs {reference_out.shape}" - assert ( - trtllm_out.dtype == reference_out.dtype - ), f"Dtype mismatch: {trtllm_out.dtype} vs {reference_out.dtype}" + assert trtllm_out.shape == reference_out.shape, ( + f"Shape mismatch: {trtllm_out.shape} vs {reference_out.shape}" + ) + assert trtllm_out.dtype == reference_out.dtype, ( + f"Dtype mismatch: {trtllm_out.dtype} vs {reference_out.dtype}" + ) # Check for NaN/Inf assert not torch.isnan(trtllm_out).any(), "TRTLLM output contains NaN" @@ -310,7 +310,7 @@ def compare_outputs(trtllm_out, reference_out, tolerance=1e-2): trt_val = trtllm_out[idx_tuple].item() ref_val = reference_out[idx_tuple].item() print( - f" [{idx_tuple}]: TRTLLM={trt_val:.6f}, Reference={ref_val:.6f}, diff={abs(trt_val-ref_val):.6f}" + f" [{idx_tuple}]: TRTLLM={trt_val:.6f}, Reference={ref_val:.6f}, diff={abs(trt_val - ref_val):.6f}" ) return all_close diff --git a/test/manual/beam_search/test_beam_search_perf_sweep.py b/test/manual/beam_search/test_beam_search_perf_sweep.py index 0e2deb579..dd1c237ce 100644 --- a/test/manual/beam_search/test_beam_search_perf_sweep.py +++ b/test/manual/beam_search/test_beam_search_perf_sweep.py @@ -72,9 +72,9 @@ class _BeamSweepBase(CustomTestCase): cls.prompts = [r.prompt for r in rows if r.prompt_len < MAX_PROMPT_LEN][ :NUM_PROMPTS ] - assert ( - len(cls.prompts) == NUM_PROMPTS - ), f"only {len(cls.prompts)} short prompts sampled" + assert len(cls.prompts) == NUM_PROMPTS, ( + f"only {len(cls.prompts)} short prompts sampled" + ) @classmethod def tearDownClass(cls): diff --git a/test/manual/chunked_prefill/test_scripted_abort.py b/test/manual/chunked_prefill/test_scripted_abort.py index 2214133c8..93169c1ba 100644 --- a/test/manual/chunked_prefill/test_scripted_abort.py +++ b/test/manual/chunked_prefill/test_scripted_abort.py @@ -53,15 +53,15 @@ class TestAbortBasic(ScriptedTestCase): "finished", "unknown", ), f"after abort r should be finished/unknown, got {r.status}" - assert ( - r.kv_pages == 0 - ), f"abort must release KV; r.kv_pages={r.kv_pages} after abort" - assert ( - r.req is None or r.req.kv.req_pool_idx is None - ), f"abort must release row; r.req={r.req} after abort" - assert ( - r.lock_refs == 0 - ), f"abort must release lock_refs; r.lock_refs={r.lock_refs}" + assert r.kv_pages == 0, ( + f"abort must release KV; r.kv_pages={r.kv_pages} after abort" + ) + assert r.req is None or r.req.kv.req_pool_idx is None, ( + f"abort must release row; r.req={r.req} after abort" + ) + assert r.lock_refs == 0, ( + f"abort must release lock_refs; r.lock_refs={r.lock_refs}" + ) def test_abort_at_chunk_0(self): self.server.execute_script(self._script_abort_at_chunk_0) @@ -436,9 +436,9 @@ class TestAbortBasic(ScriptedTestCase): assert not r.is_chunking, "aborted gap req must stay out of chunking" if r.req is not None: - assert ( - r.req.inflight_middle_chunks == 0 - ), f"inflight_middle_chunks not cleared; got {r.req.inflight_middle_chunks}" + assert r.req.inflight_middle_chunks == 0, ( + f"inflight_middle_chunks not cleared; got {r.req.inflight_middle_chunks}" + ) def test_abort_when_chunked_only_then_idle(self): self.server.execute_script(self._script_abort_when_chunked_only_then_idle) @@ -509,10 +509,10 @@ class TestAbortBasic(ScriptedTestCase): yield from _drain_until_released(t, r1) assert r1.kv_pages == 0, ( - f"force_retract + abort same yield must release KV; got " f"{r1.kv_pages}" + f"force_retract + abort same yield must release KV; got {r1.kv_pages}" ) assert r1.req is None or r1.req.kv.req_pool_idx is None, ( - f"force_retract + abort same yield must release row; got " f"{r1.req}" + f"force_retract + abort same yield must release row; got {r1.req}" ) assert r1.lock_refs == 0, ( f"force_retract + abort same yield must release lock_refs; " diff --git a/test/manual/chunked_prefill/test_scripted_chunk_size.py b/test/manual/chunked_prefill/test_scripted_chunk_size.py index 975461f67..d44145912 100644 --- a/test/manual/chunked_prefill/test_scripted_chunk_size.py +++ b/test/manual/chunked_prefill/test_scripted_chunk_size.py @@ -73,9 +73,9 @@ class TestChunkSizeDefault(ScriptedTestCase): r = t.start_req(prompt_len=1, max_new_tokens=2) yield from run_until_finished(r) assert r.finished - assert ( - r.chunks_done == 0 - ), f"single-token prompt should not chunk, got chunks_done={r.chunks_done}" + assert r.chunks_done == 0, ( + f"single-token prompt should not chunk, got chunks_done={r.chunks_done}" + ) def test_chunk_size_256_prompt_100x(self): self.server.execute_script(self._script_chunk_size_256_prompt_100x) diff --git a/test/manual/chunked_prefill/test_scripted_http_smoke.py b/test/manual/chunked_prefill/test_scripted_http_smoke.py index 74943a699..da52a7362 100644 --- a/test/manual/chunked_prefill/test_scripted_http_smoke.py +++ b/test/manual/chunked_prefill/test_scripted_http_smoke.py @@ -26,9 +26,9 @@ class TestScriptedHttpSmoke(ScriptedTestCase): break yield assert r.finished - assert ( - saw_chunking - ), "expected the req to hold the chunked_req slot at least once" + assert saw_chunking, ( + "expected the req to hold the chunked_req slot at least once" + ) def test_two_reqs_finish(self): self.server.execute_script(self._script_two_reqs_finish) diff --git a/test/manual/chunked_prefill/test_scripted_hybrid_swa.py b/test/manual/chunked_prefill/test_scripted_hybrid_swa.py index b4af4cb90..1dd80b8a9 100644 --- a/test/manual/chunked_prefill/test_scripted_hybrid_swa.py +++ b/test/manual/chunked_prefill/test_scripted_hybrid_swa.py @@ -123,9 +123,9 @@ class TestSWAHalfWindowChunk(ScriptedTestCase): r = t.start_req(prompt_len=2 * _SWA_WINDOW, max_new_tokens=4) yield from run_until_finished(r, max_steps=800) assert r.finished - assert ( - r.chunks_done >= 4 - ), f"expected >=4 chunks for 2*window / (window/2), got {r.chunks_done}" + assert r.chunks_done >= 4, ( + f"expected >=4 chunks for 2*window / (window/2), got {r.chunks_done}" + ) assert len(r.req.output_ids) == 4 @@ -171,9 +171,9 @@ class TestSWARadix(ScriptedTestCase): ) yield from run_until_finished(r2, max_steps=800) assert r2.finished - assert ( - r2.req.cached_tokens > 0 - ), f"r2 must hit the radix prefix, got cached_tokens={r2.req.cached_tokens}" + assert r2.req.cached_tokens > 0, ( + f"r2 must hit the radix prefix, got cached_tokens={r2.req.cached_tokens}" + ) if __name__ == "__main__": diff --git a/test/manual/chunked_prefill/test_scripted_invariants.py b/test/manual/chunked_prefill/test_scripted_invariants.py index 451efd9eb..68df1fa51 100644 --- a/test/manual/chunked_prefill/test_scripted_invariants.py +++ b/test/manual/chunked_prefill/test_scripted_invariants.py @@ -38,9 +38,9 @@ class TestInvariantsBasic(ScriptedTestCase): for _ in range(DEFAULT_MAX_STEPS): if r.is_chunking: observed_chunking = True - assert ( - r.kv_pages > 0 - ), f"kv_pages must be > 0 while is_chunking; got {r.kv_pages}" + assert r.kv_pages > 0, ( + f"kv_pages must be > 0 while is_chunking; got {r.kv_pages}" + ) if r.finished: break yield @@ -63,9 +63,9 @@ class TestInvariantsBasic(ScriptedTestCase): + comp.get("decode", []) + comp.get("chunked", []) ) - assert ( - r.rid in all_rids - ), f"running but not in batch_composition: {comp}" + assert r.rid in all_rids, ( + f"running but not in batch_composition: {comp}" + ) if r.finished: return yield @@ -157,9 +157,9 @@ class TestInvariantsBasic(ScriptedTestCase): t.flush_cache() yield final = t.engine_stats() - assert ( - final["kv_pool_free"] >= baseline["kv_pool_free"] - ), f"KV leak: {baseline['kv_pool_free']} -> {final['kv_pool_free']}" + assert final["kv_pool_free"] >= baseline["kv_pool_free"], ( + f"KV leak: {baseline['kv_pool_free']} -> {final['kv_pool_free']}" + ) assert final["req_pool_free"] >= baseline["req_pool_free"] def test_long_lived_engine_reps_chunked(self): @@ -212,9 +212,9 @@ class TestInvariantsBasic(ScriptedTestCase): t.flush_cache() yield final_kv = t.engine_stats()["kv_pool_free"] - assert ( - final_kv >= baseline_kv - ), f"KV leak after sustained chunked load: {baseline_kv} -> {final_kv}" + assert final_kv >= baseline_kv, ( + f"KV leak after sustained chunked load: {baseline_kv} -> {final_kv}" + ) def test_round_robin_short_and_chunked(self): self.server.execute_script(self._script_round_robin_short_and_chunked) @@ -356,9 +356,9 @@ class TestInvariantsBasic(ScriptedTestCase): ) yield from run_until_finished(r) assert r.finished - assert ( - r.chunks_done >= 2 - ), f"VERY_LONG_PROMPT_LEN should chunk; got chunks_done={r.chunks_done}" + assert r.chunks_done >= 2, ( + f"VERY_LONG_PROMPT_LEN should chunk; got chunks_done={r.chunks_done}" + ) assert len(r.req.output_ids) == n, ( f"ignore_eos=True + max_new_tokens={n} must produce exactly " f"{n} output tokens; got len(output_tokens)={len(r.req.output_ids)}" diff --git a/test/manual/chunked_prefill/test_scripted_kv_pressure.py b/test/manual/chunked_prefill/test_scripted_kv_pressure.py index 81be5f5f8..419370851 100644 --- a/test/manual/chunked_prefill/test_scripted_kv_pressure.py +++ b/test/manual/chunked_prefill/test_scripted_kv_pressure.py @@ -62,9 +62,9 @@ class TestKVPressureBasic(ScriptedTestCase): f"long req must really chunk under pinned cache; got chunks_done=" f"{r_long.chunks_done}" ) - assert ( - r_long.lock_refs == 0 - ), f"req {r_long.rid} leaked {r_long.lock_refs} lock_refs after finish" + assert r_long.lock_refs == 0, ( + f"req {r_long.rid} leaked {r_long.lock_refs} lock_refs after finish" + ) t._release_exhausted_pools() final_lock_refs = t.get_all_node_lock_refs() diff --git a/test/manual/chunked_prefill/test_scripted_lifecycle.py b/test/manual/chunked_prefill/test_scripted_lifecycle.py index 8cc136bca..f11c7a7c1 100644 --- a/test/manual/chunked_prefill/test_scripted_lifecycle.py +++ b/test/manual/chunked_prefill/test_scripted_lifecycle.py @@ -171,9 +171,9 @@ class TestLifecycleBasic(ScriptedTestCase): assert "running" in seen, f"never observed running status; seen={seen}" assert seen[-1] == "finished", f"final status must be finished; seen={seen}" finished_idx = seen.index("finished") - assert all( - s in ("finished",) for s in seen[finished_idx:] - ), f"status regressed after finish; seen={seen}" + assert all(s in ("finished",) for s in seen[finished_idx:]), ( + f"status regressed after finish; seen={seen}" + ) def test_long_prompt_only_one_decode(self): self.server.execute_script(self._script_long_prompt_only_one_decode) @@ -279,8 +279,7 @@ class TestLifecycleBasic(ScriptedTestCase): assert r1.finished and r2.finished assert r2.chunks_done == 0 assert r2.req.cached_tokens > 0, ( - f"r2 must hit r1's radix prefix; got cached_tokens=" - f"{r2.req.cached_tokens}" + f"r2 must hit r1's radix prefix; got cached_tokens={r2.req.cached_tokens}" ) assert len(r2.req.output_ids) == 2 @@ -408,9 +407,9 @@ class TestLifecycleBasic(ScriptedTestCase): t.flush_cache() yield final = t.engine_stats()["kv_pool_free"] - assert ( - final >= baseline - 1 - ), f"KV pool drift: baseline={baseline}, final={final}" + assert final >= baseline - 1, ( + f"KV pool drift: baseline={baseline}, final={final}" + ) def test_abort_all_during_chunked(self): self.server.execute_script(self._script_abort_all_during_chunked) diff --git a/test/manual/chunked_prefill/test_scripted_max_new_tokens.py b/test/manual/chunked_prefill/test_scripted_max_new_tokens.py index f1da3114f..44579e15d 100644 --- a/test/manual/chunked_prefill/test_scripted_max_new_tokens.py +++ b/test/manual/chunked_prefill/test_scripted_max_new_tokens.py @@ -143,14 +143,13 @@ class TestMaxNewTokensFirstDecodeAdjacent(ScriptedTestCase): decode_records = _decode_records(batch_log, r.rid) assert len(decode_records) == max_new_tokens, ( - f"expected {max_new_tokens} decode forwards, got " f"{len(decode_records)}" + f"expected {max_new_tokens} decode forwards, got {len(decode_records)}" ) rid_modes = [rec.mode for rec in rid_records] first_decode_pos = rid_modes.index("decode") assert first_decode_pos >= 1, ( - f"first decode must be preceded by an extend chunk; rid_modes=" - f"{rid_modes}" + f"first decode must be preceded by an extend chunk; rid_modes={rid_modes}" ) assert rid_modes[first_decode_pos - 1] == "extend", ( f"record immediately before the first decode (in this rid's " diff --git a/test/manual/chunked_prefill/test_scripted_multi_req.py b/test/manual/chunked_prefill/test_scripted_multi_req.py index d9ba9f207..7a6a5432e 100644 --- a/test/manual/chunked_prefill/test_scripted_multi_req.py +++ b/test/manual/chunked_prefill/test_scripted_multi_req.py @@ -20,9 +20,9 @@ def _drain_flush_then_assert_no_kv_leak(t: ScriptedContext, baseline: dict): t.flush_cache() yield final = t.engine_stats() - assert ( - final["kv_pool_free"] >= baseline["kv_pool_free"] - ), f"KV leak: {baseline['kv_pool_free']} -> {final['kv_pool_free']}" + assert final["kv_pool_free"] >= baseline["kv_pool_free"], ( + f"KV leak: {baseline['kv_pool_free']} -> {final['kv_pool_free']}" + ) class TestMultiReqBasic(ScriptedTestCase): @@ -39,9 +39,9 @@ class TestMultiReqBasic(ScriptedTestCase): yield assert r1.is_chunking, "r1 should still be chunking" - assert ( - not r2.is_chunking - ), "r2 must wait for r1's chunk loop to clear before chunking" + assert not r2.is_chunking, ( + "r2 must wait for r1's chunk loop to clear before chunking" + ) yield from run_until_all_finished([r1, r2]) assert r1.finished and r2.finished @@ -144,9 +144,9 @@ class TestMultiReqBasic(ScriptedTestCase): r2 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN + 8, max_new_tokens=2) yield from run_until_finished(r2) assert r1.finished and r2.finished - assert ( - r2.chunks_done < r1.chunks_done - ), "r2 reuses r1's cached prefix, so it should chunk fewer times" + assert r2.chunks_done < r1.chunks_done, ( + "r2 reuses r1's cached prefix, so it should chunk fewer times" + ) def test_trickle_per_yield_50(self): self.server.execute_script(self._script_trickle_per_yield_50) diff --git a/test/manual/chunked_prefill/test_scripted_pp.py b/test/manual/chunked_prefill/test_scripted_pp.py index 21fbafcac..55777be31 100644 --- a/test/manual/chunked_prefill/test_scripted_pp.py +++ b/test/manual/chunked_prefill/test_scripted_pp.py @@ -175,9 +175,9 @@ class TestPPPdmux(ScriptedTestCase): def _script_pp_split_prefill_chunked_no_merge_assert(t: ScriptedContext): r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) yield from run_until_finished(r, max_steps=800) - assert ( - r.finished - ), "engine died before req finished — merge_batch assert may have tripped" + assert r.finished, ( + "engine died before req finished — merge_batch assert may have tripped" + ) assert r.chunks_done >= 2, ( f"pdmux + chunked path must produce >=2 chunks to exercise " f"split_prefill_batch filter; got chunks_done={r.chunks_done}" diff --git a/test/manual/chunked_prefill/test_scripted_priority.py b/test/manual/chunked_prefill/test_scripted_priority.py index 090f4aec4..be8d087ef 100644 --- a/test/manual/chunked_prefill/test_scripted_priority.py +++ b/test/manual/chunked_prefill/test_scripted_priority.py @@ -30,9 +30,9 @@ class TestPriorityBasic(ScriptedTestCase): t.pause_generation(mode="retract") yield - assert ( - r.status == "waiting" - ), f"force-retracted chunked req must be back in waiting; got {r.status}" + assert r.status == "waiting", ( + f"force-retracted chunked req must be back in waiting; got {r.status}" + ) assert r.kv_pages == 0, f"retract must release KV; got {r.kv_pages}" t.continue_generation() diff --git a/test/manual/chunked_prefill/test_scripted_radix.py b/test/manual/chunked_prefill/test_scripted_radix.py index c868dccf5..280e651a2 100644 --- a/test/manual/chunked_prefill/test_scripted_radix.py +++ b/test/manual/chunked_prefill/test_scripted_radix.py @@ -76,9 +76,9 @@ class TestRadixBasic(ScriptedTestCase): f"after eviction r2 must re-chunk from scratch; " f"chunks_done={r2.chunks_done} cached_tokens={r2.req.cached_tokens}" ) - assert ( - r2.req.cached_tokens == 0 - ), f"eviction must clear r1's prefix; cached_tokens={r2.req.cached_tokens}" + assert r2.req.cached_tokens == 0, ( + f"eviction must clear r1's prefix; cached_tokens={r2.req.cached_tokens}" + ) assert r2.kv_pages == 0 assert r2.lock_refs == 0 @@ -219,9 +219,9 @@ class TestRadixBasic(ScriptedTestCase): r = t.start_req(prompt_len=prompt_len, max_new_tokens=2) yield from run_until_finished(r, max_steps=400) assert r.finished - assert ( - r.chunks_done == 0 - ), f"full prefix hit must skip chunked path; got chunks_done={r.chunks_done}" + assert r.chunks_done == 0, ( + f"full prefix hit must skip chunked path; got chunks_done={r.chunks_done}" + ) def test_radix_evict_race_concurrent_chunked_admit(self): self.server.execute_script( @@ -319,9 +319,9 @@ class TestRadixNoTailChunked(ScriptedTestCase): "test must observe r as the in-flight chunked_req at least once; the " "no-tail else branch was never exercised" ) - assert ( - r.kv_pages == 0 - ), f"finished chunked req must release KV; got {r.kv_pages}" + assert r.kv_pages == 0, ( + f"finished chunked req must release KV; got {r.kv_pages}" + ) class TestRadixHitCountInvariant(ScriptedTestCase): diff --git a/test/manual/chunked_prefill/test_scripted_regression.py b/test/manual/chunked_prefill/test_scripted_regression.py index 29f0c1731..4a91e0a01 100644 --- a/test/manual/chunked_prefill/test_scripted_regression.py +++ b/test/manual/chunked_prefill/test_scripted_regression.py @@ -100,9 +100,9 @@ class TestRegressionBasic(ScriptedTestCase): f"observed max={observed_max} (pre-fix bug would bump to 2 " f"at the last-chunk admit boundary)" ) - assert ( - cleared_inflight - ), "inflight_middle_chunks should be 0 once the chunk loop clears" + assert cleared_inflight, ( + "inflight_middle_chunks should be 0 once the chunk loop clears" + ) yield from run_until_finished(r) assert r.finished @@ -247,15 +247,15 @@ class TestRegressionBasic(ScriptedTestCase): t.abort(r) yield from _drain_until_released(t, r) - assert ( - r.req.kv.req_pool_idx is None - ), f"96d4749094: abort must release row; got row_idx={r.req.kv.req_pool_idx!r}" - assert ( - r.kv_pages == 0 - ), f"96d4749094: abort must release KV; got kv_pages={r.kv_pages}" - assert ( - r.lock_refs == 0 - ), f"96d4749094: abort must release lock_ref; got lock_refs={r.lock_refs}" + assert r.req.kv.req_pool_idx is None, ( + f"96d4749094: abort must release row; got row_idx={r.req.kv.req_pool_idx!r}" + ) + assert r.kv_pages == 0, ( + f"96d4749094: abort must release KV; got kv_pages={r.kv_pages}" + ) + assert r.lock_refs == 0, ( + f"96d4749094: abort must release lock_ref; got lock_refs={r.lock_refs}" + ) assert not r.is_chunking assert r.req.inflight_middle_chunks == 0 assert sum(t.get_all_node_lock_refs().values()) == baseline_refs diff --git a/test/manual/chunked_prefill/test_scripted_sampling.py b/test/manual/chunked_prefill/test_scripted_sampling.py index 11ccee95e..46942b368 100644 --- a/test/manual/chunked_prefill/test_scripted_sampling.py +++ b/test/manual/chunked_prefill/test_scripted_sampling.py @@ -34,8 +34,7 @@ class TestSamplingBasic(ScriptedTestCase): if r.rid in rec.rids and rec.mode == "decode" ] assert len(decode_records) == 0, ( - f"max_new_tokens=0 must run zero decode forwards; got " - f"{len(decode_records)}" + f"max_new_tokens=0 must run zero decode forwards; got {len(decode_records)}" ) def test_max_new_tokens_one_long_chunked(self): @@ -100,8 +99,7 @@ class TestSamplingBasic(ScriptedTestCase): assert r.chunks_done >= 2 assert len(r.req.output_ids) == 16 assert isinstance(r.req.finished_reason, FINISH_LENGTH), ( - f"ignore_eos=True must finish via length cap; got " - f"{r.req.finished_reason!r}" + f"ignore_eos=True must finish via length cap; got {r.req.finished_reason!r}" ) def test_return_logprob_top_logprobs_chunked(self): @@ -176,9 +174,9 @@ class TestSamplingBasic(ScriptedTestCase): ) yield from run_until_finished(r) assert r.finished - assert ( - r.chunks_done >= 2 - ), f"prompt should span multiple chunks, got chunks_done={r.chunks_done}" + assert r.chunks_done >= 2, ( + f"prompt should span multiple chunks, got chunks_done={r.chunks_done}" + ) assert r.req.logprob is not None input_lp = r.req.logprob.input_token_logprobs_val assert len(input_lp) == prompt_len, ( @@ -201,9 +199,9 @@ class TestSamplingBasic(ScriptedTestCase): ) yield from run_until_finished(r) assert r.finished - assert ( - r.chunks_done >= 3 - ), f"prompt should span 3+ chunks, got chunks_done={r.chunks_done}" + assert r.chunks_done >= 3, ( + f"prompt should span 3+ chunks, got chunks_done={r.chunks_done}" + ) assert r.req.logprob is not None input_lp = r.req.logprob.input_token_logprobs_val assert len(input_lp) == prompt_len - start_len, ( @@ -246,9 +244,9 @@ class TestSamplingBasic(ScriptedTestCase): ) yield from run_until_finished(r_eos, max_steps=2000) assert r_eos.finished - assert ( - r_eos.chunks_done >= 2 - ), f"scenario 1 should chunk; got chunks_done={r_eos.chunks_done}" + assert r_eos.chunks_done >= 2, ( + f"scenario 1 should chunk; got chunks_done={r_eos.chunks_done}" + ) assert isinstance(r_eos.req.finished_reason, FINISH_MATCHED_TOKEN), ( f"a stop token the model deterministically produces under greedy must " f"finish via the matched-token path; got {r_eos.req.finished_reason!r}" @@ -261,9 +259,9 @@ class TestSamplingBasic(ScriptedTestCase): ) yield from run_until_finished(r_length) assert r_length.finished - assert ( - r_length.chunks_done >= 2 - ), f"scenario 2 should chunk; got chunks_done={r_length.chunks_done}" + assert r_length.chunks_done >= 2, ( + f"scenario 2 should chunk; got chunks_done={r_length.chunks_done}" + ) assert isinstance(r_length.req.finished_reason, FINISH_LENGTH), ( f"ignore_eos=True + max_new_tokens=4 chunked must finish via " f"length cap; got {r_length.req.finished_reason!r}" diff --git a/test/manual/chunked_prefill/test_scripted_special_case.py b/test/manual/chunked_prefill/test_scripted_special_case.py index 499e9d251..29db97567 100644 --- a/test/manual/chunked_prefill/test_scripted_special_case.py +++ b/test/manual/chunked_prefill/test_scripted_special_case.py @@ -45,9 +45,9 @@ class TestSpecialCaseBasic(ScriptedTestCase): for _ in range(DEFAULT_MAX_STEPS): if r.is_chunking: saw_chunking = True - assert ( - not t.is_idle - ), "scheduler must not idle while chunked_req is in flight" + assert not t.is_idle, ( + "scheduler must not idle while chunked_req is in flight" + ) if r.finished: break yield @@ -102,9 +102,9 @@ class TestSpecialCaseBasic(ScriptedTestCase): break yield - assert ( - t.scheduler.chunked_req is None - ), f"abort must clear the chunked slot; got {t.scheduler.chunked_req!r}" + assert t.scheduler.chunked_req is None, ( + f"abort must clear the chunked slot; got {t.scheduler.chunked_req!r}" + ) assert r.kv_pages == 0 assert r.lock_refs == 0 @@ -193,19 +193,19 @@ class TestSpecialCaseBasic(ScriptedTestCase): if r1.is_chunking: saw_r1_chunking = True comp = t.batch_composition() - assert r1.rid in comp.get( - "chunked", [] - ), f"mid-chunk r1 must occupy the chunked role; got {comp!r}" - assert r1.rid not in comp.get( - "running", [] - ), f"chunked r1 must be excluded from the running role; got {comp!r}" + assert r1.rid in comp.get("chunked", []), ( + f"mid-chunk r1 must occupy the chunked role; got {comp!r}" + ) + assert r1.rid not in comp.get("running", []), ( + f"chunked r1 must be excluded from the running role; got {comp!r}" + ) if r1.finished and r2.finished: break yield assert r1.finished and r2.finished - assert ( - saw_r1_chunking - ), "r1 must have chunked at some point to exercise the exclude branch" + assert saw_r1_chunking, ( + "r1 must have chunked at some point to exercise the exclude branch" + ) @unittest.skip( "pdmux split_prefill_batch requires the pdmux topology — " @@ -257,9 +257,9 @@ class TestSpecialCaseBasic(ScriptedTestCase): t.pause_generation(mode="retract") yield - assert ( - t.scheduler.chunked_req is None - ), f"pause(retract) must clear chunked_req; got {t.scheduler.chunked_req!r}" + assert t.scheduler.chunked_req is None, ( + f"pause(retract) must clear chunked_req; got {t.scheduler.chunked_req!r}" + ) assert not r.finished, "retract must re-queue r, not finish or abort it" assert r.status == "waiting", ( f"retracted chunked req must return to the waiting queue; " @@ -268,9 +268,9 @@ class TestSpecialCaseBasic(ScriptedTestCase): t.continue_generation() yield from run_until_finished(r) - assert ( - r.finished - ), "continue_generation must drive the re-queued req to completion" + assert r.finished, ( + "continue_generation must drive the re-queued req to completion" + ) def test_retract_during_gap_inflight_middle_chunks_positive(self): self.server.execute_script( @@ -304,9 +304,9 @@ class TestSpecialCaseBasic(ScriptedTestCase): t.continue_generation() yield from run_until_finished(r, max_steps=2000) - assert ( - r.finished - ), "continue_generation must drive the re-queued req to completion" + assert r.finished, ( + "continue_generation must drive the re-queued req to completion" + ) assert r.kv_pages == 0 assert len(r.req.output_ids) == 2 @@ -344,9 +344,9 @@ class TestSpecialCaseBasic(ScriptedTestCase): break yield assert r.finished - assert ( - saw_chunking - ), "test must observe the dual-queue chunked state at least once" + assert saw_chunking, ( + "test must observe the dual-queue chunked state at least once" + ) assert saw_dedup, ( "test must observe the chunked req with a committed prefix so the " "dedup subtraction is actually exercised" @@ -491,9 +491,9 @@ class TestSpecialCaseBasic(ScriptedTestCase): break yield assert r.finished - assert ( - saw_mid_chunk - ), "test must observe the fill_ids reset boundary at least once" + assert saw_mid_chunk, ( + "test must observe the fill_ids reset boundary at least once" + ) assert r.finished def test_chunked_req_slot_cleared_when_chunk_completes(self): @@ -516,9 +516,9 @@ class TestSpecialCaseBasic(ScriptedTestCase): yield assert r.finished assert saw_chunking, "req should have occupied the chunked_req slot mid-chunk" - assert ( - s.chunked_req is None - ), f"chunked_req slot must clear after last chunk; got {s.chunked_req!r}" + assert s.chunked_req is None, ( + f"chunked_req slot must clear after last chunk; got {s.chunked_req!r}" + ) def test_second_chunked_admit_blocked_when_chunked_req_set(self): self.server.execute_script( @@ -571,9 +571,9 @@ class TestSpecialCaseBasic(ScriptedTestCase): progressed = False for _ in range(DEFAULT_MAX_STEPS): if r.is_chunking: - assert ( - not t.is_idle - ), "scheduler must not go idle while a chunked req is in flight" + assert not t.is_idle, ( + "scheduler must not go idle while a chunked req is in flight" + ) cur_chunks_done = r.chunks_done if cur_chunks_done > prev_chunks_done: progressed = True @@ -681,9 +681,9 @@ class TestSpecialCaseMixedChunk(ScriptedTestCase): prompt_token=310, ) yield from run_until(r, lambda h: h.is_chunking) - assert ( - t.last_batch_forward_mode != "MIXED" - ), f"return_logprob must disable mixed-chunk path; got {t.last_batch_forward_mode!r}" + assert t.last_batch_forward_mode != "MIXED", ( + f"return_logprob must disable mixed-chunk path; got {t.last_batch_forward_mode!r}" + ) yield from run_until_finished(r) def test_mixed_chunk_with_running_batch(self): @@ -700,9 +700,9 @@ class TestSpecialCaseMixedChunk(ScriptedTestCase): yield yield from run_until(r_chunk, lambda h: h.is_chunking) - assert ( - t.last_batch_forward_mode == "MIXED" - ), f"chunked admission with running batch must enter MIXED; got {t.last_batch_forward_mode!r}" + assert t.last_batch_forward_mode == "MIXED", ( + f"chunked admission with running batch must enter MIXED; got {t.last_batch_forward_mode!r}" + ) for _ in range(DEFAULT_MAX_STEPS * 2): if r_chunk.finished and r_dec.finished: break @@ -743,9 +743,9 @@ class TestSpecialCaseNoChunking(ScriptedTestCase): def _script_chunk_size_negative_disables_chunking(t: ScriptedContext): r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) for _ in range(DEFAULT_MAX_STEPS): - assert ( - not r.is_chunking - ), "chunked_prefill_size=-1 should disable chunked path" + assert not r.is_chunking, ( + "chunked_prefill_size=-1 should disable chunked path" + ) if r.finished: return yield @@ -825,9 +825,9 @@ class TestSpecialCaseHiCache(ScriptedTestCase): yield assert r.finished assert saw_chunking, "test must observe r mid-chunk at least once" - assert ( - first_chunk_snap is not None - ), "test must snapshot cached_tokens at the first chunk boundary" + assert first_chunk_snap is not None, ( + "test must snapshot cached_tokens at the first chunk boundary" + ) def test_hicache_cached_tokens_set_once_invariant(self): self.server.execute_script( @@ -862,9 +862,9 @@ class TestSpecialCaseHiCache(ScriptedTestCase): break yield assert r.finished - assert ( - saw_chunking - ), "test must observe the req mid-chunk (chunks_done >= 1) at least once" + assert saw_chunking, ( + "test must observe the req mid-chunk (chunks_done >= 1) at least once" + ) assert snap is not None, "test must snapshot the cached_tokens_* breakdown" @@ -1011,24 +1011,23 @@ class TestSpecialCaseRetractMerge(ScriptedTestCase): t.pause_generation(mode="retract") yield - assert ( - s.last_batch is None - ), "retract must clear last_batch after merging the extend chunk batch" + assert s.last_batch is None, ( + "retract must clear last_batch after merging the extend chunk batch" + ) assert len(s.running_batch.reqs) == 0, ( "the merged extend chunk batch must be retracted out of running_batch, " f"not stranded; got {len(s.running_batch.reqs)} reqs" ) - assert ( - r.status == "waiting" - ), f"retracted chunked req must return to the waiting queue; got {r.status!r}" + assert r.status == "waiting", ( + f"retracted chunked req must return to the waiting queue; got {r.status!r}" + ) assert r.kv_pages == 0 t.continue_generation() yield from run_until_finished(r) assert r.finished assert len(r.req.output_ids) == 2, ( - f"resumed req must emit exactly max_new_tokens; got " - f"{len(r.req.output_ids)}" + f"resumed req must emit exactly max_new_tokens; got {len(r.req.output_ids)}" ) @@ -1183,9 +1182,9 @@ class TestSpecialCaseRetractedStain(ScriptedTestCase): assert r.finished req = r.req - assert ( - req.retracted_stain is True - ), "retract must set retracted_stain so the cached-token recount is suppressed" + assert req.retracted_stain is True, ( + "retract must set retracted_stain so the cached-token recount is suppressed" + ) assert req.cached_tokens == cached_before, ( f"retracted_stain must suppress re-adding pre_len-already_computed on " f"resume; cached_tokens grew from {cached_before} to {req.cached_tokens}" @@ -1237,17 +1236,17 @@ class TestSpecialCaseMiddleChunkNoToken(ScriptedTestCase): f"middle chunk must not append an output token; got " f"output_ids len {len(r.req.output_ids)}" ) - assert ( - r.status != "finished" - ), "middle chunk must not finish the req (skip_stream_req)" + assert r.status != "finished", ( + "middle chunk must not finish the req (skip_stream_req)" + ) if r.finished: break yield assert r.finished assert saw_middle_chunk, "test must observe r mid-chunk at least once" - assert ( - len(r.req.output_ids) >= 1 - ), "output tokens must appear only after the chunked prefill completes" + assert len(r.req.output_ids) >= 1, ( + "output tokens must appear only after the chunked prefill completes" + ) if __name__ == "__main__": diff --git a/test/manual/dsv4/test_fused_compress_attn_hip.py b/test/manual/dsv4/test_fused_compress_attn_hip.py index 0e6da091d..053321b98 100644 --- a/test/manual/dsv4/test_fused_compress_attn_hip.py +++ b/test/manual/dsv4/test_fused_compress_attn_hip.py @@ -264,7 +264,6 @@ def _ref_compress( class TestFusedCompressAttn(unittest.TestCase): - def _run_test( self, ratio: int, @@ -397,7 +396,6 @@ class TestFusedCompressAttn(unittest.TestCase): class TestStateOrdering(unittest.TestCase): - def test_write_then_compress(self): """Verify write-first, compress-second matches reference.""" device = torch.device("cuda") diff --git a/test/manual/ep/test_deepep_internode.py b/test/manual/ep/test_deepep_internode.py index 1e4239606..273f1442c 100644 --- a/test/manual/ep/test_deepep_internode.py +++ b/test/manual/ep/test_deepep_internode.py @@ -149,7 +149,7 @@ def test_main( for with_topk in (False, True): if local_rank == 0: print( - f'[testing] Running with {"FP8" if isinstance(current_x, tuple) else "BF16"}, {"with" if with_topk else "without"} top-k (async={async_mode}, previous={previous_mode}) ...', + f"[testing] Running with {'FP8' if isinstance(current_x, tuple) else 'BF16'}, {'with' if with_topk else 'without'} top-k (async={async_mode}, previous={previous_mode}) ...", flush=True, end="", ) @@ -192,9 +192,9 @@ def test_main( # Checks recv_gbl_rank_prefix_sum = handle[-4] - assert gbl_num_tokens_per_rank[rank].item() == recv_x.size( - 0 - ), f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}" + assert gbl_num_tokens_per_rank[rank].item() == recv_x.size(0), ( + f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}" + ) assert ( gbl_num_tokens_per_expert.view(num_ranks, -1)[rank].tolist() == recv_num_tokens_per_expert_list @@ -316,10 +316,13 @@ def test_main( tune_args = {"x": current_x, "handle": handle, "config": config} t = bench(lambda: buffer.dispatch(**tune_args))[0] if t < best_time: - best_time, best_results = t, ( - num_sms, - nvl_chunk_size, - rdma_chunk_size, + best_time, best_results = ( + t, + ( + num_sms, + nvl_chunk_size, + rdma_chunk_size, + ), ) if local_rank == 0: print( @@ -328,7 +331,7 @@ def test_main( ) if local_rank == 0: print( - f'[tuning] Best dispatch ({"FP8" if isinstance(current_x, tuple) else "BF16"}): SMs {best_results[0]}, NVL chunk {best_results[1]}, RDMA chunk {best_results[2]}: {rdma_send_bytes / 1e9 / best_time:.2f} GB/s (RDMA), {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)', + f"[tuning] Best dispatch ({'FP8' if isinstance(current_x, tuple) else 'BF16'}): SMs {best_results[0]}, NVL chunk {best_results[1]}, RDMA chunk {best_results[2]}: {rdma_send_bytes / 1e9 / best_time:.2f} GB/s (RDMA), {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)", flush=True, ) print("", flush=True) @@ -385,10 +388,13 @@ def test_main( flush=True, ) if t < best_time: - best_time, best_results = t, ( - num_sms, - nvl_chunk_size, - rdma_chunk_size, + best_time, best_results = ( + t, + ( + num_sms, + nvl_chunk_size, + rdma_chunk_size, + ), ) if local_rank == 0: diff --git a/test/manual/ep/test_deepep_intranode.py b/test/manual/ep/test_deepep_intranode.py index ef365b010..9e117fb14 100644 --- a/test/manual/ep/test_deepep_intranode.py +++ b/test/manual/ep/test_deepep_intranode.py @@ -120,7 +120,7 @@ def test_main( for with_topk in (False, True): if local_rank == 0: print( - f'[testing] Running with {"FP8" if isinstance(current_x, tuple) else "BF16"}, {"with" if with_topk else "without"} top-k (async={async_mode}, previous={previous_mode}) ...', + f"[testing] Running with {'FP8' if isinstance(current_x, tuple) else 'BF16'}, {'with' if with_topk else 'without'} top-k (async={async_mode}, previous={previous_mode}) ...", flush=True, end="", ) @@ -162,9 +162,9 @@ def test_main( # Checks rank_prefix_matrix = handle[0] - assert gbl_num_tokens_per_rank[rank].item() == recv_x.size( - 0 - ), f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}" + assert gbl_num_tokens_per_rank[rank].item() == recv_x.size(0), ( + f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}" + ) assert ( gbl_num_tokens_per_expert.view(num_ranks, -1)[rank].tolist() == recv_num_tokens_per_expert_list @@ -280,7 +280,7 @@ def test_main( ) if local_rank == 0: print( - f'[tuning] Best dispatch ({"FP8" if isinstance(current_x, tuple) else "BF16"}): SMs {best_results[0]}, NVL chunk {best_results[1]}, {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)', + f"[tuning] Best dispatch ({'FP8' if isinstance(current_x, tuple) else 'BF16'}): SMs {best_results[0]}, NVL chunk {best_results[1]}, {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)", flush=True, ) print("", flush=True) diff --git a/test/manual/ep/test_deepep_low_latency.py b/test/manual/ep/test_deepep_low_latency.py index 1b3f34da9..47b3e55ef 100644 --- a/test/manual/ep/test_deepep_low_latency.py +++ b/test/manual/ep/test_deepep_low_latency.py @@ -36,9 +36,9 @@ def test_main( # NOTES: the integers greater than 256 exceeds the BF16 precision limit rank_offset = 128 - assert ( - num_ranks - rank_offset < 257 - ), "Too many ranks (exceeding test precision limit)" + assert num_ranks - rank_offset < 257, ( + "Too many ranks (exceeding test precision limit)" + ) x = torch.ones((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") * ( rank - rank_offset @@ -55,9 +55,9 @@ def test_main( # Randomly mask some positions for i in range(10): - topk_idx[random.randint(0, num_tokens - 1), random.randint(0, num_topk - 1)] = ( - -1 - ) + topk_idx[ + random.randint(0, num_tokens - 1), random.randint(0, num_topk - 1) + ] = -1 # Check dispatch correctness do_check = True @@ -114,9 +114,9 @@ def test_main( assert ( num_valid_tokens == (recv_layout_range & int_mask).sum().item() ), f"{num_valid_tokens} != {recv_layout_range & int_mask}.sum().item()" - assert ( - num_valid_tokens == (all_topk_idx == expert_id).sum().item() - ), f"{num_valid_tokens} != {(all_topk_idx == expert_id).sum().item()}" + assert num_valid_tokens == (all_topk_idx == expert_id).sum().item(), ( + f"{num_valid_tokens} != {(all_topk_idx == expert_id).sum().item()}" + ) # Check received data recv_x = recv_x[:num_valid_tokens] @@ -127,9 +127,10 @@ def test_main( recv_x[:, -128:] - recv_src_info.view(-1, 1) % num_tokens ).sum().item() == 0 for j in range(num_ranks): - begin_idx, count = (recv_layout_range[j] >> 32).item(), ( - recv_layout_range[j] & int_mask - ).item() + begin_idx, count = ( + (recv_layout_range[j] >> 32).item(), + (recv_layout_range[j] & int_mask).item(), + ) assert (recv_x_amin == j - rank_offset).sum().item() == ( all_topk_idx[j] == expert_id ).sum().item() @@ -145,9 +146,9 @@ def test_main( # Check combine correctness for zero_copy in (False, True): if zero_copy: - buffer.get_next_low_latency_combine_buffer(handle)[ - :, :, : - ] = simulated_gemm_x + buffer.get_next_low_latency_combine_buffer(handle)[:, :, :] = ( + simulated_gemm_x + ) out = torch.empty( (num_tokens, hidden), dtype=torch.bfloat16, device="cuda" ) @@ -203,9 +204,9 @@ def test_main( ) large_gemm_with_hook(hook) if return_recv_hook else None if zero_copy: - buffer.get_next_low_latency_combine_buffer(handle)[ - :, :, : - ] = simulated_gemm_x + buffer.get_next_low_latency_combine_buffer(handle)[:, :, :] = ( + simulated_gemm_x + ) combined_x, event, hook = buffer.low_latency_combine( simulated_gemm_x, topk_idx, diff --git a/test/manual/ep/test_elastic_scale.py b/test/manual/ep/test_elastic_scale.py index 2bbed7ceb..cc5041287 100644 --- a/test/manual/ep/test_elastic_scale.py +++ b/test/manual/ep/test_elastic_scale.py @@ -238,8 +238,7 @@ class _ElasticScaleUpEndToEndBase(CustomTestCase): join_end = rank_offset + join_tp if join_end > len(visible_devices): raise RuntimeError( - f"Scale-up requires {join_end} visible GPUs, got " - f"{len(visible_devices)}" + f"Scale-up requires {join_end} visible GPUs, got {len(visible_devices)}" ) env["CUDA_VISIBLE_DEVICES"] = ",".join(visible_devices[rank_offset:join_end]) base_joining_log = os.environ.get( diff --git a/test/manual/ep/test_flashinfer_dispatcher.py b/test/manual/ep/test_flashinfer_dispatcher.py index 8bb540f66..cadf605cc 100644 --- a/test/manual/ep/test_flashinfer_dispatcher.py +++ b/test/manual/ep/test_flashinfer_dispatcher.py @@ -16,7 +16,6 @@ from sglang.test.test_utils import CustomTestCase class TestFlashinferDispatcher(CustomTestCase): - @classmethod def setUpClass(cls): server_args = ServerArgs(model_path="dummy") @@ -125,8 +124,7 @@ class TestFlashinferDispatcher(CustomTestCase): self.assertTrue( torch.all( received_hidden_states[ - expected_source_rank - * num_tokens : (expected_source_rank + 1) + expected_source_rank * num_tokens : (expected_source_rank + 1) * num_tokens ] == 100.0 + expected_source_rank @@ -229,8 +227,7 @@ class TestFlashinferDispatcher(CustomTestCase): self.assertTrue( torch.all( received_hidden_states[ - expected_source_rank - * num_tokens : (expected_source_rank + 1) + expected_source_rank * num_tokens : (expected_source_rank + 1) * num_tokens ] == 100.0 + expected_source_rank diff --git a/test/manual/hicache/test_disaggregation_hicache.py b/test/manual/hicache/test_disaggregation_hicache.py index fe340e6f4..d1caa2b64 100644 --- a/test/manual/hicache/test_disaggregation_hicache.py +++ b/test/manual/hicache/test_disaggregation_hicache.py @@ -242,7 +242,7 @@ class TestDisaggregationDecodeWithHiCache(DisaggregationHiCacheBase): self.assertGreater( cached_tokens, previous_cached_tokens, - f"Turn {turn} should have more cached tokens than turn {turn-1}", + f"Turn {turn} should have more cached tokens than turn {turn - 1}", ) # Update context and cached tokens for next iteration diff --git a/test/manual/kernels/test_hisparse_prefetch.py b/test/manual/kernels/test_hisparse_prefetch.py index 12aad76aa..b2bee6d9d 100644 --- a/test/manual/kernels/test_hisparse_prefetch.py +++ b/test/manual/kernels/test_hisparse_prefetch.py @@ -380,14 +380,14 @@ def test_plan_then_io_cuda_graph_replay() -> None: graph.replay() torch.cuda.synchronize() # Anchor slot table matches the synchronous layer-0 result. - assert torch.equal( - out.cpu(), ref_slots[s][0].cpu() - ), f"slots differ at step {s}" + assert torch.equal(out.cpu(), ref_slots[s][0].cpu()), ( + f"slots differ at step {s}" + ) # Every layer's device buffer stays bit-identical to synchronous swap-in. for layer in range(_PIO_LAYERS): - assert torch.equal( - buf[layer].cpu(), ref_snap[s][layer].cpu() - ), f"buffer differs at step {s}, layer {layer}" + assert torch.equal(buf[layer].cpu(), ref_snap[s][layer].cpu()), ( + f"buffer differs at step {s}, layer {layer}" + ) if __name__ == "__main__": diff --git a/test/manual/lang_frontend/test_choices.py b/test/manual/lang_frontend/test_choices.py index 89e7ca1c7..1467360b6 100644 --- a/test/manual/lang_frontend/test_choices.py +++ b/test/manual/lang_frontend/test_choices.py @@ -53,7 +53,6 @@ MOCK_CHOICES_INPUT_DATA = { class TestChoices(CustomTestCase): - def test_token_length_normalized(self): """Confirm 'antidisestablishmentarianism' is selected due to high confidences for its later tokens resulting in highest token length normalized prompt logprob.""" diff --git a/test/manual/lang_frontend/test_jump_forward.py b/test/manual/lang_frontend/test_jump_forward.py index 315a50b5b..7c9ea8915 100644 --- a/test/manual/lang_frontend/test_jump_forward.py +++ b/test/manual/lang_frontend/test_jump_forward.py @@ -47,6 +47,7 @@ json_jump_forward = ( + r"""\}\n""" ) + # fmt: off @sgl.function def json_gen(s): diff --git a/test/manual/layers/attention/dsa/test_get_k_scale_triton_kernel.py b/test/manual/layers/attention/dsa/test_get_k_scale_triton_kernel.py index 2f86dff3d..78f35b085 100644 --- a/test/manual/layers/attention/dsa/test_get_k_scale_triton_kernel.py +++ b/test/manual/layers/attention/dsa/test_get_k_scale_triton_kernel.py @@ -175,7 +175,7 @@ def get_k_and_s_triton(): end_time = time.perf_counter() print( - f"_get_k_and_s_triton_kernel triton kernel infer time is {((end_time-start_time)*1000):.4f} ms\n" + f"_get_k_and_s_triton_kernel triton kernel infer time is {((end_time - start_time) * 1000):.4f} ms\n" ) diff --git a/test/manual/lora/test_lora_backend.py b/test/manual/lora/test_lora_backend.py index 5a78bd46f..ae3f7f18d 100644 --- a/test/manual/lora/test_lora_backend.py +++ b/test/manual/lora/test_lora_backend.py @@ -30,7 +30,6 @@ from sglang.test.test_utils import CustomTestCase, is_in_ci class TestLoRABackend(CustomTestCase): - def _run_backend_on_model_cases(self, model_cases: List[LoRAModelCase]): for model_case in model_cases: # If skip_long_prompt is True, filter out prompts longer than 1000 characters diff --git a/test/manual/lora/test_lora_cuda_graph.py b/test/manual/lora/test_lora_cuda_graph.py index c110d4144..3fbe3425c 100644 --- a/test/manual/lora/test_lora_cuda_graph.py +++ b/test/manual/lora/test_lora_cuda_graph.py @@ -44,7 +44,6 @@ TEST_CUDA_GRAPH_PADDING_PROMPTS = [ class TestLoRACudaGraph(CustomTestCase): - def _run_without_cuda_graph_on_model_cases(self, model_cases: List[LoRAModelCase]): # Since we have already enabled CUDA graph by default in other lora tests, # we only need to run lora tests without CUDA graph here. diff --git a/test/manual/lora/test_torch_backend.py b/test/manual/lora/test_torch_backend.py index f0c46dd3e..b88245a69 100644 --- a/test/manual/lora/test_torch_backend.py +++ b/test/manual/lora/test_torch_backend.py @@ -9,7 +9,6 @@ from sglang.test.test_utils import CustomTestCase class TestTorchNativeLoRABackend(CustomTestCase): - device = "cpu" # set duplicate weights to test merging during prepare_lora_batch diff --git a/test/manual/minimax_m3/test_npu_memory_pool.py b/test/manual/minimax_m3/test_npu_memory_pool.py index 509d5d379..b80e22038 100644 --- a/test/manual/minimax_m3/test_npu_memory_pool.py +++ b/test/manual/minimax_m3/test_npu_memory_pool.py @@ -74,7 +74,9 @@ class _FakeMHATokenToKVPool(_FakeKVCache): self.v_head_dim = ( swa_v_head_dim if swa_v_head_dim is not None - else v_head_dim if v_head_dim is not None else head_dim + else v_head_dim + if v_head_dim is not None + else head_dim ) self._create_buffers() diff --git a/test/manual/models/test_clip_models.py b/test/manual/models/test_clip_models.py index e537a7466..ad579f0f3 100644 --- a/test/manual/models/test_clip_models.py +++ b/test/manual/models/test_clip_models.py @@ -29,7 +29,6 @@ TORCH_DTYPES = [torch.float16] class TestClipModels(unittest.TestCase): - @classmethod def setUpClass(cls): mp.set_start_method("spawn", force=True) @@ -61,12 +60,12 @@ class TestClipModels(unittest.TestCase): ) print("text similarity diff", abs(text_similarity - 1)) print("image similarity diff", abs(image_similarity - 1)) - assert torch.all( - abs(text_similarity - 1) < prefill_tolerance - ), "embeddings are not all close" - assert torch.all( - abs(image_similarity - 1) < prefill_tolerance - ), "embeddings are not all close" + assert torch.all(abs(text_similarity - 1) < prefill_tolerance), ( + "embeddings are not all close" + ) + assert torch.all(abs(image_similarity - 1) < prefill_tolerance), ( + "embeddings are not all close" + ) def test_accuracy(self): for model, prefill_tolerance in MODELS: diff --git a/test/manual/models/test_gme_qwen_models.py b/test/manual/models/test_gme_qwen_models.py index 265bf8a39..5177eb15f 100644 --- a/test/manual/models/test_gme_qwen_models.py +++ b/test/manual/models/test_gme_qwen_models.py @@ -64,16 +64,16 @@ class TestQmeQwenModels(CustomTestCase): hf_text_embeddings.embed_logits[0], srt_text_embeddings.embed_logits[0] ) print("texts similarity diff", abs(similarity - 1)) - assert torch.all( - abs(similarity - 1) < prefill_tolerance - ), "embeddings are not all close" + assert torch.all(abs(similarity - 1) < prefill_tolerance), ( + "embeddings are not all close" + ) similarity = get_similarities( hf_image_embeddings.embed_logits[0], srt_image_embeddings.embed_logits[0] ) print("images similarity diff", abs(similarity - 1)) - assert torch.all( - abs(similarity - 1) < prefill_tolerance - ), "embeddings are not all close" + assert torch.all(abs(similarity - 1) < prefill_tolerance), ( + "embeddings are not all close" + ) def test_accuracy(self): for model, prefill_tolerance in MODELS: diff --git a/test/manual/models/test_kimi_k2_models.py b/test/manual/models/test_kimi_k2_models.py index 6e83ef50c..812c1589e 100644 --- a/test/manual/models/test_kimi_k2_models.py +++ b/test/manual/models/test_kimi_k2_models.py @@ -61,7 +61,7 @@ class TestKimiK2Thinking(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k (Kimi-K2-Thinking)\n" f'{metrics["score"]=:.3f}\n' + f'### test_gsm8k (Kimi-K2-Thinking)\n{metrics["score"]=:.3f}\n' ) self.assertGreater(metrics["score"], 0.95) diff --git a/test/manual/models/test_mistral_large3_basic.py b/test/manual/models/test_mistral_large3_basic.py index 2eac4f79b..7004b5a2d 100644 --- a/test/manual/models/test_mistral_large3_basic.py +++ b/test/manual/models/test_mistral_large3_basic.py @@ -67,7 +67,7 @@ class TestMistralLarge3Basic(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k (mistral-large-3)\n" f'{metrics["score"]=:.3f}\n' + f'### test_gsm8k (mistral-large-3)\n{metrics["score"]=:.3f}\n' ) self.assertGreater(metrics["score"], 0.90) @@ -79,7 +79,7 @@ class TestMistralLarge3Basic(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_bs_1_speed (mistral-large-3)\n" f"{speed=:.2f} token/s\n" + f"### test_bs_1_speed (mistral-large-3)\n{speed=:.2f} token/s\n" ) self.assertGreater(speed, 50) diff --git a/test/manual/models/test_qwen3_asr.py b/test/manual/models/test_qwen3_asr.py index c2e7b4278..114158ade 100644 --- a/test/manual/models/test_qwen3_asr.py +++ b/test/manual/models/test_qwen3_asr.py @@ -156,9 +156,9 @@ async def _stream_websocket_async( async with websockets.connect(websocket_url) as websocket: created = json.loads(await websocket.recv()) - assert ( - created.get("type") == "session.created" - ), f"expected session.created, got {created!r}" + assert created.get("type") == "session.created", ( + f"expected session.created, got {created!r}" + ) session_id = created["session"]["id"] transcription_cfg = {"model": "qwen3-asr"} @@ -200,9 +200,9 @@ async def _stream_websocket_async( if t == "conversation.item.input_audio_transcription.delta": deltas.append(resp["delta"]) elif t == "conversation.item.input_audio_transcription.completed": - assert ( - "usage" in resp - ), f"transcription.completed missing required usage field: {resp!r}" + assert "usage" in resp, ( + f"transcription.completed missing required usage field: {resp!r}" + ) assert resp["usage"].get("type") == "duration", resp["usage"] completed_msg.update(resp) return @@ -317,7 +317,7 @@ class TestQwen3ASRTranscription(CustomTestCase): self.assertEqual( results[0], results[i], - f"Request {i+1} differs from first request", + f"Request {i + 1} differs from first request", ) print(f"[Consistency] All 3 requests match: {results[0][:80]}...") diff --git a/test/manual/nightly/test_vlms_piecewise_cuda_graph.py b/test/manual/nightly/test_vlms_piecewise_cuda_graph.py index 8bacf3691..ac9b22f4a 100644 --- a/test/manual/nightly/test_vlms_piecewise_cuda_graph.py +++ b/test/manual/nightly/test_vlms_piecewise_cuda_graph.py @@ -66,7 +66,7 @@ class TestVLMPiecewiseCudaGraph(CustomTestCase): os.makedirs(output_path, exist_ok=True) # -------- compose --model_args -------- - model_args = f'model_version="{model_version}",' f"tp={tp}" + model_args = f'model_version="{model_version}",tp={tp}' # -------- build command list -------- cmd = [ diff --git a/test/manual/nightly/test_vlms_vit_cuda_graph.py b/test/manual/nightly/test_vlms_vit_cuda_graph.py index 4b2bc74d7..19472a145 100644 --- a/test/manual/nightly/test_vlms_vit_cuda_graph.py +++ b/test/manual/nightly/test_vlms_vit_cuda_graph.py @@ -69,7 +69,7 @@ class TestVLMViTCudaGraph(CustomTestCase): os.makedirs(output_path, exist_ok=True) # -------- compose --model_args -------- - model_args = f'model_version="{model_version}",' f"tp={tp}" + model_args = f'model_version="{model_version}",tp={tp}' # -------- build command list -------- cmd = [ diff --git a/test/manual/nightly/test_vlms_vit_flashinfer_cudnn.py b/test/manual/nightly/test_vlms_vit_flashinfer_cudnn.py index 1706bbda2..89d03795d 100644 --- a/test/manual/nightly/test_vlms_vit_flashinfer_cudnn.py +++ b/test/manual/nightly/test_vlms_vit_flashinfer_cudnn.py @@ -66,7 +66,7 @@ class TestVLMViTFlashinferCudnn(CustomTestCase): os.makedirs(output_path, exist_ok=True) # -------- compose --model_args -------- - model_args = f'model_version="{model_version}",' f"tp={tp}" + model_args = f'model_version="{model_version}",tp={tp}' # -------- build command list -------- cmd = [ diff --git a/test/manual/openai_server/features/test_cache_report.py b/test/manual/openai_server/features/test_cache_report.py index ec45eb119..2ea8c4530 100644 --- a/test/manual/openai_server/features/test_cache_report.py +++ b/test/manual/openai_server/features/test_cache_report.py @@ -250,12 +250,12 @@ class TestCacheReport(CustomTestCase): ) # Verify cache hit for same salt - assert ( - cached_tokens_1_second > cached_tokens_1_first - ), "Should have cache hit with same cache_salt" - assert ( - cached_tokens_1_second == prompt_tokens_1 - 1 - ), "Should cache all prompt tokens except the last one" + assert cached_tokens_1_second > cached_tokens_1_first, ( + "Should have cache hit with same cache_salt" + ) + assert cached_tokens_1_second == prompt_tokens_1 - 1, ( + "Should cache all prompt tokens except the last one" + ) # Third request with different cache_salt "salt2" - should not get cache hit response3 = self.client.chat.completions.create( @@ -269,9 +269,9 @@ class TestCacheReport(CustomTestCase): print(f"First request with salt2 - cached_tokens: {cached_tokens_2_first}") # Verify no cache hit for different salt (should be similar to first request with salt1) - assert ( - cached_tokens_2_first <= cached_tokens_1_first + self.min_cached - ), "Different cache_salt should not share cache" + assert cached_tokens_2_first <= cached_tokens_1_first + self.min_cached, ( + "Different cache_salt should not share cache" + ) # Fourth request with same cache_salt "salt2" - should now get cache hit response4 = self.client.chat.completions.create( @@ -285,9 +285,9 @@ class TestCacheReport(CustomTestCase): print(f"Second request with salt2 - cached_tokens: {cached_tokens_2_second}") # Verify cache hit for salt2 - assert ( - cached_tokens_2_second > cached_tokens_2_first - ), "Should have cache hit with same cache_salt for salt2" + assert cached_tokens_2_second > cached_tokens_2_first, ( + "Should have cache hit with same cache_salt for salt2" + ) if __name__ == "__main__": diff --git a/test/manual/openai_server/function_call/test_tool_choice.py b/test/manual/openai_server/function_call/test_tool_choice.py index 6e5eae797..28178d4da 100644 --- a/test/manual/openai_server/function_call/test_tool_choice.py +++ b/test/manual/openai_server/function_call/test_tool_choice.py @@ -23,7 +23,6 @@ from sglang.test.test_utils import ( class TestToolChoiceLlama32(CustomTestCase): - @classmethod def setUpClass(cls): # Mark flaky tests for this model @@ -379,15 +378,15 @@ class TestToolChoiceLlama32(CustomTestCase): # Update function name if present (first chunk) if tool_call_delta.function and tool_call_delta.function.name: - tool_calls_by_index[tool_index]["function"][ - "name" - ] = tool_call_delta.function.name + tool_calls_by_index[tool_index]["function"]["name"] = ( + tool_call_delta.function.name + ) # Accumulate arguments (all chunks) if tool_call_delta.function and tool_call_delta.function.arguments: - tool_calls_by_index[tool_index]["function"][ - "arguments" - ] += tool_call_delta.function.arguments + tool_calls_by_index[tool_index]["function"]["arguments"] += ( + tool_call_delta.function.arguments + ) self.assertGreater(len(tool_calls_by_index), 0) diff --git a/test/manual/perf/test_bench_one_batch_1gpu.py b/test/manual/perf/test_bench_one_batch_1gpu.py index 6d8ece897..10f51330d 100644 --- a/test/manual/perf/test_bench_one_batch_1gpu.py +++ b/test/manual/perf/test_bench_one_batch_1gpu.py @@ -17,7 +17,6 @@ from sglang.test.test_utils import ( class TestBenchOneBatch1GPU(CustomTestCase): - def test_bs1_small(self): _, output_throughput, _ = run_bench_one_batch( DEFAULT_SMALL_MODEL_NAME_FOR_TEST, ["--cuda-graph-max-bs-decode", "2"] diff --git a/test/manual/prefill_only/test_cross_encoder_models.py b/test/manual/prefill_only/test_cross_encoder_models.py index 0c7a867ce..85f3b3787 100644 --- a/test/manual/prefill_only/test_cross_encoder_models.py +++ b/test/manual/prefill_only/test_cross_encoder_models.py @@ -20,7 +20,6 @@ TORCH_DTYPES = [torch.float32] class TestCrossEncoderModels(CustomTestCase): - @classmethod def setUpClass(cls): mp.set_start_method("spawn", force=True) @@ -55,9 +54,9 @@ class TestCrossEncoderModels(CustomTestCase): for i in range(len(srt_scores)): score_difference = abs(hf_scores[i] - srt_scores[i]) - assert ( - score_difference < score_tolerance - ), "cross encoder scores are not all close" + assert score_difference < score_tolerance, ( + "cross encoder scores are not all close" + ) def preprocess_prompts(self, prompt): processed_prompts = [] diff --git a/test/manual/prefill_only/test_encoder_embedding_models.py b/test/manual/prefill_only/test_encoder_embedding_models.py index cbc8178e6..143c99ec9 100644 --- a/test/manual/prefill_only/test_encoder_embedding_models.py +++ b/test/manual/prefill_only/test_encoder_embedding_models.py @@ -37,7 +37,6 @@ sgl_to_st_ratio = [] class TestEncoderEmbeddingModels(CustomTestCase): - @classmethod def setUpClass(cls): mp.set_start_method("spawn", force=True) @@ -115,9 +114,9 @@ class TestEncoderEmbeddingModels(CustomTestCase): # print("similarity diff", abs(similarity - 1)) if len(truncated_prompts[i]) <= 1000: - assert torch.all( - abs(similarity - 1) < prefill_tolerance - ), "embeddings are not all close" + assert torch.all(abs(similarity - 1) < prefill_tolerance), ( + "embeddings are not all close" + ) def test_prefill_logits(self): models_to_test = MODELS diff --git a/test/manual/quant/test_block_fp8.py b/test/manual/quant/test_block_fp8.py index b45eeaa6e..8356a3a93 100644 --- a/test/manual/quant/test_block_fp8.py +++ b/test/manual/quant/test_block_fp8.py @@ -46,9 +46,9 @@ def native_per_token_group_quant_fp8( quantized tensor along with the scaling factor used for quantization. Note that only `torch.float8_e4m3fn` is supported for now. """ - assert ( - x.shape[-1] % group_size == 0 - ), "the last dimension of `x` cannot be divisible by `group_size`" + assert x.shape[-1] % group_size == 0, ( + "the last dimension of `x` cannot be divisible by `group_size`" + ) assert x.is_contiguous(), "`x` is not contiguous" finfo = torch.finfo(dtype) @@ -343,7 +343,6 @@ def native_w8a8_block_fp8_matmul(A, B, As, Bs, block_size, output_dtype=torch.fl class TestW8A8BlockFP8Matmul(CustomTestCase): - if not _is_cuda: OUT_DTYPES = [torch.float32, torch.half, torch.bfloat16] M = [1, 7, 83, 512, 2048] diff --git a/test/manual/quant/test_block_fp8_deep_gemm_blackwell.py b/test/manual/quant/test_block_fp8_deep_gemm_blackwell.py index 1d138ee41..41aba3e70 100644 --- a/test/manual/quant/test_block_fp8_deep_gemm_blackwell.py +++ b/test/manual/quant/test_block_fp8_deep_gemm_blackwell.py @@ -171,7 +171,6 @@ def block_quant_dequant( class TestDeepGemmBlackwell(CustomTestCase): - if not _is_cuda: OUT_DTYPES = [torch.float32, torch.half, torch.bfloat16] M = [1, 7, 83, 512, 2048] diff --git a/test/manual/quant/test_deepseek_v32_fp4_4gpu.py b/test/manual/quant/test_deepseek_v32_fp4_4gpu.py index 2752a0b8c..df7e8e967 100644 --- a/test/manual/quant/test_deepseek_v32_fp4_4gpu.py +++ b/test/manual/quant/test_deepseek_v32_fp4_4gpu.py @@ -67,7 +67,7 @@ class TestDeepseekV32FP4DP(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k (deepseek-v3-fp4)\n" f'{metrics["score"]=:.3f}\n' + f'### test_gsm8k (deepseek-v3-fp4)\n{metrics["score"]=:.3f}\n' ) self.assertGreater(metrics["score"], 0.93) @@ -135,7 +135,7 @@ class TestDeepseekV32FP4TP(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k (deepseek-v3-fp4)\n" f'{metrics["score"]=:.3f}\n' + f'### test_gsm8k (deepseek-v3-fp4)\n{metrics["score"]=:.3f}\n' ) self.assertGreater(metrics["score"], 0.93) diff --git a/test/manual/quant/test_deepseek_v3_fp4_4gpu_trtllm.py b/test/manual/quant/test_deepseek_v3_fp4_4gpu_trtllm.py index dc7fb9c61..742882d4e 100644 --- a/test/manual/quant/test_deepseek_v3_fp4_4gpu_trtllm.py +++ b/test/manual/quant/test_deepseek_v3_fp4_4gpu_trtllm.py @@ -70,7 +70,7 @@ class TestDeepseekV3FP4(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k (deepseek-v3-fp4)\n" f'{metrics["score"]=:.3f}\n' + f'### test_gsm8k (deepseek-v3-fp4)\n{metrics["score"]=:.3f}\n' ) self.assertGreater(metrics["score"], 0.93) @@ -83,7 +83,7 @@ class TestDeepseekV3FP4(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_bs_1_speed (deepseek-v3-fp4)\n" f"{speed=:.2f} token/s\n" + f"### test_bs_1_speed (deepseek-v3-fp4)\n{speed=:.2f} token/s\n" ) self.assertGreater(speed, 120) diff --git a/test/manual/quant/test_eval_fp8_accuracy.py b/test/manual/quant/test_eval_fp8_accuracy.py index 329e2dad8..d6b8b9eea 100644 --- a/test/manual/quant/test_eval_fp8_accuracy.py +++ b/test/manual/quant/test_eval_fp8_accuracy.py @@ -46,7 +46,6 @@ class TestEvalFP8Accuracy(CustomTestCase): class TestEvalFP8DynamicQuantAccuracy(CustomTestCase): - def _run_test(self, model, other_args, expected_score): base_url = DEFAULT_URL_FOR_TEST other_args = other_args or [] diff --git a/test/manual/scheduler/test_no_chunked_prefill.py b/test/manual/scheduler/test_no_chunked_prefill.py index 59869ff59..c6b457552 100644 --- a/test/manual/scheduler/test_no_chunked_prefill.py +++ b/test/manual/scheduler/test_no_chunked_prefill.py @@ -9,7 +9,6 @@ from sglang.test.test_utils import ( class TestNoChunkedPrefill(CustomTestCase): - def test_no_chunked_prefill(self): run_mmlu_test( disable_radix_cache=False, enable_mixed_chunk=False, chunked_prefill_size=-1 diff --git a/test/manual/spec/eagle/test_eagle3_basic.py b/test/manual/spec/eagle/test_eagle3_basic.py index dc2b31f90..b0c1956d2 100644 --- a/test/manual/spec/eagle/test_eagle3_basic.py +++ b/test/manual/spec/eagle/test_eagle3_basic.py @@ -59,5 +59,4 @@ class TestEagle3Basic(EagleServerBase): if __name__ == "__main__": - unittest.main() diff --git a/test/manual/test_create_custom_4d_mask.py b/test/manual/test_create_custom_4d_mask.py index 8627ed863..6c79b83b6 100644 --- a/test/manual/test_create_custom_4d_mask.py +++ b/test/manual/test_create_custom_4d_mask.py @@ -386,8 +386,8 @@ class TestPerformance(unittest.TestCase): speedup = t_ref / t_new dev_tag = "CUDA" if "cuda" in str(self.device) else "CPU" print( - f" [{dev_tag}] B={batch_size:3d} S={seq_len:5d} img%={int(image_fraction*100):3d}%" - f" ref={t_ref*1e3:.2f}ms new={t_new*1e3:.2f}ms speedup={speedup:.2f}x" + f" [{dev_tag}] B={batch_size:3d} S={seq_len:5d} img%={int(image_fraction * 100):3d}%" + f" ref={t_ref * 1e3:.2f}ms new={t_new * 1e3:.2f}ms speedup={speedup:.2f}x" ) self.assertGreaterEqual( speedup, diff --git a/test/manual/test_fim_completion.py b/test/manual/test_fim_completion.py index 6efdfe776..4db9562e6 100644 --- a/test/manual/test_fim_completion.py +++ b/test/manual/test_fim_completion.py @@ -57,9 +57,9 @@ class TestFimCompletion(CustomTestCase): assert response.id assert response.created assert response.object == "text_completion" - assert ( - response.usage.prompt_tokens == num_prompt_tokens - ), f"{response.usage.prompt_tokens} vs {num_prompt_tokens}" + assert response.usage.prompt_tokens == num_prompt_tokens, ( + f"{response.usage.prompt_tokens} vs {num_prompt_tokens}" + ) assert response.usage.completion_tokens > 0 assert response.usage.total_tokens > 0 diff --git a/test/manual/test_get_weights_by_name.py b/test/manual/test_get_weights_by_name.py index fa97c7df8..23528ca27 100644 --- a/test/manual/test_get_weights_by_name.py +++ b/test/manual/test_get_weights_by_name.py @@ -30,7 +30,6 @@ def _process_return(ret): class TestGetWeightsByName(CustomTestCase): - def init_hf_model(self, model_name, tie_word_embeddings): self.hf_model = AutoModelForCausalLM.from_pretrained( model_name, torch_dtype="bfloat16", tie_word_embeddings=tie_word_embeddings diff --git a/test/manual/test_kda_target_verify.py b/test/manual/test_kda_target_verify.py index a25ad5452..03389e42f 100644 --- a/test/manual/test_kda_target_verify.py +++ b/test/manual/test_kda_target_verify.py @@ -104,15 +104,15 @@ def test_kda_target_verify_equivalence(): state_diff = (cached_state - decode_state).abs().max().item() status = "OK" if state_diff < 1e-5 else "FAIL" print(f" step={step} req={req_idx}: diff={state_diff:.6e} [{status}]") - assert ( - state_diff < 1e-5 - ), f"Intermediate state mismatch at step={step}, req={req_idx}: {state_diff}" + assert state_diff < 1e-5, ( + f"Intermediate state mismatch at step={step}, req={req_idx}: {state_diff}" + ) ssm_unchanged_diff = (ssm_states_verify - ssm_states_base).abs().max().item() print(f"SSM state in-place change (should be 0): {ssm_unchanged_diff:.6e}") - assert ( - ssm_unchanged_diff == 0.0 - ), f"target_verify modified ssm_states in-place! diff: {ssm_unchanged_diff}" + assert ssm_unchanged_diff == 0.0, ( + f"target_verify modified ssm_states in-place! diff: {ssm_unchanged_diff}" + ) print("\nPASSED: KDA target_verify matches sequential decode!") diff --git a/test/manual/test_logprobs.py b/test/manual/test_logprobs.py index 9256d1b4b..af77a181a 100644 --- a/test/manual/test_logprobs.py +++ b/test/manual/test_logprobs.py @@ -218,7 +218,6 @@ def generate_baseline( class TestLogprobsDense(unittest.TestCase): - @classmethod def setUpClass(cls): """Set up the test class - initialize the engine once for all tests.""" diff --git a/test/manual/test_modelopt.py b/test/manual/test_modelopt.py index ef6a959ec..69a90e385 100644 --- a/test/manual/test_modelopt.py +++ b/test/manual/test_modelopt.py @@ -16,7 +16,6 @@ from sglang.test.test_utils import ( class TestEvalFP8ModelOptQuantAccuracy(CustomTestCase): - def _run_test(self, model, other_args, expected_score): base_url = DEFAULT_URL_FOR_TEST other_args = other_args or [] diff --git a/test/manual/test_models_from_modelscope.py b/test/manual/test_models_from_modelscope.py index dacaae30a..077d48c11 100644 --- a/test/manual/test_models_from_modelscope.py +++ b/test/manual/test_models_from_modelscope.py @@ -9,7 +9,6 @@ from sglang.test.test_utils import CustomTestCase class TestDownloadFromModelScope(CustomTestCase): - @classmethod def setUpClass(cls): cls.model = "iic/nlp_lstmcrf_word-segmentation_chinese-news" diff --git a/test/manual/test_ray_engine.py b/test/manual/test_ray_engine.py index 0eeb996ba..e53f343bc 100644 --- a/test/manual/test_ray_engine.py +++ b/test/manual/test_ray_engine.py @@ -155,7 +155,6 @@ def _cleanup(actor, pg): @unittest.skipUnless(_has_ray, "ray is not installed") @unittest.skipUnless(_NUM_GPUS >= 1, "requires at least 1 GPU") class TestRayEngineOfflineTP1(unittest.TestCase): - @classmethod def setUpClass(cls): if not ray.is_initialized(): @@ -196,7 +195,6 @@ class TestRayEngineOfflineTP1(unittest.TestCase): @unittest.skipUnless(_has_ray, "ray is not installed") @unittest.skipUnless(_NUM_GPUS >= 2, "requires at least 2 GPUs") class TestRayEngineOfflineTP2(unittest.TestCase): - @classmethod def setUpClass(cls): if not ray.is_initialized(): @@ -231,7 +229,6 @@ class TestRayEngineOfflineTP2(unittest.TestCase): @unittest.skipUnless(_has_ray, "ray is not installed") @unittest.skipUnless(_NUM_GPUS >= 2, "requires at least 2 GPUs") class TestRayEngineOfflinePP2(unittest.TestCase): - @classmethod def setUpClass(cls): if not ray.is_initialized(): @@ -347,7 +344,6 @@ class TestRayEngineOfflineDPAttention(unittest.TestCase): @unittest.skipUnless(_has_ray, "ray is not installed") @unittest.skipUnless(_NUM_GPUS >= 1, "requires at least 1 GPU") class TestRayEngineErrors(unittest.TestCase): - @classmethod def setUpClass(cls): if not ray.is_initialized(): diff --git a/test/manual/test_sagemaker_server.py b/test/manual/test_sagemaker_server.py index 7f54b0919..2be86f75b 100644 --- a/test/manual/test_sagemaker_server.py +++ b/test/manual/test_sagemaker_server.py @@ -69,9 +69,9 @@ class TestSageMakerServer(CustomTestCase): ret_num_top_logprobs = len( response["choices"][0]["logprobs"]["content"][0]["top_logprobs"] ) - assert ( - ret_num_top_logprobs == logprobs - ), f"{ret_num_top_logprobs} vs {logprobs}" + assert ret_num_top_logprobs == logprobs, ( + f"{ret_num_top_logprobs} vs {logprobs}" + ) assert len(response["choices"]) == parallel_sample_num assert response["choices"][0]["message"]["role"] == "assistant" @@ -155,18 +155,18 @@ class TestSageMakerServer(CustomTestCase): .get("content")[0] .get("top_logprobs") ) - assert ( - ret_num_top_logprobs == logprobs - ), f"{ret_num_top_logprobs} vs {logprobs}" + assert ret_num_top_logprobs == logprobs, ( + f"{ret_num_top_logprobs} vs {logprobs}" + ) assert isinstance(data["content"], str) assert line["id"] assert line["created"] for index in [i for i in range(parallel_sample_num)]: - assert not is_firsts.get( - index, True - ), f"index {index} is not found in the response" + assert not is_firsts.get(index, True), ( + f"index {index} is not found in the response" + ) def test_chat_completion(self): for logprobs in [None, 5]: diff --git a/test/manual/test_schedule_policy.py b/test/manual/test_schedule_policy.py index c04481a46..7258e31f4 100644 --- a/test/manual/test_schedule_policy.py +++ b/test/manual/test_schedule_policy.py @@ -25,7 +25,6 @@ def _make_req(rid, origin_input_text, origin_input_ids, sampling_params=None, ** class TestSchedulePolicy(CustomTestCase): - def setUp(self): self.tree_cache = RadixCache.create_simulated() diff --git a/test/manual/test_srt_engine_with_quant_args.py b/test/manual/test_srt_engine_with_quant_args.py index 78049741c..79a5917aa 100644 --- a/test/manual/test_srt_engine_with_quant_args.py +++ b/test/manual/test_srt_engine_with_quant_args.py @@ -5,7 +5,6 @@ from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST, CustomTest class TestSRTEngineWithQuantArgs(CustomTestCase): - def test_1_quantization_args(self): # we only test fp8 because other methods are currently dependent on vllm. We can add other methods back to test after vllm dependency is resolved. diff --git a/test/manual/test_tokenizer_batch_encode.py b/test/manual/test_tokenizer_batch_encode.py index 31d877304..715c280df 100644 --- a/test/manual/test_tokenizer_batch_encode.py +++ b/test/manual/test_tokenizer_batch_encode.py @@ -38,7 +38,6 @@ class TestTokenizerBatchEncode(unittest.TestCase): "sglang.srt.utils.hf_transformers_utils.get_tokenizer" ) as mock_tokenizer, ): - mock_tokenizer.return_value = Mock(vocab_size=32000) publish(self.server_args, role="tokenizer") self.tokenizer_manager = TokenizerManager(self.server_args, self.port_args) diff --git a/test/manual/test_triton_attention_rocm_mla.py b/test/manual/test_triton_attention_rocm_mla.py index 7c89d521b..5755cee63 100644 --- a/test/manual/test_triton_attention_rocm_mla.py +++ b/test/manual/test_triton_attention_rocm_mla.py @@ -14,7 +14,6 @@ from sglang.test.test_utils import CustomTestCase class TestTritonAttentionMLA(CustomTestCase): - def _set_all_seeds(self, seed): """Set all random seeds for reproducibility.""" random.seed(seed) diff --git a/test/manual/test_vlm_accuracy.py b/test/manual/test_vlm_accuracy.py index 88fdfec01..3fe13bde0 100644 --- a/test/manual/test_vlm_accuracy.py +++ b/test/manual/test_vlm_accuracy.py @@ -212,8 +212,7 @@ class TestMiniCPMV2_6Logits(VisionLLMLogitsBase): # per image if len(pixel_b) != len(tgt_b): raise ValueError( - "Inconsistent N lengths, found: " - f"{len(pixel_b)} vs {len(tgt_b)}" + f"Inconsistent N lengths, found: {len(pixel_b)} vs {len(tgt_b)}" ) for pixel_n, tgt_n in zip(pixel_b, tgt_b): pixel_values_flat += [pixel_n] diff --git a/test/manual/test_weight_cache_e2e.py b/test/manual/test_weight_cache_e2e.py index 5cd757f11..7cdf628c0 100644 --- a/test/manual/test_weight_cache_e2e.py +++ b/test/manual/test_weight_cache_e2e.py @@ -266,7 +266,7 @@ def main(): print( f"Daemon pp_rank={pp_rank} tp_rank={tp_rank} ready " - f"({time.time()-start:.0f}s)" + f"({time.time() - start:.0f}s)" ) if error_found: break @@ -278,7 +278,7 @@ def main(): sys.exit(1) print( - f"\nAll {total_ranks} daemons ready! Total load time: {time.time()-start:.1f}s" + f"\nAll {total_ranks} daemons ready! Total load time: {time.time() - start:.1f}s" ) # Query config from daemon (pp_rank=0, tp_rank=0) diff --git a/test/manual/test_whisper_cuda_graph.py b/test/manual/test_whisper_cuda_graph.py index 72d6da16b..25146e241 100644 --- a/test/manual/test_whisper_cuda_graph.py +++ b/test/manual/test_whisper_cuda_graph.py @@ -97,14 +97,14 @@ class TestWhisperCudaGraph(CustomTestCase): result = self._transcribe() self.assertIn("text", result) results.append(result["text"]) - print(f"Request {i+1}: {result['text'][:80]}...") + print(f"Request {i + 1}: {result['text'][:80]}...") # All transcriptions of the same audio should be identical for i in range(1, len(results)): self.assertEqual( results[0], results[i], - f"Transcription {i+1} differs from first transcription", + f"Transcription {i + 1} differs from first transcription", ) def test_transcription_quality(self): diff --git a/test/manual/vlm/verify_single_image_gather.py b/test/manual/vlm/verify_single_image_gather.py index 5c298e200..e7af929d1 100755 --- a/test/manual/vlm/verify_single_image_gather.py +++ b/test/manual/vlm/verify_single_image_gather.py @@ -76,9 +76,9 @@ def main(): if rank == 0: print( f"world={world} owner={owner} shape=[{n_tok},{hidden}] " - f"(~{n_tok*hidden*2/1e6:.0f}MB) | all_ranks_bitwise_ok={bool(flags.item())} " + f"(~{n_tok * hidden * 2 / 1e6:.0f}MB) | all_ranks_bitwise_ok={bool(flags.item())} " f"(A==truth={eq_truth} A==B={eq_ab}) | all_gather {ta:.3f}ms " - f"broadcast {tb:.3f}ms speedup {ta/tb:.2f}x", + f"broadcast {tb:.3f}ms speedup {ta / tb:.2f}x", flush=True, ) dist.destroy_process_group() diff --git a/test/registered/8-gpu-models/test_deepseek_v32_indexcache.py b/test/registered/8-gpu-models/test_deepseek_v32_indexcache.py index f710996b0..dbe61e287 100644 --- a/test/registered/8-gpu-models/test_deepseek_v32_indexcache.py +++ b/test/registered/8-gpu-models/test_deepseek_v32_indexcache.py @@ -61,7 +61,7 @@ class TestDeepseekV32IndexTopkPattern(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k (deepseek-v32)\n" f'{metrics["accuracy"]=:.3f}\n' + f'### test_gsm8k (deepseek-v32)\n{metrics["accuracy"]=:.3f}\n' ) self.assertGreater(metrics["accuracy"], 0.935) @@ -108,7 +108,7 @@ class TestDeepseekV32IndexFreq(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k (deepseek-v32)\n" f'{metrics["accuracy"]=:.3f}\n' + f'### test_gsm8k (deepseek-v32)\n{metrics["accuracy"]=:.3f}\n' ) self.assertGreater(metrics["accuracy"], 0.935) diff --git a/test/registered/amd/accuracy/mi30x/test_deepseek_r1_eval_amd.py b/test/registered/amd/accuracy/mi30x/test_deepseek_r1_eval_amd.py index cdd8cf2a2..d2740823c 100644 --- a/test/registered/amd/accuracy/mi30x/test_deepseek_r1_eval_amd.py +++ b/test/registered/amd/accuracy/mi30x/test_deepseek_r1_eval_amd.py @@ -246,9 +246,9 @@ class TestDeepSeekR1EvalAMD(unittest.TestCase): for config in self.models: display_name = config.get_display_name() with self.subTest(model=display_name): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {display_name}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi30x/test_deepseek_v32_dp_eval_amd.py b/test/registered/amd/accuracy/mi30x/test_deepseek_v32_dp_eval_amd.py index 44004abdf..556d95e18 100644 --- a/test/registered/amd/accuracy/mi30x/test_deepseek_v32_dp_eval_amd.py +++ b/test/registered/amd/accuracy/mi30x/test_deepseek_v32_dp_eval_amd.py @@ -98,8 +98,7 @@ class TestDeepseekV32DP(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k (deepseek-v32 DP MI325)\n" - f'{metrics["accuracy"]=:.3f}\n' + f'### test_gsm8k (deepseek-v32 DP MI325)\n{metrics["accuracy"]=:.3f}\n' ) self.assertGreater(metrics["accuracy"], GSM8K_ACCURACY_THRESHOLD) @@ -112,8 +111,7 @@ class TestDeepseekV32DP(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_bs_1_speed (deepseek-v32 DP MI325)\n" - f"{speed=:.2f} token/s\n" + f"### test_bs_1_speed (deepseek-v32 DP MI325)\n{speed=:.2f} token/s\n" ) self.assertGreater(speed, 10) diff --git a/test/registered/amd/accuracy/mi30x/test_deepseek_v32_eval_amd.py b/test/registered/amd/accuracy/mi30x/test_deepseek_v32_eval_amd.py index 1a0521591..736435955 100644 --- a/test/registered/amd/accuracy/mi30x/test_deepseek_v32_eval_amd.py +++ b/test/registered/amd/accuracy/mi30x/test_deepseek_v32_eval_amd.py @@ -182,9 +182,9 @@ class TestDeepSeekV32EvalAMD(unittest.TestCase): for config in self.models: display_name = config.get_display_name() with self.subTest(model=display_name): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {display_name}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi30x/test_deepseek_v32_tc_eval_amd.py b/test/registered/amd/accuracy/mi30x/test_deepseek_v32_tc_eval_amd.py index 8e464e181..ecab84b8a 100644 --- a/test/registered/amd/accuracy/mi30x/test_deepseek_v32_tc_eval_amd.py +++ b/test/registered/amd/accuracy/mi30x/test_deepseek_v32_tc_eval_amd.py @@ -99,8 +99,7 @@ class TestDeepseekV32TC(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k (deepseek-v32 TC MI325)\n" - f'{metrics["accuracy"]=:.3f}\n' + f'### test_gsm8k (deepseek-v32 TC MI325)\n{metrics["accuracy"]=:.3f}\n' ) self.assertGreater(metrics["accuracy"], GSM8K_ACCURACY_THRESHOLD) @@ -113,8 +112,7 @@ class TestDeepseekV32TC(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_bs_1_speed (deepseek-v32 TC MI325)\n" - f"{speed=:.2f} token/s\n" + f"### test_bs_1_speed (deepseek-v32 TC MI325)\n{speed=:.2f} token/s\n" ) self.assertGreater(speed, 10) diff --git a/test/registered/amd/accuracy/mi30x/test_glm51_eval_amd.py b/test/registered/amd/accuracy/mi30x/test_glm51_eval_amd.py index ee5948a10..60a724bba 100644 --- a/test/registered/amd/accuracy/mi30x/test_glm51_eval_amd.py +++ b/test/registered/amd/accuracy/mi30x/test_glm51_eval_amd.py @@ -172,9 +172,9 @@ class TestGLM51EvalAMD(unittest.TestCase): for config in self.models: display_name = config.get_display_name() with self.subTest(model=display_name): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {display_name}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi30x/test_glm51_hisparse_eval_mi30x.py b/test/registered/amd/accuracy/mi30x/test_glm51_hisparse_eval_mi30x.py index cca2da359..9f7391c9e 100644 --- a/test/registered/amd/accuracy/mi30x/test_glm51_hisparse_eval_mi30x.py +++ b/test/registered/amd/accuracy/mi30x/test_glm51_hisparse_eval_mi30x.py @@ -80,8 +80,7 @@ class TestGLM51HiSparseEvalAMD(unittest.TestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k (glm-5.1 hisparse mi30x)\n" - f'{metrics["score"]=:.3f}\n' + f'### test_gsm8k (glm-5.1 hisparse mi30x)\n{metrics["score"]=:.3f}\n' ) self.assertGreater(metrics["score"], 0.93) diff --git a/test/registered/amd/accuracy/mi30x/test_glm5_eval_amd.py b/test/registered/amd/accuracy/mi30x/test_glm5_eval_amd.py index 3295dce46..1971edc9b 100644 --- a/test/registered/amd/accuracy/mi30x/test_glm5_eval_amd.py +++ b/test/registered/amd/accuracy/mi30x/test_glm5_eval_amd.py @@ -182,9 +182,9 @@ class TestGLM5EvalAMD(unittest.TestCase): for config in self.models: display_name = config.get_display_name() with self.subTest(model=display_name): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {display_name}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi30x/test_gpt_oss_eval_amd.py b/test/registered/amd/accuracy/mi30x/test_gpt_oss_eval_amd.py index 1651c5682..cd514e530 100644 --- a/test/registered/amd/accuracy/mi30x/test_gpt_oss_eval_amd.py +++ b/test/registered/amd/accuracy/mi30x/test_gpt_oss_eval_amd.py @@ -184,9 +184,9 @@ class TestGptOssEvalAMD(unittest.TestCase): for config in self.models: with self.subTest(model=config.model_path): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {config.model_path}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi30x/test_grok_eval_amd.py b/test/registered/amd/accuracy/mi30x/test_grok_eval_amd.py index 065cc2c17..89ed2ae66 100644 --- a/test/registered/amd/accuracy/mi30x/test_grok_eval_amd.py +++ b/test/registered/amd/accuracy/mi30x/test_grok_eval_amd.py @@ -222,9 +222,9 @@ class TestGrokEvalAMD(unittest.TestCase): for config in self.models: with self.subTest(model=config.model_path): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {config.model_path}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi30x/test_gsm8k_eval_amd.py b/test/registered/amd/accuracy/mi30x/test_gsm8k_eval_amd.py index ed8ceb739..e826450a8 100644 --- a/test/registered/amd/accuracy/mi30x/test_gsm8k_eval_amd.py +++ b/test/registered/amd/accuracy/mi30x/test_gsm8k_eval_amd.py @@ -181,9 +181,9 @@ def check_model_scores(results): line = f"| {model} | {tp_size} | {score:.3f} | {threshold_str} | {startup_str} | {eval_str} | {total_str} | {status} |\n" summary += line - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print("SUMMARY - TP=2 Instruction Models (gsm8k)") - print(f"{'='*60}") + print(f"{'=' * 60}") print(summary) print(f"\n📊 Final Statistics:") print(f" Passed: {passed_count}") @@ -219,19 +219,19 @@ class TestNightlyGsm8KEval(unittest.TestCase): all_results = [] total_test_start = time.time() - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print("AMD GSM8K Evaluation Test (TP=2 Instruction Models)") - print(f"{'='*60}") + print(f"{'=' * 60}") print(f"Benchmark: gsm8k (chat completions)") - print(f"{'='*60}\n") + print(f"{'=' * 60}\n") for model_group, is_fp8, is_tp2 in self.model_groups: for model in model_group: with self.subTest(model=model): tp_size = 2 if is_tp2 else 1 - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {model} (TP={tp_size}, FP8={is_fp8})") - print(f"{'='*60}") + print(f"{'=' * 60}") model_start = time.time() startup_time = None @@ -326,7 +326,7 @@ class TestNightlyGsm8KEval(unittest.TestCase): # Check all scores after collecting all results check_model_scores(all_results) print( - f"\n⏱️ Total test runtime: {total_test_time:.1f}s ({total_test_time/60:.1f} min)" + f"\n⏱️ Total test runtime: {total_test_time:.1f}s ({total_test_time / 60:.1f} min)" ) diff --git a/test/registered/amd/accuracy/mi30x/test_minimax_m25_eval_amd.py b/test/registered/amd/accuracy/mi30x/test_minimax_m25_eval_amd.py index 6ee047e83..20d49ffb6 100644 --- a/test/registered/amd/accuracy/mi30x/test_minimax_m25_eval_amd.py +++ b/test/registered/amd/accuracy/mi30x/test_minimax_m25_eval_amd.py @@ -179,9 +179,9 @@ class TestMiniMaxM25EvalAMD(unittest.TestCase): for config in self.models: display_name = config.get_display_name() with self.subTest(model=display_name): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {display_name}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi30x/test_minimax_m27_eval_amd.py b/test/registered/amd/accuracy/mi30x/test_minimax_m27_eval_amd.py index c33e93eb8..e3795d346 100644 --- a/test/registered/amd/accuracy/mi30x/test_minimax_m27_eval_amd.py +++ b/test/registered/amd/accuracy/mi30x/test_minimax_m27_eval_amd.py @@ -179,9 +179,9 @@ class TestMiniMaxM27EvalAMD(unittest.TestCase): for config in self.models: display_name = config.get_display_name() with self.subTest(model=display_name): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {display_name}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi30x/test_vlms_mmmu_eval_amd.py b/test/registered/amd/accuracy/mi30x/test_vlms_mmmu_eval_amd.py index 6cb9b8807..13eb87ff4 100644 --- a/test/registered/amd/accuracy/mi30x/test_vlms_mmmu_eval_amd.py +++ b/test/registered/amd/accuracy/mi30x/test_vlms_mmmu_eval_amd.py @@ -172,14 +172,14 @@ class TestNightlyVLMMmmuEvalAMD(unittest.TestCase): all_results = [] total_test_start = time.time() - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print("AMD VLM MMMU Evaluation Test") - print(f"{'='*60}") + print(f"{'=' * 60}") print(f"Benchmark: MMMU (100 samples)") print(f"Models to test: {len(self.models)}") for m in self.models: print(f" - {m['model_path']} (TP={m['tp_size']})") - print(f"{'='*60}\n") + print(f"{'=' * 60}\n") for model_config in self.models: model_path = model_config["model_path"] @@ -189,9 +189,9 @@ class TestNightlyVLMMmmuEvalAMD(unittest.TestCase): error_message = None with self.subTest(model=model_path): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {model_path} (TP={tp_size})") - print(f"{'='*60}") + print(f"{'=' * 60}") model_start = time.time() startup_time = None @@ -358,21 +358,21 @@ class TestNightlyVLMMmmuEvalAMD(unittest.TestCase): summary += f"| {model} | {tp_size} | {score_str} | {threshold:.2f} | {startup_str} | {eval_str} | {total_str} | {status} |\n" - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print("SUMMARY - AMD VLM MMMU Evaluation") - print(f"{'='*60}") + print(f"{'=' * 60}") print(summary) print(f"\n📊 Final Statistics:") print(f" Passed: {passed_count}") print(f" Failed: {failed_count}") print( - f"\n⏱️ Total test runtime: {total_test_time:.1f}s ({total_test_time/60:.1f} min)" + f"\n⏱️ Total test runtime: {total_test_time:.1f}s ({total_test_time / 60:.1f} min)" ) if is_in_ci(): write_github_step_summary( f"### TestNightlyVLMMmmuEvalAMD\n{summary}\n\n" - f"**Total Runtime:** {total_test_time:.1f}s ({total_test_time/60:.1f} min)" + f"**Total Runtime:** {total_test_time:.1f}s ({total_test_time / 60:.1f} min)" ) if failed_models: diff --git a/test/registered/amd/accuracy/mi35x/test_deepseek_r1_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_deepseek_r1_eval_mi35x.py index 5b363d03a..ddd26d654 100644 --- a/test/registered/amd/accuracy/mi35x/test_deepseek_r1_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_deepseek_r1_eval_mi35x.py @@ -182,9 +182,9 @@ class TestDeepSeekR1EvalMI35x(unittest.TestCase): for config in self.models: display_name = config.get_display_name() with self.subTest(model=display_name): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {display_name}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi35x/test_deepseek_r1_mxfp4_ar_fusion_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_deepseek_r1_mxfp4_ar_fusion_eval_mi35x.py index 2e19fd800..81433281d 100644 --- a/test/registered/amd/accuracy/mi35x/test_deepseek_r1_mxfp4_ar_fusion_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_deepseek_r1_mxfp4_ar_fusion_eval_mi35x.py @@ -180,9 +180,9 @@ class TestDeepSeekR1MXFP4ArFusionEvalMI35x(unittest.TestCase): for config in self.models: display_name = config.get_display_name() with self.subTest(model=display_name): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {display_name}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi35x/test_deepseek_r1_mxfp4_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_deepseek_r1_mxfp4_eval_mi35x.py index f33d67f1c..e6a05d00d 100644 --- a/test/registered/amd/accuracy/mi35x/test_deepseek_r1_mxfp4_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_deepseek_r1_mxfp4_eval_mi35x.py @@ -178,9 +178,9 @@ class TestDeepSeekR1MXFP4EvalMI35x(unittest.TestCase): for config in self.models: display_name = config.get_display_name() with self.subTest(model=display_name): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {display_name}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi35x/test_deepseek_r1_mxfp4_kv_fp8_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_deepseek_r1_mxfp4_kv_fp8_eval_mi35x.py index d41fb5d1a..6f5b04390 100644 --- a/test/registered/amd/accuracy/mi35x/test_deepseek_r1_mxfp4_kv_fp8_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_deepseek_r1_mxfp4_kv_fp8_eval_mi35x.py @@ -181,9 +181,9 @@ class TestDeepSeekR1MXFP4KvFp8EvalMI35x(unittest.TestCase): for config in self.models: display_name = config.get_display_name() with self.subTest(model=display_name): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {display_name}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi35x/test_deepseek_v32_dp_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_deepseek_v32_dp_eval_mi35x.py index 74c11f572..19bcdb145 100644 --- a/test/registered/amd/accuracy/mi35x/test_deepseek_v32_dp_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_deepseek_v32_dp_eval_mi35x.py @@ -89,8 +89,7 @@ class TestDeepseekV32DP(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k (deepseek-v32 DP MI35x)\n" - f'{metrics["accuracy"]=:.3f}\n' + f'### test_gsm8k (deepseek-v32 DP MI35x)\n{metrics["accuracy"]=:.3f}\n' ) self.assertGreater(metrics["accuracy"], GSM8K_ACCURACY_THRESHOLD) @@ -103,8 +102,7 @@ class TestDeepseekV32DP(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_bs_1_speed (deepseek-v32 DP MI35x)\n" - f"{speed=:.2f} token/s\n" + f"### test_bs_1_speed (deepseek-v32 DP MI35x)\n{speed=:.2f} token/s\n" ) self.assertGreater(speed, 10) diff --git a/test/registered/amd/accuracy/mi35x/test_deepseek_v32_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_deepseek_v32_eval_mi35x.py index 4ac12931f..59af819a2 100644 --- a/test/registered/amd/accuracy/mi35x/test_deepseek_v32_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_deepseek_v32_eval_mi35x.py @@ -183,9 +183,9 @@ class TestDeepSeekV32EvalMI35x(unittest.TestCase): for config in self.models: display_name = config.get_display_name() with self.subTest(model=display_name): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {display_name}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi35x/test_glm51_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_glm51_eval_mi35x.py index e6a88e2e6..8e44f2f7a 100644 --- a/test/registered/amd/accuracy/mi35x/test_glm51_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_glm51_eval_mi35x.py @@ -172,9 +172,9 @@ class TestGLM51EvalMI35x(unittest.TestCase): for config in self.models: display_name = config.get_display_name() with self.subTest(model=display_name): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {display_name}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi35x/test_glm51_hisparse_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_glm51_hisparse_eval_mi35x.py index 9dc8a0740..00b9b2fba 100644 --- a/test/registered/amd/accuracy/mi35x/test_glm51_hisparse_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_glm51_hisparse_eval_mi35x.py @@ -80,8 +80,7 @@ class TestGLM51HiSparseEvalMI35x(unittest.TestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k (glm-5.1 hisparse mi35x)\n" - f'{metrics["score"]=:.3f}\n' + f'### test_gsm8k (glm-5.1 hisparse mi35x)\n{metrics["score"]=:.3f}\n' ) self.assertGreater(metrics["score"], 0.93) diff --git a/test/registered/amd/accuracy/mi35x/test_glm5_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_glm5_eval_mi35x.py index 4c1801091..5d9ae4505 100644 --- a/test/registered/amd/accuracy/mi35x/test_glm5_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_glm5_eval_mi35x.py @@ -182,9 +182,9 @@ class TestGLM5EvalMI35x(unittest.TestCase): for config in self.models: display_name = config.get_display_name() with self.subTest(model=display_name): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {display_name}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi35x/test_glm5_mxfp4_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_glm5_mxfp4_eval_mi35x.py index 82b422bb1..3d7ae621f 100644 --- a/test/registered/amd/accuracy/mi35x/test_glm5_mxfp4_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_glm5_mxfp4_eval_mi35x.py @@ -185,9 +185,9 @@ class TestGLM5MXFP4EvalMI35x(unittest.TestCase): for config in self.models: display_name = config.get_display_name() with self.subTest(model=display_name): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {display_name}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi35x/test_gpt_oss_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_gpt_oss_eval_mi35x.py index d371bb3a5..48483945b 100644 --- a/test/registered/amd/accuracy/mi35x/test_gpt_oss_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_gpt_oss_eval_mi35x.py @@ -196,9 +196,9 @@ class TestGptOssEvalMI35x(unittest.TestCase): for config in self.models: with self.subTest(model=config.model_path): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {config.model_path}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi35x/test_gpt_oss_w4a8_mxfp4_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_gpt_oss_w4a8_mxfp4_eval_mi35x.py index e19b02dd9..0461281f6 100644 --- a/test/registered/amd/accuracy/mi35x/test_gpt_oss_w4a8_mxfp4_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_gpt_oss_w4a8_mxfp4_eval_mi35x.py @@ -185,9 +185,9 @@ class TestGptOssW4A8Mxfp4EvalMI35x(unittest.TestCase): for config in self.models: with self.subTest(model=config.model_path): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {config.model_path}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi35x/test_kimi_k25_aiter_mla_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_kimi_k25_aiter_mla_eval_mi35x.py index dd4a6d862..7c726cf3c 100644 --- a/test/registered/amd/accuracy/mi35x/test_kimi_k25_aiter_mla_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_kimi_k25_aiter_mla_eval_mi35x.py @@ -132,9 +132,9 @@ class TestKimiK25AiterMlaEvalMI35x(unittest.TestCase): for config in self.models: display_name = config.get_display_name() with self.subTest(model=display_name): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {display_name}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi35x/test_kimi_k25_mxfp4_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_kimi_k25_mxfp4_eval_mi35x.py index 44ee20410..084dedd80 100644 --- a/test/registered/amd/accuracy/mi35x/test_kimi_k25_mxfp4_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_kimi_k25_mxfp4_eval_mi35x.py @@ -125,9 +125,9 @@ class TestKimiK25MXFP4AiterMlaEvalMI35x(unittest.TestCase): for config in self.models: display_name = config.get_display_name() with self.subTest(model=display_name): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {display_name}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi35x/test_kimi_k27_code_mxfp4_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_kimi_k27_code_mxfp4_eval_mi35x.py index 42d0e5166..bc08a4711 100644 --- a/test/registered/amd/accuracy/mi35x/test_kimi_k27_code_mxfp4_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_kimi_k27_code_mxfp4_eval_mi35x.py @@ -109,9 +109,9 @@ class TestKimiK27CodeMXFP4AiterMlaEvalMI35x(unittest.TestCase): for config in self.models: display_name = config.get_display_name() with self.subTest(model=display_name): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {display_name}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi35x/test_minimax_m25_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_minimax_m25_eval_mi35x.py index b229cdae5..a88ad93f0 100644 --- a/test/registered/amd/accuracy/mi35x/test_minimax_m25_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_minimax_m25_eval_mi35x.py @@ -179,9 +179,9 @@ class TestMiniMaxM25EvalMI35x(unittest.TestCase): for config in self.models: display_name = config.get_display_name() with self.subTest(model=display_name): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {display_name}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi35x/test_minimax_m25_tp4_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_minimax_m25_tp4_eval_mi35x.py index 95f76a138..30e30666f 100644 --- a/test/registered/amd/accuracy/mi35x/test_minimax_m25_tp4_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_minimax_m25_tp4_eval_mi35x.py @@ -188,9 +188,9 @@ class TestMiniMaxM25TP4EvalMI35x(unittest.TestCase): for config in self.models: display_name = config.get_display_name() with self.subTest(model=display_name): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {display_name}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi35x/test_minimax_m27_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_minimax_m27_eval_mi35x.py index d2703f3f3..b61219d2d 100644 --- a/test/registered/amd/accuracy/mi35x/test_minimax_m27_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_minimax_m27_eval_mi35x.py @@ -179,9 +179,9 @@ class TestMiniMaxM27EvalMI35x(unittest.TestCase): for config in self.models: display_name = config.get_display_name() with self.subTest(model=display_name): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {display_name}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi35x/test_minimax_m3_tp4_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_minimax_m3_tp4_eval_mi35x.py index 00fc99b11..7cc72382b 100644 --- a/test/registered/amd/accuracy/mi35x/test_minimax_m3_tp4_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_minimax_m3_tp4_eval_mi35x.py @@ -212,9 +212,9 @@ class TestMiniMaxM3TP4EvalMI35x(unittest.TestCase): for config in self.models: display_name = config.get_display_name() with self.subTest(model=display_name): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {display_name}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi35x/test_qwen35_mxfp4_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_qwen35_mxfp4_eval_mi35x.py index 458dab6a2..3b4c9cf63 100644 --- a/test/registered/amd/accuracy/mi35x/test_qwen35_mxfp4_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_qwen35_mxfp4_eval_mi35x.py @@ -196,7 +196,7 @@ class TestQwen35Mxfp4MI35x(CustomTestCase): if is_in_ci(): write_github_step_summary( f"### gsm8k accuracy ({MXFP4_MODEL_PATH})\n" - f'score={metrics["score"]:.3f} ' + f"score={metrics['score']:.3f} " f"(threshold {MXFP4_ACC_THRESHOLD})\n" ) self.assertGreater(metrics["score"], MXFP4_ACC_THRESHOLD) @@ -233,7 +233,7 @@ class TestQwen35MoeMxfp4MI35x(CustomTestCase): if is_in_ci(): write_github_step_summary( f"### gsm8k run-through ({MOE_MXFP4_MODEL_PATH}, fuse_gate)\n" - f'score={metrics["score"]:.3f} (run-through, no gate)\n' + f"score={metrics['score']:.3f} (run-through, no gate)\n" ) # Ran e2e and returned parseable answers -> fuse_gate path is healthy. self.assertGreater(metrics["score"], 0.0) diff --git a/test/registered/amd/accuracy/mi35x/test_qwen3_coder_next_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_qwen3_coder_next_eval_mi35x.py index daedb6bb1..f253ccddb 100644 --- a/test/registered/amd/accuracy/mi35x/test_qwen3_coder_next_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_qwen3_coder_next_eval_mi35x.py @@ -202,9 +202,9 @@ class TestQwen3CoderNextEvalMI35x(unittest.TestCase): for config in self.models: display_name = config.get_display_name() with self.subTest(model=display_name): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {display_name}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi35x/test_qwen3_moe_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_qwen3_moe_eval_mi35x.py index a3195eb9b..e68c5d4aa 100644 --- a/test/registered/amd/accuracy/mi35x/test_qwen3_moe_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_qwen3_moe_eval_mi35x.py @@ -164,9 +164,9 @@ class TestQwen3MoeEvalMI35x(unittest.TestCase): for config in self.models: with self.subTest(model=config.model_path): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {config.model_path}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi45x/test_deepseek_r1_0528_mxfp4_eval_mi45x.py b/test/registered/amd/accuracy/mi45x/test_deepseek_r1_0528_mxfp4_eval_mi45x.py index 95d3983e8..3c79531a4 100644 --- a/test/registered/amd/accuracy/mi45x/test_deepseek_r1_0528_mxfp4_eval_mi45x.py +++ b/test/registered/amd/accuracy/mi45x/test_deepseek_r1_0528_mxfp4_eval_mi45x.py @@ -193,9 +193,9 @@ class TestDeepSeekR10528MXFP4EvalMI45x(unittest.TestCase): for config in self.models: with self.subTest(model=config.model_path): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {config.model_path}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi45x/test_deepseek_v4_flash_eval_mi45x.py b/test/registered/amd/accuracy/mi45x/test_deepseek_v4_flash_eval_mi45x.py index dce843b8a..f0b84774c 100644 --- a/test/registered/amd/accuracy/mi45x/test_deepseek_v4_flash_eval_mi45x.py +++ b/test/registered/amd/accuracy/mi45x/test_deepseek_v4_flash_eval_mi45x.py @@ -196,9 +196,9 @@ class TestDeepSeekV4FlashEvalMI45x(unittest.TestCase): for config in self.models: with self.subTest(model=config.model_path): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {config.model_path}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/accuracy/mi45x/test_gpt_oss_w4a8_mxfp4_eval_mi45x.py b/test/registered/amd/accuracy/mi45x/test_gpt_oss_w4a8_mxfp4_eval_mi45x.py index 127c93a39..fac070060 100644 --- a/test/registered/amd/accuracy/mi45x/test_gpt_oss_w4a8_mxfp4_eval_mi45x.py +++ b/test/registered/amd/accuracy/mi45x/test_gpt_oss_w4a8_mxfp4_eval_mi45x.py @@ -182,9 +182,9 @@ class TestGptOssW4A8Mxfp4EvalMI45x(unittest.TestCase): for config in self.models: with self.subTest(model=config.model_path): - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing: {config.model_path}") - print(f"{'='*60}") + print(f"{'=' * 60}") env = os.environ.copy() for key, value in config.env_vars.items(): diff --git a/test/registered/amd/disaggregation/test_disaggregation_basic.py b/test/registered/amd/disaggregation/test_disaggregation_basic.py index 97bfbf208..0229af441 100644 --- a/test/registered/amd/disaggregation/test_disaggregation_basic.py +++ b/test/registered/amd/disaggregation/test_disaggregation_basic.py @@ -134,12 +134,12 @@ class TestDisaggregationAccuracy(PDDisaggregationServerBase): input_logprobs = j["meta_info"]["input_token_logprobs"] output_logprobs = j["meta_info"]["output_token_logprobs"] - assert ( - len(output_logprobs) == completion_tokens - ), f"output_logprobs and completion_tokens should have the same length, but got {len(output_logprobs)} and {completion_tokens}" - assert ( - len(input_logprobs) > 0 - ), f"input_logprobs should have at least one token, but got {len(input_logprobs)}" + assert len(output_logprobs) == completion_tokens, ( + f"output_logprobs and completion_tokens should have the same length, but got {len(output_logprobs)} and {completion_tokens}" + ) + assert len(input_logprobs) > 0, ( + f"input_logprobs should have at least one token, but got {len(input_logprobs)}" + ) def test_structured_output(self): json_schema = json.dumps( diff --git a/test/registered/amd/test_deepseek_r1_mxfp4_8gpu.py b/test/registered/amd/test_deepseek_r1_mxfp4_8gpu.py index 15921eb33..a6a980e06 100644 --- a/test/registered/amd/test_deepseek_r1_mxfp4_8gpu.py +++ b/test/registered/amd/test_deepseek_r1_mxfp4_8gpu.py @@ -76,7 +76,7 @@ class TestDeepseekR1MXFP4(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k (deepseek-r1-mxfp4)\n" f'{metrics["accuracy"]=:.3f}\n' + f'### test_gsm8k (deepseek-r1-mxfp4)\n{metrics["accuracy"]=:.3f}\n' ) self.assertGreater(metrics["accuracy"], 0.94) @@ -88,7 +88,7 @@ class TestDeepseekR1MXFP4(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_bs_1_speed (deepseek-r1-mxfp4)\n" f"{speed=:.2f} token/s\n" + f"### test_bs_1_speed (deepseek-r1-mxfp4)\n{speed=:.2f} token/s\n" ) self.assertGreater(speed, 75) diff --git a/test/registered/amd/test_deepseek_v32_basic.py b/test/registered/amd/test_deepseek_v32_basic.py index dbdd4620d..e4118c32a 100644 --- a/test/registered/amd/test_deepseek_v32_basic.py +++ b/test/registered/amd/test_deepseek_v32_basic.py @@ -72,7 +72,7 @@ class TestDeepseekV32DP(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k (deepseek-v32)\n" f'{metrics["accuracy"]=:.3f}\n' + f'### test_gsm8k (deepseek-v32)\n{metrics["accuracy"]=:.3f}\n' ) self.assertGreater(metrics["accuracy"], 0.935) @@ -84,7 +84,7 @@ class TestDeepseekV32DP(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_bs_1_speed (deepseek-v32)\n" f"{speed=:.2f} token/s\n" + f"### test_bs_1_speed (deepseek-v32)\n{speed=:.2f} token/s\n" ) if is_in_amd_ci(): self.assertGreater(speed, 10) @@ -140,7 +140,7 @@ class TestDeepseekV32TP(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k (deepseek-v32)\n" f'{metrics["accuracy"]=:.3f}\n' + f'### test_gsm8k (deepseek-v32)\n{metrics["accuracy"]=:.3f}\n' ) self.assertGreater(metrics["accuracy"], 0.935) @@ -152,7 +152,7 @@ class TestDeepseekV32TP(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_bs_1_speed (deepseek-v32)\n" f"{speed=:.2f} token/s\n" + f"### test_bs_1_speed (deepseek-v32)\n{speed=:.2f} token/s\n" ) if is_in_amd_ci(): self.assertGreater(speed, 15) diff --git a/test/registered/amd/test_deepseek_v3_basic.py b/test/registered/amd/test_deepseek_v3_basic.py index e87952c72..dfa8ae351 100644 --- a/test/registered/amd/test_deepseek_v3_basic.py +++ b/test/registered/amd/test_deepseek_v3_basic.py @@ -60,7 +60,7 @@ class TestDeepseekV3Basic(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k (deepseek-v3)\n" f'{metrics["accuracy"]=:.3f}\n' + f'### test_gsm8k (deepseek-v3)\n{metrics["accuracy"]=:.3f}\n' ) self.assertGreater(metrics["accuracy"], 0.935) @@ -72,7 +72,7 @@ class TestDeepseekV3Basic(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_bs_1_speed (deepseek-v3)\n" f"{speed=:.2f} token/s\n" + f"### test_bs_1_speed (deepseek-v3)\n{speed=:.2f} token/s\n" ) if is_in_amd_ci(): self.assertGreater(speed, 12) diff --git a/test/registered/amd/test_deepseek_v3_basic_kv_fp8.py b/test/registered/amd/test_deepseek_v3_basic_kv_fp8.py index 601c07cee..173066144 100644 --- a/test/registered/amd/test_deepseek_v3_basic_kv_fp8.py +++ b/test/registered/amd/test_deepseek_v3_basic_kv_fp8.py @@ -64,7 +64,7 @@ class TestDeepseekV3BasicKvFp8(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k (deepseek-v3 kv-fp8)\n" f'{metrics["accuracy"]=:.3f}\n' + f'### test_gsm8k (deepseek-v3 kv-fp8)\n{metrics["accuracy"]=:.3f}\n' ) self.assertGreater(metrics["accuracy"], 0.93) @@ -76,7 +76,7 @@ class TestDeepseekV3BasicKvFp8(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_bs_1_speed (deepseek-v3 kv-fp8)\n" f"{speed=:.2f} token/s\n" + f"### test_bs_1_speed (deepseek-v3 kv-fp8)\n{speed=:.2f} token/s\n" ) if is_in_amd_ci(): self.assertGreater(speed, 40) diff --git a/test/registered/amd/test_kimi_k25_mxfp4.py b/test/registered/amd/test_kimi_k25_mxfp4.py index e48dd8b2d..7c8ac2dc4 100644 --- a/test/registered/amd/test_kimi_k25_mxfp4.py +++ b/test/registered/amd/test_kimi_k25_mxfp4.py @@ -93,7 +93,7 @@ class TestKimiK25MXFP4(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k (Kimi-K2.5-MXFP4)\n" f'{metrics["accuracy"]=:.3f}\n' + f'### test_gsm8k (Kimi-K2.5-MXFP4)\n{metrics["accuracy"]=:.3f}\n' ) self.assertGreater(metrics["accuracy"], 0.92) @@ -105,7 +105,7 @@ class TestKimiK25MXFP4(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_bs_1_speed (Kimi-K2.5-MXFP4)\n" f"{speed=:.2f} token/s\n" + f"### test_bs_1_speed (Kimi-K2.5-MXFP4)\n{speed=:.2f} token/s\n" ) if is_in_amd_ci(): self.assertGreater(speed, 30) diff --git a/test/registered/amd/test_kimi_k2_instruct.py b/test/registered/amd/test_kimi_k2_instruct.py index bb2c57c3d..5e99ac576 100644 --- a/test/registered/amd/test_kimi_k2_instruct.py +++ b/test/registered/amd/test_kimi_k2_instruct.py @@ -73,8 +73,7 @@ class TestKimiK2Instruct0905(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k (Kimi-K2-Instruct-0905)\n" - f'{metrics["accuracy"]=:.3f}\n' + f'### test_gsm8k (Kimi-K2-Instruct-0905)\n{metrics["accuracy"]=:.3f}\n' ) self.assertGreater(metrics["accuracy"], 0.94) @@ -86,8 +85,7 @@ class TestKimiK2Instruct0905(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_bs_1_speed (Kimi-K2-Instruct-0905)\n" - f"{speed=:.2f} token/s\n" + f"### test_bs_1_speed (Kimi-K2-Instruct-0905)\n{speed=:.2f} token/s\n" ) if is_in_amd_ci(): self.assertGreater(speed, 30) diff --git a/test/registered/amd/test_moriep_small.py b/test/registered/amd/test_moriep_small.py index c5808a220..8e46c4cf6 100644 --- a/test/registered/amd/test_moriep_small.py +++ b/test/registered/amd/test_moriep_small.py @@ -103,7 +103,6 @@ mtp_args = [ class TestPureDP(CustomTestCase): - @classmethod def setUpClass(cls): cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST @@ -148,7 +147,6 @@ class TestPureDP(CustomTestCase): class TestMTP(CustomTestCase): - @classmethod def setUpClass(cls): cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST @@ -199,7 +197,6 @@ class TestMTP(CustomTestCase): class TestNormal(CustomTestCase): - @classmethod def setUpClass(cls): cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST @@ -247,7 +244,6 @@ class TestNormal(CustomTestCase): class TestLowLatency(CustomTestCase): - @classmethod def setUpClass(cls): cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST @@ -297,7 +293,6 @@ class TestLowLatency(CustomTestCase): class TestTBOwithNormal(CustomTestCase): - @classmethod def setUpClass(cls): cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST @@ -346,7 +341,6 @@ class TestTBOwithNormal(CustomTestCase): class TestTBOwithLowLatency(CustomTestCase): - @classmethod def setUpClass(cls): cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST @@ -397,7 +391,6 @@ class TestTBOwithLowLatency(CustomTestCase): class TestMTPwithTBONormal(CustomTestCase): - @classmethod def setUpClass(cls): cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST @@ -457,7 +450,6 @@ class TestMTPwithTBONormal(CustomTestCase): class TestMTPwithTBOLowLatency(CustomTestCase): - @classmethod def setUpClass(cls): cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST diff --git a/test/registered/amd/test_qwen3_coder_next_8gpu.py b/test/registered/amd/test_qwen3_coder_next_8gpu.py index 28644ea62..e8a30e5aa 100644 --- a/test/registered/amd/test_qwen3_coder_next_8gpu.py +++ b/test/registered/amd/test_qwen3_coder_next_8gpu.py @@ -82,7 +82,7 @@ class TestQwen3CoderNext(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k (qwen3-coder-next)\n" f'{metrics["accuracy"]=:.3f}\n' + f'### test_gsm8k (qwen3-coder-next)\n{metrics["accuracy"]=:.3f}\n' ) self.assertGreater(metrics["accuracy"], 0.90) @@ -95,7 +95,7 @@ class TestQwen3CoderNext(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_bs_1_speed (qwen3-coder-next)\n" f"{speed=:.2f} token/s\n" + f"### test_bs_1_speed (qwen3-coder-next)\n{speed=:.2f} token/s\n" ) # self.assertGreater(speed, 50) diff --git a/test/registered/amd/test_qwen3_instruct.py b/test/registered/amd/test_qwen3_instruct.py index eacd1d74b..7f9f41c33 100644 --- a/test/registered/amd/test_qwen3_instruct.py +++ b/test/registered/amd/test_qwen3_instruct.py @@ -68,7 +68,7 @@ class TestQwen3Instruct2507(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k ({self.model})\n" f'{metrics["accuracy"]=:.3f}\n' + f'### test_gsm8k ({self.model})\n{metrics["accuracy"]=:.3f}\n' ) self.assertGreater(metrics["accuracy"], 0.95) @@ -80,7 +80,7 @@ class TestQwen3Instruct2507(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_bs_1_speed ({self.model})\n" f"{speed=:.2f} token/s\n" + f"### test_bs_1_speed ({self.model})\n{speed=:.2f} token/s\n" ) if is_in_amd_ci(): self.assertGreater(speed, 50) diff --git a/test/registered/amd/test_qwen3_instruct_fp8.py b/test/registered/amd/test_qwen3_instruct_fp8.py index 7c2978e2f..3dc257abc 100644 --- a/test/registered/amd/test_qwen3_instruct_fp8.py +++ b/test/registered/amd/test_qwen3_instruct_fp8.py @@ -68,7 +68,7 @@ class TestQwen3Instruct2507FP8(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k ({self.model})\n" f'{metrics["accuracy"]=:.3f}\n' + f'### test_gsm8k ({self.model})\n{metrics["accuracy"]=:.3f}\n' ) self.assertGreater(metrics["accuracy"], 0.95) @@ -80,7 +80,7 @@ class TestQwen3Instruct2507FP8(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_bs_1_speed ({self.model})\n" f"{speed=:.2f} token/s\n" + f"### test_bs_1_speed ({self.model})\n{speed=:.2f} token/s\n" ) if is_in_amd_ci(): self.assertGreater(speed, 40) diff --git a/test/registered/amd/test_qwen3_instruct_mxfp4.py b/test/registered/amd/test_qwen3_instruct_mxfp4.py index f1e170946..7505555b7 100644 --- a/test/registered/amd/test_qwen3_instruct_mxfp4.py +++ b/test/registered/amd/test_qwen3_instruct_mxfp4.py @@ -72,7 +72,7 @@ class TestQwen3Instruct2507MXFP4(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k ({self.model})\n" f'{metrics["accuracy"]=:.3f}\n' + f'### test_gsm8k ({self.model})\n{metrics["accuracy"]=:.3f}\n' ) self.assertGreater(metrics["accuracy"], 0.93) @@ -84,7 +84,7 @@ class TestQwen3Instruct2507MXFP4(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_bs_1_speed ({self.model})\n" f"{speed=:.2f} token/s\n" + f"### test_bs_1_speed ({self.model})\n{speed=:.2f} token/s\n" ) if is_in_amd_ci(): self.assertGreater(speed, 60) diff --git a/test/registered/attention/test_kda_decode_flashinfer.py b/test/registered/attention/test_kda_decode_flashinfer.py index fbea645eb..3d15c9fcc 100644 --- a/test/registered/attention/test_kda_decode_flashinfer.py +++ b/test/registered/attention/test_kda_decode_flashinfer.py @@ -193,9 +193,9 @@ def test_kda_decode_flashinfer_matches_triton(batch_size): idx = d["cache_indices"].long() s_err = (st_fi[idx].float() - st_ref[idx].float()).abs() assert s_err.max().item() < 1e-1, f"decode state max diff {s_err.max().item():.2e}" - assert ( - s_err.mean().item() < 1e-2 - ), f"decode state mean diff {s_err.mean().item():.2e}" + assert s_err.mean().item() < 1e-2, ( + f"decode state mean diff {s_err.mean().item():.2e}" + ) @pytest.mark.parametrize("batch_size,num_spec", [(1, 7), (8, 7), (32, 3)]) @@ -248,12 +248,12 @@ def test_kda_target_verify_flashinfer_checkpoint_states( assert torch.isfinite(got).all(), "verify checkpoint states have non-finite values" s_err = (got - ref_states).abs() # bf16 recurrent state; same tolerance as the decode committed-state check. - assert ( - s_err.max().item() < 1e-1 - ), f"checkpoint state max diff {s_err.max().item():.2e}" - assert ( - s_err.mean().item() < 1e-2 - ), f"checkpoint state mean diff {s_err.mean().item():.2e}" + assert s_err.max().item() < 1e-1, ( + f"checkpoint state max diff {s_err.max().item():.2e}" + ) + assert s_err.mean().item() < 1e-2, ( + f"checkpoint state mean diff {s_err.mean().item():.2e}" + ) def test_kda_target_verify_flashinfer_rejects_tree_spec(): diff --git a/test/registered/attention/test_kda_prefill_flashkda.py b/test/registered/attention/test_kda_prefill_flashkda.py index 215843063..13d4e9d70 100644 --- a/test/registered/attention/test_kda_prefill_flashkda.py +++ b/test/registered/attention/test_kda_prefill_flashkda.py @@ -127,9 +127,9 @@ def test_flashkda_matches_triton_safe_gate(seq_lens): # bf16 cross-implementation noise (chunk=16 CUTLASS vs chunk=64 Triton); # measured cos ~0.985 output / ~0.9999 state on H20-3e and B200. assert _cos(ref_out, out) > 0.95, f"output cos too low: {_cos(ref_out, out):.4f}" - assert ( - _cos(ref_state, st_fk[d["idx"]]) > 0.99 - ), f"state cos too low: {_cos(ref_state, st_fk[d['idx']]):.4f}" + assert _cos(ref_state, st_fk[d["idx"]]) > 0.99, ( + f"state cos too low: {_cos(ref_state, st_fk[d['idx']]):.4f}" + ) def test_flashkda_falls_back_without_lower_bound(): @@ -191,9 +191,9 @@ def test_flashkda_spec_verify_falls_back(): assert torch.isfinite(out).all() # Took the Triton fallback (not FlashKDA) -> matches chunk_kda closely. If # FlashKDA had run, the cross-impl cos would be ~0.985 and this would fail. - assert ( - _cos(ref_out, out) > 0.999 - ), f"spec-decode did not fall back: {_cos(ref_out, out):.4f}" + assert _cos(ref_out, out) > 0.999, ( + f"spec-decode did not fall back: {_cos(ref_out, out):.4f}" + ) if __name__ == "__main__": diff --git a/test/registered/attention/test_triton_attention_kernels.py b/test/registered/attention/test_triton_attention_kernels.py index ff90b20a7..01fdc595e 100644 --- a/test/registered/attention/test_triton_attention_kernels.py +++ b/test/registered/attention/test_triton_attention_kernels.py @@ -152,7 +152,6 @@ def decode_attention_fwd_torch( class TestTritonAttention(CustomTestCase): - def _set_all_seeds(self, seed): """Set all random seeds for reproducibility.""" random.seed(seed) diff --git a/test/registered/attention/test_triton_sliding_window.py b/test/registered/attention/test_triton_sliding_window.py index 86eb0e569..e23b1e528 100644 --- a/test/registered/attention/test_triton_sliding_window.py +++ b/test/registered/attention/test_triton_sliding_window.py @@ -42,9 +42,12 @@ class TestSlidingWindowAttentionTriton(CustomTestCase): cls.short_context_prompt = "The capital of France is" # Test prompt longer than window size - cls.long_context_prompt = """ + cls.long_context_prompt = ( + """ Once upon a time, there was a mountain. In the mountain, there was a temple. In the temple, there was an old monk telling a story. The story was: - """ * 100 + """ + * 100 + ) cls.long_context_prompt += "\nNow, summarize the story in one sentence:" def _test_mmlu(self): diff --git a/test/registered/attention/test_wave_attention_kernels.py b/test/registered/attention/test_wave_attention_kernels.py index ebbd4f391..2734a4450 100644 --- a/test/registered/attention/test_wave_attention_kernels.py +++ b/test/registered/attention/test_wave_attention_kernels.py @@ -29,7 +29,6 @@ register_amd_ci(est_time=60, suite="stage-a-test-1-gpu-small-amd") class TestWaveAttention(unittest.TestCase): - def _set_all_seeds(self, seed): """Set all random seeds for reproducibility.""" random.seed(seed) diff --git a/test/registered/attention/unittests/dense/test_extend_init_contract.py b/test/registered/attention/unittests/dense/test_extend_init_contract.py index 62e621392..4c3e6a921 100644 --- a/test/registered/attention/unittests/dense/test_extend_init_contract.py +++ b/test/registered/attention/unittests/dense/test_extend_init_contract.py @@ -52,7 +52,6 @@ _EXTEND_CASE = DenseAttentionCase( @unittest.skipIf(not torch.cuda.is_available(), "CUDA is required") class TestExtendInitContract(CustomTestCase): - def _make_case(self, backend: str) -> DenseAttentionCase: return DenseAttentionCase( name=f"extend_no_prefix_{backend}", diff --git a/test/registered/attention/unittests/dense/test_trtllm_mha.py b/test/registered/attention/unittests/dense/test_trtllm_mha.py index 87b0fcc84..0d12a6c08 100644 --- a/test/registered/attention/unittests/dense/test_trtllm_mha.py +++ b/test/registered/attention/unittests/dense/test_trtllm_mha.py @@ -179,9 +179,10 @@ class TestTRTLLMMHADenseAttentionBackendCorrectness(CustomTestCase): def test_projected_dense_decode_cases(self): for case_index, case in enumerate(self.DECODE_CASES): splits = 2 if case_index == 0 else 1 - with self.subTest( - case=case.name, backend=case.backend - ), envs.SGLANG_TRTLLM_MHA_DECODE_SEQ_LEN_SPLITS.override(splits): + with ( + self.subTest(case=case.name, backend=case.backend), + envs.SGLANG_TRTLLM_MHA_DECODE_SEQ_LEN_SPLITS.override(splits), + ): run_dense_attention_case( self, case, @@ -230,9 +231,10 @@ class TestTRTLLMMHADenseAttentionBackendCorrectness(CustomTestCase): def test_runner_mode_frozen_kv_mtp_cuda_graph_runner_cases(self): for case in self.FROZEN_KV_MTP_RUNNER_CASES: - with self.subTest( - case=case.name, backend=case.backend - ), envs.SGLANG_TRTLLM_MHA_DECODE_SEQ_LEN_SPLITS.override(2): + with ( + self.subTest(case=case.name, backend=case.backend), + envs.SGLANG_TRTLLM_MHA_DECODE_SEQ_LEN_SPLITS.override(2), + ): run_dense_frozen_kv_mtp_cuda_graph_runner_case( self, case, diff --git a/test/registered/attention/unittests/mla/test_flashmla.py b/test/registered/attention/unittests/mla/test_flashmla.py index d8a4bf7e0..c45e54032 100644 --- a/test/registered/attention/unittests/mla/test_flashmla.py +++ b/test/registered/attention/unittests/mla/test_flashmla.py @@ -379,8 +379,9 @@ class TestFlashMLAAttentionBackendCorrectness(CustomTestCase): ) fixture = self._build_target_verify_metadata_fixture(case) - with torch.no_grad(), forward_context( - ForwardContext(attn_backend=fixture.backend) + with ( + torch.no_grad(), + forward_context(ForwardContext(attn_backend=fixture.backend)), ): fixture.backend.init_forward_metadata(fixture.forward_batch) diff --git a/test/registered/backends/test_flashinfer_fusion_preflight.py b/test/registered/backends/test_flashinfer_fusion_preflight.py index ea06d0124..a0636a9cf 100644 --- a/test/registered/backends/test_flashinfer_fusion_preflight.py +++ b/test/registered/backends/test_flashinfer_fusion_preflight.py @@ -62,9 +62,7 @@ def _run_rank(rank, world_size, port, scenario, result_q): free, _total = torch.cuda.mem_get_info(rank) target = max(free - (1 << 30), 0) - granularity_flag = ( - cuda_driver.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED - ) + granularity_flag = cuda_driver.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED err, gran = cuda_driver.cuMemGetAllocationGranularity( prop, granularity_flag, diff --git a/test/registered/core/test_srt_endpoint.py b/test/registered/core/test_srt_endpoint.py index 2d67f2bff..bb82925ce 100644 --- a/test/registered/core/test_srt_endpoint.py +++ b/test/registered/core/test_srt_endpoint.py @@ -259,7 +259,6 @@ class TestSRTEndpoint(CustomTestCase): for logprob_start_len in [0, 500, 2500, 5000, 25000]: for return_logprob in [True, False]: for top_logprobs_num in [0, 5]: - if logprob_start_len >= input_len: continue diff --git a/test/registered/core/test_srt_engine.py b/test/registered/core/test_srt_engine.py index c3e86e5d6..a142420bf 100644 --- a/test/registered/core/test_srt_engine.py +++ b/test/registered/core/test_srt_engine.py @@ -25,7 +25,6 @@ register_amd_ci(est_time=261, suite="stage-b-test-1-gpu-small-amd") class TestSRTEngine(CustomTestCase): - def test_1_engine_runtime_consistency(self): prompt = "Today is a sunny day and I like" model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST diff --git a/test/registered/cp/test_dsa_prefill_cp.py b/test/registered/cp/test_dsa_prefill_cp.py index 357360d0e..1aff2df21 100644 --- a/test/registered/cp/test_dsa_prefill_cp.py +++ b/test/registered/cp/test_dsa_prefill_cp.py @@ -78,8 +78,7 @@ class TestDSACPV2Interleave(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_a_gsm8k (dsa-cp-v2-interleave)\n" - f'{metrics["score"]=:.3f}\n' + f'### test_a_gsm8k (dsa-cp-v2-interleave)\n{metrics["score"]=:.3f}\n' ) self.assertGreater(metrics["score"], 0.935) diff --git a/test/registered/cpu/test_cpu_graph.py b/test/registered/cpu/test_cpu_graph.py index 0ed2bcaf2..9d624c597 100644 --- a/test/registered/cpu/test_cpu_graph.py +++ b/test/registered/cpu/test_cpu_graph.py @@ -25,7 +25,6 @@ register_cpu_ci(est_time=315, suite="base-b-tp-test-cpu") class TestCPUGraph(CustomTestCase): - @intel_amx_benchmark( extra_args=[ "--batch-size", diff --git a/test/registered/cpu/test_extend.py b/test/registered/cpu/test_extend.py index 21d6474fb..40807a920 100644 --- a/test/registered/cpu/test_extend.py +++ b/test/registered/cpu/test_extend.py @@ -71,7 +71,6 @@ class TestExtendAttention(CustomTestCase): start_q, start_kv = 0, 0 for seq_idx in range(seq_lens.shape[0]): - extend_seq_len_q = extend_seq_lens[seq_idx] prefill_seq_len_q = extend_prefix_lens[seq_idx] diff --git a/test/registered/cpu/test_flash_attn.py b/test/registered/cpu/test_flash_attn.py index f2efd2eb7..1ac9f7a5a 100644 --- a/test/registered/cpu/test_flash_attn.py +++ b/test/registered/cpu/test_flash_attn.py @@ -81,7 +81,6 @@ def flash_attn_non_varlen_ref( class TestFlashAttn(CustomTestCase): - @parametrize( batch=[4], max_seqlen_q=[35, 96], diff --git a/test/registered/cpu/test_gemm.py b/test/registered/cpu/test_gemm.py index 64dfe34dc..3011f6374 100644 --- a/test/registered/cpu/test_gemm.py +++ b/test/registered/cpu/test_gemm.py @@ -33,7 +33,6 @@ class Mod(nn.Module): class TestGemm(CustomTestCase): - @parametrize( M=[1, 101], N=[16, 32 * 13], diff --git a/test/registered/cpu/test_intel_amx_attention_backend_a.py b/test/registered/cpu/test_intel_amx_attention_backend_a.py index a5ddf21a7..7fe178812 100644 --- a/test/registered/cpu/test_intel_amx_attention_backend_a.py +++ b/test/registered/cpu/test_intel_amx_attention_backend_a.py @@ -25,7 +25,6 @@ register_cpu_ci(est_time=685, suite="base-b-tp-test-cpu") class TestIntelAMXAttnBackend(CustomTestCase): - @intel_amx_benchmark( extra_args=["--batch-size", "4", "--mem-fraction-static", "0.3"], min_throughput=10, diff --git a/test/registered/cpu/test_intel_amx_attention_backend_b.py b/test/registered/cpu/test_intel_amx_attention_backend_b.py index fcea8fe71..9c088237c 100644 --- a/test/registered/cpu/test_intel_amx_attention_backend_b.py +++ b/test/registered/cpu/test_intel_amx_attention_backend_b.py @@ -18,7 +18,6 @@ register_cpu_ci(est_time=47, suite="base-b-test-cpu") class TestIntelAMXAttnBackendQuant(CustomTestCase): - @intel_amx_benchmark( extra_args=["--batch-size", "4", "--mem-fraction-static", "0.3"], min_throughput=150, diff --git a/test/registered/cpu/test_intel_amx_attention_backend_c.py b/test/registered/cpu/test_intel_amx_attention_backend_c.py index ce1621f59..3445b2d09 100644 --- a/test/registered/cpu/test_intel_amx_attention_backend_c.py +++ b/test/registered/cpu/test_intel_amx_attention_backend_c.py @@ -18,7 +18,6 @@ register_cpu_ci(est_time=477, suite="base-b-tp-test-cpu") class TestIntelAMXAttnBackendQuant(CustomTestCase): - @intel_amx_benchmark( extra_args=[ "--batch-size", diff --git a/test/registered/cpu/test_moe.py b/test/registered/cpu/test_moe.py index 1c7fc7f26..4811f51f3 100644 --- a/test/registered/cpu/test_moe.py +++ b/test/registered/cpu/test_moe.py @@ -140,7 +140,6 @@ def make_mxfp4_weights(e, out_dim, in_dim, dtype, with_bias=False): class TestFusedExperts: - def test_unsupported_activation_is_rejected(self): m, n, k, e, topk = 2, 32, 32, 4, 2 a = torch.randn((m, k), dtype=dtype) / 10 diff --git a/test/registered/cpu/test_norm.py b/test/registered/cpu/test_norm.py index 0ebfbb818..f03825d09 100644 --- a/test/registered/cpu/test_norm.py +++ b/test/registered/cpu/test_norm.py @@ -19,7 +19,6 @@ eps = 1e-6 class TestNorm: - def _forward_native( self, x: torch.Tensor, @@ -201,7 +200,6 @@ class TestNorm: class TestFusedRMSNormGated: - def _forward_native( self, hidden_states: torch.Tensor, @@ -236,7 +234,6 @@ class TestFusedRMSNormGated: class TestFusedQKRMSNorm: - @pytest.mark.parametrize("dtype", DTYPES, ids=DTYPE_IDS) @pytest.mark.parametrize( "batch_size,q_size,k_size,v_size", @@ -317,7 +314,6 @@ class TestFusedQKRMSNorm: class TestLayerNorm: - def _forward_native( self, x: torch.Tensor, @@ -382,7 +378,6 @@ class TestLayerNorm: class TestFusedQKGemmaRMSNorm: - def _gemma_rmsnorm_per_head_native( self, x: torch.Tensor, diff --git a/test/registered/cpu/test_rank_consensus_checker.py b/test/registered/cpu/test_rank_consensus_checker.py index 254baa618..0bc7b421e 100644 --- a/test/registered/cpu/test_rank_consensus_checker.py +++ b/test/registered/cpu/test_rank_consensus_checker.py @@ -226,9 +226,9 @@ class TestAssertSame(RankConsensusCheckerTestCase): err = err_box.get() shutdown() - assert isinstance( - err, RuntimeError - ), f"Expected RuntimeError from stray-thread assert_same, got {err!r}" + assert isinstance(err, RuntimeError), ( + f"Expected RuntimeError from stray-thread assert_same, got {err!r}" + ) def test_assert_same_rejects_non_scheduler_thread(self): """Check that assert_same() must be called in the scheduler thread. Otherwise report error.""" diff --git a/test/registered/cpu/test_server_args_backend.py b/test/registered/cpu/test_server_args_backend.py index 9822397eb..5c710fa84 100644 --- a/test/registered/cpu/test_server_args_backend.py +++ b/test/registered/cpu/test_server_args_backend.py @@ -56,17 +56,20 @@ class TestServerArgsIBDeviceValidation(unittest.TestCase): real_isdir = os.path.isdir real_listdir = os.listdir - with patch( - "sglang.srt.arg_groups.validation_hook.os.path.isdir", - side_effect=lambda path: ( - True if path == "/sys/class/infiniband" else real_isdir(path) + with ( + patch( + "sglang.srt.arg_groups.validation_hook.os.path.isdir", + side_effect=lambda path: ( + True if path == "/sys/class/infiniband" else real_isdir(path) + ), ), - ), patch( - "sglang.srt.arg_groups.validation_hook.os.listdir", - side_effect=lambda path: ( - available_devices - if path == "/sys/class/infiniband" - else real_listdir(path) + patch( + "sglang.srt.arg_groups.validation_hook.os.listdir", + side_effect=lambda path: ( + available_devices + if path == "/sys/class/infiniband" + else real_listdir(path) + ), ), ): return validate_ib_devices(device_str) diff --git a/test/registered/cpu/test_spec_kernels.py b/test/registered/cpu/test_spec_kernels.py index 542c33a5d..bbb9131b2 100644 --- a/test/registered/cpu/test_spec_kernels.py +++ b/test/registered/cpu/test_spec_kernels.py @@ -635,26 +635,122 @@ class TestBuildTreeKernelEfficient(CustomTestCase): # fmt: off torch.tensor( [ - [29889, 29974, 29945, 29900, 29974, 29922, 29930, 29958, - 29889, 29974, 29930, 29945, 29974, 29922, 29930, 29958], - [22550, 4136, 16492, 8439, 29871, 2, 3001, 13, - 2, 13, 29906, 29946, 2, 13, 29871, 259], + [ + 29889, + 29974, + 29945, + 29900, + 29974, + 29922, + 29930, + 29958, + 29889, + 29974, + 29930, + 29945, + 29974, + 29922, + 29930, + 29958, + ], + [ + 22550, + 4136, + 16492, + 8439, + 29871, + 2, + 3001, + 13, + 2, + 13, + 29906, + 29946, + 2, + 13, + 29871, + 259, + ], ], ), torch.tensor( [ - [29946, 29945, 29953, 29906, 29896, 29945, 29900, 29906, - 29896, 29945, 29906, 29953, 29896, 29945, 29906, 29946], - [29871, 2, 29901, 29889, 29871, 2, 395, 259, - 29901, 29871, 2, 29889, 3001, 1234, 7146, 2186], + [ + 29946, + 29945, + 29953, + 29906, + 29896, + 29945, + 29900, + 29906, + 29896, + 29945, + 29906, + 29953, + 29896, + 29945, + 29906, + 29946, + ], + [ + 29871, + 2, + 29901, + 29889, + 29871, + 2, + 395, + 259, + 29901, + 29871, + 2, + 29889, + 3001, + 1234, + 7146, + 2186, + ], ], ), torch.tensor( [ - [29946, 29974, 29945, 29930, 29889, 29922, 29974, 29930, - 29974, 29946, 29930, 29922, 29889, 29974, 29945, 29922], - [29941, 29906, 2, 29946, 29871, 450, 319, 14990, - 29946, 29941, 2, 29906, 29871, 2, 3001, 13], + [ + 29946, + 29974, + 29945, + 29930, + 29889, + 29922, + 29974, + 29930, + 29974, + 29946, + 29930, + 29922, + 29889, + 29974, + 29945, + 29922, + ], + [ + 29941, + 29906, + 2, + 29946, + 29871, + 450, + 319, + 14990, + 29946, + 29941, + 2, + 29906, + 29871, + 2, + 3001, + 13, + ], ], ), # fmt: on @@ -772,15 +868,38 @@ class TestReconstructIndicesFromTreeMask(CustomTestCase): tree_mask = torch.tensor( # fmt: off [ - 1, 0, 0, 0, - 1, 1, 0, 0, - 1, 0, 1, 0, - 1, 0, 1, 1, - - 1, 0, 0, 0, - 1, 1, 0, 0, - 1, 1, 1, 0, - 1, 1, 1, 1, + 1, + 0, + 0, + 0, + 1, + 1, + 0, + 0, + 1, + 0, + 1, + 0, + 1, + 0, + 1, + 1, + 1, + 0, + 0, + 0, + 1, + 1, + 0, + 0, + 1, + 1, + 1, + 0, + 1, + 1, + 1, + 1, ], # fmt: on dtype=torch.bool, diff --git a/test/registered/dcp/test_dcp_layout_unit.py b/test/registered/dcp/test_dcp_layout_unit.py index ee20e0fae..a4b277412 100644 --- a/test/registered/dcp/test_dcp_layout_unit.py +++ b/test/registered/dcp/test_dcp_layout_unit.py @@ -411,10 +411,13 @@ class TestGetDcpLens(CustomTestCase): ) # The allocator widens from get_parallel(), not from the injected # server_args stand-in -- drive the cause, not the effect. - with patch( - "sglang.srt.mem_cache.kv_cache_configurator.current_platform.is_out_of_tree", - return_value=False, - ), rc.get_parallel().override(attn_dcp_size=dcp_size): + with ( + patch( + "sglang.srt.mem_cache.kv_cache_configurator.current_platform.is_out_of_tree", + return_value=False, + ), + rc.get_parallel().override(attn_dcp_size=dcp_size), + ): allocators[dcp_size] = ( KVCacheConfigurator._build_token_to_kv_pool_allocator( configurator, diff --git a/test/registered/dcp/test_qwen3p5_triton_dcp.py b/test/registered/dcp/test_qwen3p5_triton_dcp.py index 030faed71..1a0e48a57 100644 --- a/test/registered/dcp/test_qwen3p5_triton_dcp.py +++ b/test/registered/dcp/test_qwen3p5_triton_dcp.py @@ -87,7 +87,7 @@ class TestQwen35TritonDCPGsm8k(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_a_gsm8k (qwen3.5-triton-dcp4)\n" f'{metrics["score"]=:.3f}\n' + f'### test_a_gsm8k (qwen3.5-triton-dcp4)\n{metrics["score"]=:.3f}\n' ) self.assertGreater(metrics["score"], GSM8K_ACCURACY_THRESHOLD) diff --git a/test/registered/dcp/test_reduce_scatter_along_dim.py b/test/registered/dcp/test_reduce_scatter_along_dim.py index cbad393e6..481b46dc4 100644 --- a/test/registered/dcp/test_reduce_scatter_along_dim.py +++ b/test/registered/dcp/test_reduce_scatter_along_dim.py @@ -85,8 +85,7 @@ def multiprocess_test(file: str, nproc: int, timeout: int = 120) -> None: ) from e assert result.returncode == 0, ( - f"torchrun (nproc={nproc}) failed with rc={result.returncode}\n" - f"{result.stdout}" + f"torchrun (nproc={nproc}) failed with rc={result.returncode}\n{result.stdout}" ) diff --git a/test/registered/debug_utils/comparator/tensor_comparator/test_formatter.py b/test/registered/debug_utils/comparator/tensor_comparator/test_formatter.py index d3c4534da..41fdffb51 100644 --- a/test/registered/debug_utils/comparator/tensor_comparator/test_formatter.py +++ b/test/registered/debug_utils/comparator/tensor_comparator/test_formatter.py @@ -473,8 +473,7 @@ class TestFormatComparisonRichMinimal: assert_rich_tags_balanced(result) assert result == ( - "[red]❌[/] [bold red]hidden_states [/] " - "rel_diff=5.00e-01" + "[red]❌[/] [bold red]hidden_states [/] rel_diff=5.00e-01" ) def test_shape_mismatch(self) -> None: @@ -1046,7 +1045,7 @@ class TestFormatAbsDiffPercentilesRich: result: str = _format_abs_diff_percentiles_rich(diff) assert result == ( - "p1=1.00e-04 p5=1.00e-04 p50=2.00e-04 " "p95=4.00e-04 p99=5.00e-04" + "p1=1.00e-04 p5=1.00e-04 p50=2.00e-04 p95=4.00e-04 p99=5.00e-04" ) def test_high_p99_coloring(self) -> None: @@ -1123,7 +1122,7 @@ class TestFormatReplicatedChecks: result: str = format_replicated_checks(checks) assert result == ( - "Replicated checks:\n" " ✅ axis=tp group=0 idx=1 vs 0: n/a diff" + "Replicated checks:\n ✅ axis=tp group=0 idx=1 vs 0: n/a diff" ) diff --git a/test/registered/debug_utils/comparator/test_entrypoint.py b/test/registered/debug_utils/comparator/test_entrypoint.py index 361b50c3f..422c38cd1 100644 --- a/test/registered/debug_utils/comparator/test_entrypoint.py +++ b/test/registered/debug_utils/comparator/test_entrypoint.py @@ -3299,9 +3299,9 @@ def _create_thd_cp_zigzag_dumps( # Dump each rank for cp_rank in range(cp_size): rank_tensor: torch.Tensor = torch.cat(rank_segments[cp_rank], dim=0) - assert ( - rank_tensor.shape[0] == total_per_rank - ), f"rank {cp_rank}: expected {total_per_rank} tokens, got {rank_tensor.shape[0]}" + assert rank_tensor.shape[0] == total_per_rank, ( + f"rank {cp_rank}: expected {total_per_rank} tokens, got {rank_tensor.shape[0]}" + ) _create_rank_dump( directory, @@ -4008,11 +4008,13 @@ class TestEntrypointMetaOverride: baseline_path, target_path = self._create_single_rank_pair(tmp_path) yaml_path: Path = tmp_path / "override.yaml" - yaml_path.write_text(textwrap.dedent("""\ + yaml_path.write_text( + textwrap.dedent("""\ overrides: - match: "hidden" dims: "t h" - """)) + """) + ) argv = _make_argv( baseline_path, @@ -4148,13 +4150,15 @@ class TestEntrypointMetaOverride: baseline_path, target_path = self._create_single_rank_pair(tmp_path) yaml_path: Path = tmp_path / "override.yaml" - yaml_path.write_text(textwrap.dedent("""\ + yaml_path.write_text( + textwrap.dedent("""\ overrides: - match: "hidden" dims: "t h" - match: "hidden" dims: "a b" - """)) + """) + ) argv = _make_argv( baseline_path, @@ -4169,11 +4173,13 @@ class TestEntrypointMetaOverride: baseline_path, target_path = self._create_single_rank_pair(tmp_path) yaml_path: Path = tmp_path / "override.yaml" - yaml_path.write_text(textwrap.dedent("""\ + yaml_path.write_text( + textwrap.dedent("""\ overrides: - match: "hidden" dims: "a b" - """)) + """) + ) argv = _make_argv( baseline_path, diff --git a/test/registered/debug_utils/comparator/test_meta_overrider.py b/test/registered/debug_utils/comparator/test_meta_overrider.py index 05d23f88f..5e642f49c 100644 --- a/test/registered/debug_utils/comparator/test_meta_overrider.py +++ b/test/registered/debug_utils/comparator/test_meta_overrider.py @@ -194,11 +194,13 @@ class TestFromArgsAndConfig: def test_cli_before_yaml(self, tmp_path: Path) -> None: """CLI rules are ordered before YAML rules (CLI wins on conflict).""" yaml_path = tmp_path / "override.yaml" - yaml_path.write_text(textwrap.dedent("""\ + yaml_path.write_text( + textwrap.dedent("""\ overrides: - match: "hidden" dims: "FROM_YAML" - """)) + """) + ) overrider = MetaOverrider.from_args_and_config( override_dims=["hidden:FROM_CLI"], @@ -256,14 +258,16 @@ class TestLoadYamlRules: def test_valid_yaml(self, tmp_path: Path) -> None: """Valid YAML with override rules loads correctly.""" yaml_path = tmp_path / "override.yaml" - yaml_path.write_text(textwrap.dedent("""\ + yaml_path.write_text( + textwrap.dedent("""\ overrides: - match: "hidden" dims: "b s h d" - match: "logits" dims: "b s v[tp]" side: baseline - """)) + """) + ) rules = _load_yaml_rules(yaml_path) assert len(rules) == 2 assert rules[0].dims == "b s h d" diff --git a/test/registered/debug_utils/comparator/test_output_types.py b/test/registered/debug_utils/comparator/test_output_types.py index d6eec78b4..ba4e9f8a0 100644 --- a/test/registered/debug_utils/comparator/test_output_types.py +++ b/test/registered/debug_utils/comparator/test_output_types.py @@ -591,9 +591,7 @@ class TestFormatAlignerPlan: ) result: str = _format_aligner_plan(_wrap_plan(plan)) - assert result == ( - "Aligner Plan:\n" " baseline: (no steps)\n" " target: (no steps)" - ) + assert result == ("Aligner Plan:\n baseline: (no steps)\n target: (no steps)") def test_unsharder(self) -> None: unsharder: UnsharderPlan = UnsharderPlan( @@ -614,9 +612,7 @@ class TestFormatAlignerPlan: result: str = _format_aligner_plan(_wrap_plan(plan)) assert result == ( - "Aligner Plan:\n" - " baseline: (no steps)\n" - " target: [step=0: unsharder(tp)]" + "Aligner Plan:\n baseline: (no steps)\n target: [step=0: unsharder(tp)]" ) def test_reorderer(self) -> None: diff --git a/test/registered/debug_utils/source_patcher/test_source_editor.py b/test/registered/debug_utils/source_patcher/test_source_editor.py index 38898e6d8..0d95ee53b 100644 --- a/test/registered/debug_utils/source_patcher/test_source_editor.py +++ b/test/registered/debug_utils/source_patcher/test_source_editor.py @@ -12,7 +12,7 @@ class TestApplyEdits: """Tests for the apply_edits() source text transformation function.""" def test_single_line_match_to_multiline_replacement(self) -> None: - source = "def foo():\n" " x = compute()\n" " return x\n" + source = "def foo():\n x = compute()\n return x\n" edits = [ EditSpec( match="x = compute()", @@ -20,12 +20,10 @@ class TestApplyEdits: ) ] result = apply_edits(source=source, edits=edits) - assert result == ( - "def foo():\n" " x = compute()\n" " print(x)\n" " return x\n" - ) + assert result == ("def foo():\n x = compute()\n print(x)\n return x\n") def test_pure_insertion(self) -> None: - source = "def foo():\n" " a = 1\n" " b = 2\n" + source = "def foo():\n a = 1\n b = 2\n" edits = [ EditSpec( match="a = 1", @@ -33,10 +31,10 @@ class TestApplyEdits: ) ] result = apply_edits(source=source, edits=edits) - assert result == ("def foo():\n" " a = 1\n" " print(a)\n" " b = 2\n") + assert result == ("def foo():\n a = 1\n print(a)\n b = 2\n") def test_pure_deletion_via_empty_replacement(self) -> None: - source = "def foo():\n" " debug_log()\n" " return 42\n" + source = "def foo():\n debug_log()\n return 42\n" edits = [ EditSpec( match="debug_log()", @@ -44,10 +42,10 @@ class TestApplyEdits: ) ] result = apply_edits(source=source, edits=edits) - assert result == ("def foo():\n" " return 42\n") + assert result == ("def foo():\n return 42\n") def test_deletion_fewer_lines(self) -> None: - source = "def foo():\n" " a = 1\n" " b = 2\n" " c = 3\n" + source = "def foo():\n a = 1\n b = 2\n c = 3\n" edits = [ EditSpec( match="a = 1\nb = 2", @@ -55,7 +53,7 @@ class TestApplyEdits: ) ] result = apply_edits(source=source, edits=edits) - assert result == ("def foo():\n" " ab = 3\n" " c = 3\n") + assert result == ("def foo():\n ab = 3\n c = 3\n") def test_multiline_match_to_multiline_replacement(self) -> None: source = ( @@ -133,15 +131,7 @@ class TestApplyEdits: def test_not_found_diagnostic_single_window_with_marker(self) -> None: """first line is present once but full match doesn't fit: one window with '>' on the match-region line.""" - source = ( - "line0\n" - "line1\n" - "line2\n" - "anchor()\n" - "wrong_next()\n" - "line5\n" - "line6\n" - ) + source = "line0\nline1\nline2\nanchor()\nwrong_next()\nline5\nline6\n" edits = [EditSpec(match="anchor()\nright_next()", replacement="x")] with pytest.raises(PatchApplicationError) as exc_info: apply_edits(source=source, edits=edits) @@ -232,24 +222,22 @@ class TestApplyEdits: assert "filler9" not in msg def test_match_found_multiple_times_raises(self) -> None: - source = "def foo():\n" " print(1)\n" " print(1)\n" + source = "def foo():\n print(1)\n print(1)\n" edits = [EditSpec(match="print(1)", replacement="print(2)")] with pytest.raises(PatchApplicationError, match="multiple"): apply_edits(source=source, edits=edits) def test_multiple_edits_applied_sequentially(self) -> None: - source = "def foo():\n" " a = 1\n" " b = 2\n" " return a + b\n" + source = "def foo():\n a = 1\n b = 2\n return a + b\n" edits = [ EditSpec(match="a = 1", replacement="a = 10"), EditSpec(match="b = 2", replacement="b = 20"), ] result = apply_edits(source=source, edits=edits) - assert result == ( - "def foo():\n" " a = 10\n" " b = 20\n" " return a + b\n" - ) + assert result == ("def foo():\n a = 10\n b = 20\n return a + b\n") def test_strip_matching_ignores_leading_trailing_whitespace(self) -> None: - source = "def foo():\n" " x = compute()\n" " return x\n" + source = "def foo():\n x = compute()\n return x\n" edits = [ EditSpec( match=" x = compute() ", @@ -257,11 +245,11 @@ class TestApplyEdits: ) ] result = apply_edits(source=source, edits=edits) - assert result == ("def foo():\n" " x = replaced()\n" " return x\n") + assert result == ("def foo():\n x = replaced()\n return x\n") def test_replacement_indented_text_realigned(self) -> None: """replacement text with its own indentation gets realigned to match source.""" - source = "def foo():\n" " x = compute()\n" " return x\n" + source = "def foo():\n x = compute()\n return x\n" edits = [ EditSpec( match="x = compute()", @@ -270,15 +258,12 @@ class TestApplyEdits: ] result = apply_edits(source=source, edits=edits) assert result == ( - "def foo():\n" - " x = compute()\n" - " print(x)\n" - " return x\n" + "def foo():\n x = compute()\n print(x)\n return x\n" ) def test_replacement_with_existing_indent_realigned(self) -> None: """replacement text already has indentation that should be rebased.""" - source = "def foo():\n" " if True:\n" " x = 1\n" " return x\n" + source = "def foo():\n if True:\n x = 1\n return x\n" edits = [ EditSpec( match="x = 1", @@ -296,12 +281,10 @@ class TestApplyEdits: ) def test_append_keeps_match_and_adds_after(self) -> None: - source = "def foo():\n" " x = compute()\n" " return x\n" + source = "def foo():\n x = compute()\n return x\n" edits = [EditSpec(match="x = compute()", append="print(x)")] result = apply_edits(source=source, edits=edits) - assert result == ( - "def foo():\n" " x = compute()\n" " print(x)\n" " return x\n" - ) + assert result == ("def foo():\n x = compute()\n print(x)\n return x\n") def test_append_multiline_match(self) -> None: source = ( @@ -330,21 +313,18 @@ class TestApplyEdits: ) def test_prepend_adds_before_match(self) -> None: - source = "def foo():\n" " x = compute()\n" " return x\n" + source = "def foo():\n x = compute()\n return x\n" edits = [EditSpec(match="x = compute()", prepend="print('before')")] result = apply_edits(source=source, edits=edits) assert result == ( - "def foo():\n" - " print('before')\n" - " x = compute()\n" - " return x\n" + "def foo():\n print('before')\n x = compute()\n return x\n" ) def test_prepend_multiline(self) -> None: - source = "def foo():\n" " return x\n" + source = "def foo():\n return x\n" edits = [EditSpec(match="return x", prepend="a = 1\nb = 2")] result = apply_edits(source=source, edits=edits) - assert result == ("def foo():\n" " a = 1\n" " b = 2\n" " return x\n") + assert result == ("def foo():\n a = 1\n b = 2\n return x\n") def test_prepend_deep_indent(self) -> None: source = ( @@ -365,11 +345,7 @@ class TestApplyEdits: def test_prepend_multiline_match(self) -> None: source = ( - "def foo():\n" - " result = call(\n" - " a=1,\n" - " )\n" - " return result\n" + "def foo():\n result = call(\n a=1,\n )\n return result\n" ) edits = [ EditSpec( @@ -401,13 +377,13 @@ class TestApplyEdits: def test_second_edit_sees_result_of_first(self) -> None: """Edits are applied sequentially; second edit matches modified source.""" - source = "def foo():\n" " x = 1\n" " return x\n" + source = "def foo():\n x = 1\n return x\n" edits = [ EditSpec(match="x = 1", replacement="x = 1\ny = 2"), EditSpec(match="y = 2", replacement="y = 20"), ] result = apply_edits(source=source, edits=edits) - assert result == ("def foo():\n" " x = 1\n" " y = 20\n" " return x\n") + assert result == ("def foo():\n x = 1\n y = 20\n return x\n") if __name__ == "__main__": diff --git a/test/registered/debug_utils/test_dumper.py b/test/registered/debug_utils/test_dumper.py index f5158c086..639c12053 100644 --- a/test/registered/debug_utils/test_dumper.py +++ b/test/registered/debug_utils/test_dumper.py @@ -734,9 +734,9 @@ def _assert_files(filenames, *, exist=(), not_exist=()): for p in exist: assert any(p in f for f in filenames), f"{p} not found in {filenames}" for p in not_exist: - assert not any( - p in f for f in filenames - ), f"{p} should not exist in {filenames}" + assert not any(p in f for f in filenames), ( + f"{p} should not exist in {filenames}" + ) def _load_dump(path: Path) -> dict: @@ -750,9 +750,9 @@ def _find_dump_file(tmpdir, *, rank: int = 0, name: str) -> Path: for f in Path(tmpdir).glob("*/*.pt") if f"rank={rank}" in f.name and name in f.name ] - assert ( - len(matches) == 1 - ), f"Expected 1 file matching rank={rank} name={name}, got {matches}" + assert len(matches) == 1, ( + f"Expected 1 file matching rank={rank} name={name}, got {matches}" + ) return matches[0] @@ -1657,9 +1657,9 @@ class TestZmqPortIsolation: ) resp.raise_for_status() states = resp.json() - assert ( - len(states) == 2 - ), f"Instance {i} (port {port}): expected 2 ranks, got {len(states)}" + assert len(states) == 2, ( + f"Instance {i} (port {port}): expected 2 ranks, got {len(states)}" + ) finally: for event in stop_events: event.set() @@ -1719,9 +1719,9 @@ class TestDumperHttp: val = state for k in keys: val = val[k] - assert ( - val == expected - ), f"rank {rank}: {path}={val!r}, expected {expected!r}" + assert val == expected, ( + f"rank {rank}: {path}={val!r}, expected {expected!r}" + ) def test_configure_enable_toggle(self, dumper_http_url: str): for enable in [True, False]: @@ -1915,9 +1915,9 @@ class TestNonIntrusiveDumper(_NonIntrusiveTestBase): ) dumped_output = captured[f"{P}model.mutator.output"]["value"] - assert ( - dumped_output == 999.0 - ).all(), "post-hook should capture outputs after forward" + assert (dumped_output == 999.0).all(), ( + "post-hook should capture outputs after forward" + ) def test_hooks_all_module_levels(self, tmp_path): class Attention(torch.nn.Module): @@ -2374,9 +2374,9 @@ class TestDumperE2E: states = requests.post(f"{base_url}/dumper/get_state", json={}).json() assert len(states) == 2 for rank, state in enumerate(states): - assert ( - state["config"]["enable"] is True - ), f"rank {rank}: enable should be True after configure" + assert state["config"]["enable"] is True, ( + f"rank {rank}: enable should be True after configure" + ) assert state["config"]["dir"] == dump_dir resp = requests.post( @@ -2403,16 +2403,16 @@ class TestDumperE2E: ) for rank in range(2): - assert any( - f"rank={rank}" in f for f in filenames - ), f"No dump files for rank {rank}" + assert any(f"rank={rank}" in f for f in filenames), ( + f"No dump files for rank {rank}" + ) sample_file = dump_files[0] loaded = torch.load(sample_file, map_location="cpu", weights_only=False) assert isinstance(loaded, dict), f"Expected dict, got {type(loaded)}" - assert ( - "value" in loaded and "meta" in loaded - ), f"Missing value/meta keys: {loaded.keys()}" + assert "value" in loaded and "meta" in loaded, ( + f"Missing value/meta keys: {loaded.keys()}" + ) assert "name" in loaded["meta"] assert "rank" in loaded["meta"] assert "step" in loaded["meta"] @@ -2438,22 +2438,22 @@ class TestDumperE2E: "attn_cp_size", ] for key in expected_keys: - assert ( - key in par - ), f"Missing {key} in sglang_parallel_info, got: {sorted(par)}" + assert key in par, ( + f"Missing {key} in sglang_parallel_info, got: {sorted(par)}" + ) rids_files = [f for f in dump_files if "name=rids" in f.name] rids_loaded = torch.load( rids_files[0], map_location="cpu", weights_only=False ) rids_value = rids_loaded["value"] - assert isinstance( - rids_value, list - ), f"rids should be a list, got {type(rids_value)}" + assert isinstance(rids_value, list), ( + f"rids should be a list, got {type(rids_value)}" + ) assert len(rids_value) > 0, "rids should be non-empty" - assert all( - isinstance(r, str) for r in rids_value - ), f"each rid should be a str, got {[type(r) for r in rids_value]}" + assert all(isinstance(r, str) for r in rids_value), ( + f"each rid should be a str, got {[type(r) for r in rids_value]}" + ) finally: kill_process_tree(proc.pid) @@ -2912,9 +2912,9 @@ class TestRecomputeStatus: model(torch.randn(2, 4)) for key, data in captured.items(): - assert ( - "recompute_status" in data["meta"] - ), f"missing recompute_status in {key}" + assert "recompute_status" in data["meta"], ( + f"missing recompute_status in {key}" + ) assert data["meta"]["recompute_status"] == "disabled" def test_detect_recompute_status_default(self) -> None: @@ -3553,8 +3553,7 @@ class TestGrafterDistributed: # worker prepends tmp_path to sys.path so import_module sees it. module_name = "_xform_user_basic" (tmp_path / f"{module_name}.py").write_text( - "def transform(graft_input):\n" - " return graft_input.received_list[0] * 2\n" + "def transform(graft_input):\n return graft_input.received_list[0] * 2\n" ) graft_port = find_available_port(29610) _run_graft_test( @@ -3646,7 +3645,9 @@ class TestGrafterDistributed: 7.0, 7.0, 7.0, - ], f"target should be unchanged after shape-mismatch graft, got {target.tolist()}" + ], ( + f"target should be unchanged after shape-mismatch graft, got {target.tolist()}" + ) finally: if grafter._pg is not None: dist.destroy_process_group(grafter._pg) @@ -3694,7 +3695,9 @@ class TestGrafterDistributed: 9.0, 9.0, 9.0, - ], f"target must be unchanged when transform throws, got {target.tolist()}" + ], ( + f"target must be unchanged when transform throws, got {target.tolist()}" + ) output = captured.getvalue() assert "transform/copy_ raised RuntimeError" in output, output assert "intentional test error" in output, output @@ -3783,9 +3786,9 @@ class TestGrafterDistributed: grafter.maybe_intercept(value=target, tags={"name": "x"}) output = captured.getvalue() if rank == 0: - assert ( - "WARNING" in output - ), f"expected WARNING in rank 0 output: {output}" + assert "WARNING" in output, ( + f"expected WARNING in rank 0 output: {output}" + ) assert "has not completed after 2s" in output, output finally: if grafter._pg is not None: @@ -3852,9 +3855,9 @@ class TestGrafterDistributed: pg_after_first = grafter._pg assert pg_after_first is not None grafter.maybe_intercept(value=t2, tags={"name": "x"}) - assert ( - grafter._pg is pg_after_first - ), "_pg must be cached across calls, not re-initialized" + assert grafter._pg is pg_after_first, ( + "_pg must be cached across calls, not re-initialized" + ) else: target1 = torch.zeros(3, device="cuda:1") target2 = torch.zeros(3, device="cuda:1") @@ -4190,9 +4193,9 @@ def _e2e_transform(graft_input): the transform is just identity. Real workflows would compute a non-trivial override (scale, reshape, decode, ...) using the extras. """ - assert ( - graft_input.received_extras_list[0]["my_extra_key"] == "my_extra_value" - ), graft_input.received_extras_list + assert graft_input.received_extras_list[0]["my_extra_key"] == "my_extra_value", ( + graft_input.received_extras_list + ) return graft_input.received_list[0] diff --git a/test/registered/debug_utils/test_engine_dumper_comparator_e2e.py b/test/registered/debug_utils/test_engine_dumper_comparator_e2e.py index e38173f37..f6208d1f1 100644 --- a/test/registered/debug_utils/test_engine_dumper_comparator_e2e.py +++ b/test/registered/debug_utils/test_engine_dumper_comparator_e2e.py @@ -290,7 +290,7 @@ def _run_e2e_scenario( print(f"Comparator debug output: {debug_file}") assert result.returncode == 0, ( - f"Comparator failed (rc={result.returncode}). " f"Debug output: {debug_file}" + f"Comparator failed (rc={result.returncode}). Debug output: {debug_file}" ) diff --git a/test/registered/debug_utils/test_tensor_dump_forward_hook.py b/test/registered/debug_utils/test_tensor_dump_forward_hook.py index 5074f6a92..1661df8c3 100644 --- a/test/registered/debug_utils/test_tensor_dump_forward_hook.py +++ b/test/registered/debug_utils/test_tensor_dump_forward_hook.py @@ -31,7 +31,6 @@ TEST_HIDDEN_SIZE = 32 class SimpleModel(nn.Module): - def __init__(self) -> None: super().__init__() self.hidden_size = TEST_HIDDEN_SIZE diff --git a/test/registered/disaggregation/test_disaggregation_basic.py b/test/registered/disaggregation/test_disaggregation_basic.py index ea5a49cad..b994a144e 100644 --- a/test/registered/disaggregation/test_disaggregation_basic.py +++ b/test/registered/disaggregation/test_disaggregation_basic.py @@ -76,12 +76,12 @@ class TestDisaggregationAccuracy(PauseResumeInPlaceMixin, PDDisaggregationServer input_logprobs = j["meta_info"]["input_token_logprobs"] output_logprobs = j["meta_info"]["output_token_logprobs"] - assert ( - len(output_logprobs) == completion_tokens - ), f"output_logprobs and completion_tokens should have the same length, but got {len(output_logprobs)} and {completion_tokens}" - assert ( - len(input_logprobs) > 0 - ), f"input_logprobs should have at least one token, but got {len(input_logprobs)}" + assert len(output_logprobs) == completion_tokens, ( + f"output_logprobs and completion_tokens should have the same length, but got {len(output_logprobs)} and {completion_tokens}" + ) + assert len(input_logprobs) > 0, ( + f"input_logprobs should have at least one token, but got {len(input_logprobs)}" + ) def test_chat_completion_top_logprobs(self): client = openai.Client(api_key="empty", base_url=f"{self.lb_url}/v1") diff --git a/test/registered/disaggregation/test_kimi_linear_pd_dcp4.py b/test/registered/disaggregation/test_kimi_linear_pd_dcp4.py index e30e08672..844f0b319 100644 --- a/test/registered/disaggregation/test_kimi_linear_pd_dcp4.py +++ b/test/registered/disaggregation/test_kimi_linear_pd_dcp4.py @@ -461,8 +461,7 @@ class TestKimiLinearPDDCP4(GSM8KMixin, PDDisaggregationServerBase): reference, actual, label=( - f"niah prompt_tokens={LONG_CONTEXT_TOKENS} " - f"depth={needle_depth}" + f"niah prompt_tokens={LONG_CONTEXT_TOKENS} depth={needle_depth}" ), ) diff --git a/test/registered/dp_attn/test_dp_attention_bcg_kl.py b/test/registered/dp_attn/test_dp_attention_bcg_kl.py index 272783279..b6b7978cd 100644 --- a/test/registered/dp_attn/test_dp_attention_bcg_kl.py +++ b/test/registered/dp_attn/test_dp_attention_bcg_kl.py @@ -138,9 +138,9 @@ def test_prefill_cache_hit(base_url, input_ids, max_new_tokens, is_hicache=False hit_label = "L1 hits" if is_hicache else "cache hits" print(f" {hit_label}: {len(new_input_ids)}/{len(input_ids)}") - assert ( - len(new_input_ids) > len(input_ids) // 2 - ), f"too few {hit_label}: {len(new_input_ids)}/{len(input_ids)}" + assert len(new_input_ids) > len(input_ids) // 2, ( + f"too few {hit_label}: {len(new_input_ids)}/{len(input_ids)}" + ) input_logprobs = _get_input_logprobs(base_url, new_input_ids, output_logprobs) return _compute_kl(input_logprobs, output_logprobs) @@ -170,9 +170,9 @@ def test_decode_cache_hit(base_url, input_ids, max_new_tokens, is_hicache=False) hit_label = "L1 decode hits" if is_hicache else "cache hits" print(f" {hit_label}: {len(new_input_ids)}/{len(turn2_ids)}") - assert ( - len(new_input_ids) > len(turn2_ids) // 2 - ), f"too few {hit_label}: {len(new_input_ids)}/{len(turn2_ids)}" + assert len(new_input_ids) > len(turn2_ids) // 2, ( + f"too few {hit_label}: {len(new_input_ids)}/{len(turn2_ids)}" + ) input_logprobs = _get_input_logprobs(base_url, new_input_ids, output_logprobs) return _compute_kl(input_logprobs, output_logprobs) @@ -281,8 +281,7 @@ class _DPAttentionPrefillCudaGraphKLMixin: return time.sleep(0.5) self.fail( - f"No {self.prefill_backend} prefill CUDA graph replay was logged " - f"for {case}" + f"No {self.prefill_backend} prefill CUDA graph replay was logged for {case}" ) def test_prefill_and_decode_cache_hit_kl_is_zero(self): diff --git a/test/registered/eplb/test_lplb_distributed.py b/test/registered/eplb/test_lplb_distributed.py index 6bdc81aec..444a61a50 100644 --- a/test/registered/eplb/test_lplb_distributed.py +++ b/test/registered/eplb/test_lplb_distributed.py @@ -173,16 +173,16 @@ def test_solve_ipm_matches_torch_reference(): max_diff = (cuda_x - torch_x).abs().max().item() print( f"\n[ipm-compare] converged={converged} max|cuda-torch|={max_diff:.3e} " - f"cuda={[round(v,4) for v in cuda_x.tolist()]} " - f"torch={[round(v,4) for v in torch_x.tolist()]}" + f"cuda={[round(v, 4) for v in cuda_x.tolist()]} " + f"torch={[round(v, 4) for v in torch_x.tolist()]}" ) assert converged, ( "IPM returned the 0.5 non-convergence sentinel — the comparison would " "be trivial. Adjust the LP instance so it converges." ) - assert torch.allclose( - cuda_x, torch_x, atol=1e-2, rtol=1e-2 - ), f"fused IPM diverges from torch reference: max abs diff {max_diff:.3e}" + assert torch.allclose(cuda_x, torch_x, atol=1e-2, rtol=1e-2), ( + f"fused IPM diverges from torch reference: max abs diff {max_diff:.3e}" + ) @pytest.mark.skipif( @@ -363,9 +363,9 @@ def _check_all_ranks_empty(rank: int, world_size: int, device: torch.device): assert torch.isfinite(actual).all(), f"rank {rank}: non-finite for empty-batch" assert (actual >= 0).all(), f"rank {rank}: negative for empty-batch" - assert torch.allclose( - actual, expected, atol=1e-4, rtol=1e-3 - ), f"rank {rank}: empty-batch output disagrees with all-zero oracle" + assert torch.allclose(actual, expected, atol=1e-4, rtol=1e-3), ( + f"rank {rank}: empty-batch output disagrees with all-zero oracle" + ) def _check_solver_determinism(rank: int, world_size: int, device: torch.device): diff --git a/test/registered/function_call/test_kimik2_detector.py b/test/registered/function_call/test_kimik2_detector.py index b3e3011e0..3b182c8f3 100644 --- a/test/registered/function_call/test_kimik2_detector.py +++ b/test/registered/function_call/test_kimik2_detector.py @@ -309,7 +309,7 @@ class TestKimiK2DetectorSpecialTokenLeakage(unittest.TestCase): def test_no_leak_on_error_fallback(self): """On parse errors, normal_text fallback has tokens stripped.""" cleaned = _strip_special_tokens( - "leaked<|tool_calls_section_begin|>" "<|tool_call_end|>content" + "leaked<|tool_calls_section_begin|><|tool_call_end|>content" ) self.assertEqual(cleaned, "leakedcontent") @@ -992,7 +992,7 @@ class TestKimiK2EndToEnd(unittest.TestCase): "<|tool_call_begin|>functions.get_weather:1" f'<|tool_call_argument_begin|>{{"city":', ' "Bad", "valid": fasle', - "<|tool_call_end|>" "<|tool_calls_section_end|>" + good_section_1, + "<|tool_call_end|><|tool_calls_section_end|>" + good_section_1, ], } diff --git a/test/registered/gemm/test_linear_bf16_fp32_hpc.py b/test/registered/gemm/test_linear_bf16_fp32_hpc.py index 7dcc4d2bc..7034a45ff 100644 --- a/test/registered/gemm/test_linear_bf16_fp32_hpc.py +++ b/test/registered/gemm/test_linear_bf16_fp32_hpc.py @@ -31,7 +31,6 @@ _ROUTER_SHAPES = ((6144, 768), (3072, 384)) "requires HPC-Ops (https://github.com/Tencent/hpc-ops) and a Hopper GPU", ) class TestLinearBf16Fp32Hpc(CustomTestCase): - @classmethod def setUpClass(cls): mark_hpc_bf16xfp32_gemm_enabled() diff --git a/test/registered/kernels/benchmark/communication/bench_custom_all_reduce.py b/test/registered/kernels/benchmark/communication/bench_custom_all_reduce.py index f0a4bb5df..d42846f2c 100644 --- a/test/registered/kernels/benchmark/communication/bench_custom_all_reduce.py +++ b/test/registered/kernels/benchmark/communication/bench_custom_all_reduce.py @@ -246,7 +246,7 @@ def benchmark(message_KB: int, provider: str): ) if provider == "aot" and world_size not in AOT_SUPPORTED_WORLD_SIZES: marker.skip( - f"AOT custom_all_reduce needs world_size in " f"{AOT_SUPPORTED_WORLD_SIZES}" + f"AOT custom_all_reduce needs world_size in {AOT_SUPPORTED_WORLD_SIZES}" ) _init_all_backends() backend = BACKEND_FACTORY[provider]() diff --git a/test/registered/kernels/benchmark/communication/bench_tp_qknorm.py b/test/registered/kernels/benchmark/communication/bench_tp_qknorm.py index 0895d65a3..e5f4be66f 100644 --- a/test/registered/kernels/benchmark/communication/bench_tp_qknorm.py +++ b/test/registered/kernels/benchmark/communication/bench_tp_qknorm.py @@ -88,7 +88,7 @@ def _precompile_kernels(num_gpus: List[int]) -> None: p.join() if p.exitcode != 0: raise RuntimeError( - f"TP QKNorm precompile failed for {world_size=} " f"(exit {p.exitcode})" + f"TP QKNorm precompile failed for {world_size=} (exit {p.exitcode})" ) diff --git a/test/registered/kernels/benchmark/diffusion/bench_fused_norm_scale_shift.py b/test/registered/kernels/benchmark/diffusion/bench_fused_norm_scale_shift.py index 56628bb37..ef6d089bb 100644 --- a/test/registered/kernels/benchmark/diffusion/bench_fused_norm_scale_shift.py +++ b/test/registered/kernels/benchmark/diffusion/bench_fused_norm_scale_shift.py @@ -128,12 +128,12 @@ def bench_fused_scale_residual_norm_scale_shift( if __name__ == "__main__": - print(f"\n{'='*80}") + print(f"\n{'=' * 80}") print("Benchmark: fused_norm_scale_shift") - print(f"{'='*80}\n") + print(f"{'=' * 80}\n") bench_fused_norm_scale_shift.run(print_data=True) - print(f"\n{'='*80}") + print(f"\n{'=' * 80}") print("Benchmark: fused_scale_residual_norm_scale_shift") - print(f"{'='*80}\n") + print(f"{'=' * 80}\n") bench_fused_scale_residual_norm_scale_shift.run(print_data=True) diff --git a/test/registered/kernels/benchmark/diffusion/bench_residual_gate_add.py b/test/registered/kernels/benchmark/diffusion/bench_residual_gate_add.py index 9cd3f8674..4e66961ce 100644 --- a/test/registered/kernels/benchmark/diffusion/bench_residual_gate_add.py +++ b/test/registered/kernels/benchmark/diffusion/bench_residual_gate_add.py @@ -81,9 +81,7 @@ def benchmark() -> None: repeats = 5 if is_in_ci() else 20 rounds = 5 if is_in_ci() else 13 - print( - "| workload | gate | torch us | triton us | cuda us | reference | " "ref/cuda |" - ) + print("| workload | gate | torch us | triton us | cuda us | reference | ref/cuda |") print("|---|---|---:|---:|---:|---|---:|") for workload in workloads: diff --git a/test/registered/kernels/benchmark/kv_canary/bench_verify.py b/test/registered/kernels/benchmark/kv_canary/bench_verify.py index 36d33e014..19f7a5e3c 100644 --- a/test/registered/kernels/benchmark/kv_canary/bench_verify.py +++ b/test/registered/kernels/benchmark/kv_canary/bench_verify.py @@ -76,7 +76,9 @@ def _verify_num_slots(case: BenchCase) -> int: return max(2, case.bs * per_req_slots + 1) -def _build_verify_inputs(case: BenchCase, *, device: torch.device) -> Tuple[ +def _build_verify_inputs( + case: BenchCase, *, device: torch.device +) -> Tuple[ torch.Tensor, VerifyPlan, torch.Tensor, diff --git a/test/registered/kernels/ops/activation/test_activation.py b/test/registered/kernels/ops/activation/test_activation.py index ba57d0667..a191cdbdb 100644 --- a/test/registered/kernels/ops/activation/test_activation.py +++ b/test/registered/kernels/ops/activation/test_activation.py @@ -133,9 +133,9 @@ def test_activation_filter_expert( if kept.any(): torch.testing.assert_close(out[kept], expected[kept], atol=atol, rtol=rtol) if token_skip.any(): - assert torch.isnan( - out[token_skip] - ).all(), "filter_expert kernel touched rows whose expert_id is -1" + assert torch.isnan(out[token_skip]).all(), ( + "filter_expert kernel touched rows whose expert_id is -1" + ) @pytest.mark.parametrize("op_name", OPS) diff --git a/test/registered/kernels/ops/attention/test_cutedsl_gdn.py b/test/registered/kernels/ops/attention/test_cutedsl_gdn.py index b7fce1054..4887e7607 100644 --- a/test/registered/kernels/ops/attention/test_cutedsl_gdn.py +++ b/test/registered/kernels/ops/attention/test_cutedsl_gdn.py @@ -267,8 +267,9 @@ def test_cutedsl_gdn_performance(B: int): # Benchmark triton_times, cutedsl_times = [], [] for _ in range(bench_iters): - start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event( - enable_timing=True + start, end = ( + torch.cuda.Event(enable_timing=True), + torch.cuda.Event(enable_timing=True), ) start.record() if graph_triton: @@ -279,8 +280,9 @@ def test_cutedsl_gdn_performance(B: int): torch.cuda.synchronize() triton_times.append(start.elapsed_time(end)) - start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event( - enable_timing=True + start, end = ( + torch.cuda.Event(enable_timing=True), + torch.cuda.Event(enable_timing=True), ) with torch.cuda.stream(torch_stream): start.record() diff --git a/test/registered/kernels/ops/attention/test_dsv32_indexer_fusion.py b/test/registered/kernels/ops/attention/test_dsv32_indexer_fusion.py index f12bc5e98..f8c2039a0 100644 --- a/test/registered/kernels/ops/attention/test_dsv32_indexer_fusion.py +++ b/test/registered/kernels/ops/attention/test_dsv32_indexer_fusion.py @@ -178,9 +178,9 @@ def test_q_rope_quant_matches_reference(pos_dtype): # scale step at the bottom of the range. deq = q_fp8.float() * scale err = (deq - ref).abs() - assert ( - err <= 0.0625 * ref.abs() + scale - ).all(), f"max fp8 dequant error {err.max().item()}" + assert (err <= 0.0625 * ref.abs() + scale).all(), ( + f"max fp8 dequant error {err.max().item()}" + ) # ---------------------------------------------------------------------------- diff --git a/test/registered/kernels/ops/attention/test_flash_attention_4_sm120.py b/test/registered/kernels/ops/attention/test_flash_attention_4_sm120.py index 235856657..986dae9ff 100644 --- a/test/registered/kernels/ops/attention/test_flash_attention_4_sm120.py +++ b/test/registered/kernels/ops/attention/test_flash_attention_4_sm120.py @@ -1965,9 +1965,7 @@ def test_sm120_paged_decode_transpose_is_cache_order_independent(): num_pages, device="cuda", dtype=torch.int64, - ).to( - torch.int32 - )[None] + ).to(torch.int32)[None] cache_seqlens = torch.tensor( [max_seqlen], device="cuda", @@ -2097,9 +2095,7 @@ def test_sm120_paged_decode_graph_pdl_is_correct_and_eager_reusable( num_pages, device="cuda", dtype=torch.int64, - ).to( - torch.int32 - )[None] + ).to(torch.int32)[None] cache_seqlens = torch.tensor( [max_seqlen], device="cuda", diff --git a/test/registered/kernels/ops/attention/test_flash_mla_backends.py b/test/registered/kernels/ops/attention/test_flash_mla_backends.py index 17867fb39..5ccb25be4 100644 --- a/test/registered/kernels/ops/attention/test_flash_mla_backends.py +++ b/test/registered/kernels/ops/attention/test_flash_mla_backends.py @@ -152,9 +152,7 @@ def _build_kvcache( nope_dequant = ( nope_fp8.view(num_pages, page_size, _NUM_TILES, _TILE_SIZE) * scale_e8m0.view(num_pages, page_size, _NUM_TILES, 1) - ).view( - num_pages, page_size, _NOPE_DIM - ) # float32 + ).view(num_pages, page_size, _NOPE_DIM) # float32 ref_per_token = torch.cat( [nope_dequant.to(torch.bfloat16), rope_bf16_vals], dim=-1 ) # (num_pages, page_size, 512) bf16 diff --git a/test/registered/kernels/ops/attention/test_fused_metadata_copy.py b/test/registered/kernels/ops/attention/test_fused_metadata_copy.py index 30197a1e2..805352dda 100644 --- a/test/registered/kernels/ops/attention/test_fused_metadata_copy.py +++ b/test/registered/kernels/ops/attention/test_fused_metadata_copy.py @@ -396,29 +396,29 @@ def test_fused_metadata_copy(bs, forward_mode, has_real_page_table, has_flashmla ) # Compare results - assert torch.equal( - dst_ref["cache_seqlens"], dst_fused["cache_seqlens"] - ), "cache_seqlens mismatch" - assert torch.equal( - dst_ref["cu_seqlens_k"], dst_fused["cu_seqlens_k"] - ), "cu_seqlens_k mismatch" - assert torch.equal( - dst_ref["page_table_1"], dst_fused["page_table_1"] - ), "page_table_1 mismatch" - assert torch.equal( - dst_ref["dsa_cache_seqlens"], dst_fused["dsa_cache_seqlens"] - ), "dsa_cache_seqlens mismatch" + assert torch.equal(dst_ref["cache_seqlens"], dst_fused["cache_seqlens"]), ( + "cache_seqlens mismatch" + ) + assert torch.equal(dst_ref["cu_seqlens_k"], dst_fused["cu_seqlens_k"]), ( + "cu_seqlens_k mismatch" + ) + assert torch.equal(dst_ref["page_table_1"], dst_fused["page_table_1"]), ( + "page_table_1 mismatch" + ) + assert torch.equal(dst_ref["dsa_cache_seqlens"], dst_fused["dsa_cache_seqlens"]), ( + "dsa_cache_seqlens mismatch" + ) assert torch.equal( dst_ref["dsa_seqlens_expanded"], dst_fused["dsa_seqlens_expanded"] ), "dsa_seqlens_expanded mismatch" - assert torch.equal( - dst_ref["dsa_cu_seqlens_k"], dst_fused["dsa_cu_seqlens_k"] - ), "dsa_cu_seqlens_k mismatch" + assert torch.equal(dst_ref["dsa_cu_seqlens_k"], dst_fused["dsa_cu_seqlens_k"]), ( + "dsa_cu_seqlens_k mismatch" + ) if has_real_page_table: - assert torch.equal( - dst_ref["real_page_table"], dst_fused["real_page_table"] - ), "real_page_table mismatch" + assert torch.equal(dst_ref["real_page_table"], dst_fused["real_page_table"]), ( + "real_page_table mismatch" + ) if has_flashmla: assert torch.equal( @@ -643,7 +643,11 @@ def test_fused_metadata_copy_multi_dtype_validation(): # Create source tensors - one with WRONG dtype cache_seqlens_src_wrong = torch.randint( - 1, max_len, (bs,), dtype=torch.int64, device=device # Wrong dtype! + 1, + max_len, + (bs,), + dtype=torch.int64, + device=device, # Wrong dtype! ) cu_seqlens_k_src = torch.zeros(bs + 1, dtype=torch.int32, device=device) page_indices_src = torch.randint( @@ -829,7 +833,7 @@ def test_fused_metadata_copy_multi(bs, has_real_page_table, has_flashmla): f"\n[VERIFY] bs={bs}, real_page_table={has_real_page_table}, flashmla={has_flashmla}" ) print( - f"[VERIFY] Fused time: {fused_time*1000:.3f}ms, Loop time: {loop_time*1000:.3f}ms, Speedup: {speedup:.2f}x" + f"[VERIFY] Fused time: {fused_time * 1000:.3f}ms, Loop time: {loop_time * 1000:.3f}ms, Speedup: {speedup:.2f}x" ) max_diff = 0.0 @@ -1063,7 +1067,7 @@ def test_fused_metadata_copy_multi_large_batch(bs): speedup = loop_time / fused_time if fused_time > 0 else 0 print( - f"\n[PERF] Large batch (bs={bs}): Fused={fused_time*1000:.3f}ms, Loop={loop_time*1000:.3f}ms, Speedup={speedup:.2f}x" + f"\n[PERF] Large batch (bs={bs}): Fused={fused_time * 1000:.3f}ms, Loop={loop_time * 1000:.3f}ms, Speedup={speedup:.2f}x" ) # Verify correctness @@ -1076,9 +1080,9 @@ def test_fused_metadata_copy_multi_large_batch(bs): ): for key in dst_ref: if dst_ref[key] is not None and dst_fused[key] is not None: - assert torch.equal( - dst_ref[key], dst_fused[key] - ), f"Backend {backend_idx} {key} mismatch" + assert torch.equal(dst_ref[key], dst_fused[key]), ( + f"Backend {backend_idx} {key} mismatch" + ) if __name__ == "__main__": diff --git a/test/registered/kernels/ops/attention/test_fused_store_index_cache.py b/test/registered/kernels/ops/attention/test_fused_store_index_cache.py index bcd62a808..a7ecb21cb 100644 --- a/test/registered/kernels/ops/attention/test_fused_store_index_cache.py +++ b/test/registered/kernels/ops/attention/test_fused_store_index_cache.py @@ -392,9 +392,9 @@ def test_roundtrip_reconstruction(num_tokens: int): per_row_energy = reconstructed.abs().sum(dim=-1) orig_energy = original.abs().sum(dim=-1) mask = orig_energy > 0.1 - assert ( - per_row_energy[mask] > 0.01 - ).all(), "Some tokens have zero reconstruction — kernel may not be writing output" + assert (per_row_energy[mask] > 0.01).all(), ( + "Some tokens have zero reconstruction — kernel may not be writing output" + ) # TEST 4: Boundary conditions diff --git a/test/registered/kernels/ops/attention/test_minimax_decode_topk_page_table.py b/test/registered/kernels/ops/attention/test_minimax_decode_topk_page_table.py index a7ecc8ab9..23200ab49 100644 --- a/test/registered/kernels/ops/attention/test_minimax_decode_topk_page_table.py +++ b/test/registered/kernels/ops/attention/test_minimax_decode_topk_page_table.py @@ -159,9 +159,9 @@ def test_dp_flattened_page_table(nkv, bs, seq_len): row = b * nkv + h # effective KV length = sum of valid tokens over selected blocks exp_kv = sum(min(block, seq_len - c * block) for c in blocks) - assert ( - int(cache[row]) == exp_kv - ), f"row {row}: {int(cache[row])} != {exp_kv}" + assert int(cache[row]) == exp_kv, ( + f"row {row}: {int(cache[row])} != {exp_kv}" + ) # page table: each block -> ppb pages via req_to_token, head-minor encoded for e in range(len(blocks) * ppb): c = blocks[e // ppb] @@ -169,9 +169,9 @@ def test_dp_flattened_page_table(nkv, bs, seq_len): if tok >= max_kv: tok = max_kv - 1 exp = int(r2t_cpu[b, tok]) // ps * nkv + h - assert ( - int(pt[row, e]) == exp - ), f"row {row} e {e}: {int(pt[row,e])} != {exp}" + assert int(pt[row, e]) == exp, ( + f"row {row} e {e}: {int(pt[row, e])} != {exp}" + ) if __name__ == "__main__": diff --git a/test/registered/kernels/ops/attention/test_paged_mqa_metadata.py b/test/registered/kernels/ops/attention/test_paged_mqa_metadata.py index d16f09b08..49682ed75 100644 --- a/test/registered/kernels/ops/attention/test_paged_mqa_metadata.py +++ b/test/registered/kernels/ops/attention/test_paged_mqa_metadata.py @@ -58,12 +58,12 @@ def paged_mqa_metadata_ref( seq_lens: torch.Tensor, num_sm: int, page_size: int ) -> torch.Tensor: assert page_size == 64, f"page_size must be 64, got {page_size}" - assert ( - seq_lens.dtype == torch.int32 - ), f"seq_lens dtype must be int32, got {seq_lens.dtype}" - assert ( - seq_lens.dim() == 1 - ), f"seq_lens must be 1-D, got shape {tuple(seq_lens.shape)}" + assert seq_lens.dtype == torch.int32, ( + f"seq_lens dtype must be int32, got {seq_lens.dtype}" + ) + assert seq_lens.dim() == 1, ( + f"seq_lens must be 1-D, got shape {tuple(seq_lens.shape)}" + ) device = seq_lens.device batch_size = int(seq_lens.shape[0]) diff --git a/test/registered/kernels/ops/attention/test_pos_enc.py b/test/registered/kernels/ops/attention/test_pos_enc.py index 03d9b8105..58e3647c6 100644 --- a/test/registered/kernels/ops/attention/test_pos_enc.py +++ b/test/registered/kernels/ops/attention/test_pos_enc.py @@ -481,7 +481,7 @@ def test_performance( print( f"\nPerformance Test - Batch={batch_size}, SeqLen={seq_len}, Tokens={total_tokens}" ) - print(f"JIT: {jit_time*1000:.9f}ms, SGL: {sgl_time*1000:.9f}ms") + print(f"JIT: {jit_time * 1000:.9f}ms, SGL: {sgl_time * 1000:.9f}ms") if sgl_time > 0: speedup = sgl_time / jit_time if jit_time > 0 else float("inf") print(f"Speedup (SGL/JIT): {speedup:.2f}x") diff --git a/test/registered/kernels/ops/attention/test_q8kv8_sparse_prefill_backend.py b/test/registered/kernels/ops/attention/test_q8kv8_sparse_prefill_backend.py index 7347bada1..c74f30e65 100644 --- a/test/registered/kernels/ops/attention/test_q8kv8_sparse_prefill_backend.py +++ b/test/registered/kernels/ops/attention/test_q8kv8_sparse_prefill_backend.py @@ -442,9 +442,10 @@ def test_q8kv8_sparse_prefill_helper_builds_fp8_workspace_matching_bf16_path( bf16_capture = _Capture() q8_capture = _Capture() - with _patched_sparse_kernels( - bf16_capture, q8_capture - ), _patched_compressed_sparse_cache_paths(compress_ratio): + with ( + _patched_sparse_kernels(bf16_capture, q8_capture), + _patched_compressed_sparse_cache_paths(compress_ratio), + ): bf16_out = backend._forward_prefill_sparse( q=q, layer_id=0, diff --git a/test/registered/kernels/ops/attention/test_qprep_bf16_fp8_sm90.py b/test/registered/kernels/ops/attention/test_qprep_bf16_fp8_sm90.py index 05263707d..5a7858be7 100644 --- a/test/registered/kernels/ops/attention/test_qprep_bf16_fp8_sm90.py +++ b/test/registered/kernels/ops/attention/test_qprep_bf16_fp8_sm90.py @@ -109,9 +109,9 @@ def test_qprep_fp64_reference_parity(h_k): ).transpose(0, 1) err_tri = (ref8[..., :N_LORA].to(torch.float64) - ref64).abs().mean().item() err_cuda = (out8[..., :N_LORA].to(torch.float64) - ref64).abs().mean().item() - assert ( - err_cuda <= 1.05 * err_tri - ), f"CUDA fp64-ref mean |err| {err_cuda:.4e} exceeds Triton's {err_tri:.4e}" + assert err_cuda <= 1.05 * err_tri, ( + f"CUDA fp64-ref mean |err| {err_cuda:.4e} exceeds Triton's {err_tri:.4e}" + ) if __name__ == "__main__": diff --git a/test/registered/kernels/ops/attention/test_sm120_paged_mqa_logits.py b/test/registered/kernels/ops/attention/test_sm120_paged_mqa_logits.py index 06ea5211d..23f71e8cb 100644 --- a/test/registered/kernels/ops/attention/test_sm120_paged_mqa_logits.py +++ b/test/registered/kernels/ops/attention/test_sm120_paged_mqa_logits.py @@ -149,9 +149,9 @@ def _compare( # Invalid positions in SM120 output must be -inf positions = torch.arange(max_seq_len, device=sm120.device) invalid = positions.unsqueeze(0) >= seq_lens.unsqueeze(1) - assert torch.all( - torch.isinf(sm120[invalid]) & (sm120[invalid] < 0) - ), "SM120 output must fill invalid positions with -inf" + assert torch.all(torch.isinf(sm120[invalid]) & (sm120[invalid] < 0)), ( + "SM120 output must fill invalid positions with -inf" + ) class TestSM120PagedMqaLogitsTorch(CustomTestCase): diff --git a/test/registered/kernels/ops/attention/test_topk_v2.py b/test/registered/kernels/ops/attention/test_topk_v2.py index b22b3d380..cae08c08c 100644 --- a/test/registered/kernels/ops/attention/test_topk_v2.py +++ b/test/registered/kernels/ops/attention/test_topk_v2.py @@ -96,9 +96,9 @@ def _assert_topk_close(scores_cpu, ref_raw, our_raw, bs, seq_lens, k): print( f"b={i} L={L} k={k}: more={list(more)[:4]} less={list(less)[:4]} mv={mv[:3]} lv={lv[:3]}" ) - assert len(our) == min( - k, L - ), f"b={i} L={L} k={k}: {len(our)} valid != {min(k, L)}" + assert len(our) == min(k, L), ( + f"b={i} L={L} k={k}: {len(our)} valid != {min(k, L)}" + ) assert bad <= MAX_PERMIT_ERROR, f"{bad=} > {MAX_PERMIT_ERROR}" diff --git a/test/registered/kernels/ops/communication/test_amd_deterministic_custom_allreduce.py b/test/registered/kernels/ops/communication/test_amd_deterministic_custom_allreduce.py index 64522f430..bf6ac3498 100644 --- a/test/registered/kernels/ops/communication/test_amd_deterministic_custom_allreduce.py +++ b/test/registered/kernels/ops/communication/test_amd_deterministic_custom_allreduce.py @@ -95,7 +95,7 @@ def worker(world_size, rank, port): input_size_bytes = base_input.numel() * base_input.element_size() if input_size_bytes > custom_ar.max_size and rank == 0: print( - f"Warning: Input size ({input_size_bytes/(1024*1024):.1f} MB) exceeds buffer size ({custom_ar.max_size/(1024*1024):.1f} MB)" + f"Warning: Input size ({input_size_bytes / (1024 * 1024):.1f} MB) exceeds buffer size ({custom_ar.max_size / (1024 * 1024):.1f} MB)" ) print(" Using unregistered mode (will copy to buffer)") @@ -105,9 +105,9 @@ def worker(world_size, rank, port): # TEST 1: Deterministic kernel (same batch size) - should be DETERMINISTIC # ========================================================================= if rank == 0: - print(f"\n{'='*70}") + print(f"\n{'=' * 70}") print("TEST 1: Deterministic kernel (same batch size)") - print(f"{'='*70}") + print(f"{'=' * 70}") dist.barrier() results_allreduce_only = [] @@ -125,7 +125,7 @@ def worker(world_size, rank, port): if rank == 0: print( - f" Trial {trial+1:2d}: sum={checksum:.6f}, first5={first_vals.tolist()}" + f" Trial {trial + 1:2d}: sum={checksum:.6f}, first5={first_vals.tolist()}" ) # Check determinism @@ -135,7 +135,7 @@ def worker(world_size, rank, port): for i, (s, vals) in enumerate(results_allreduce_only[1:], 1): if abs(ref_sum - s) > 1e-3 or not torch.allclose(ref_vals, vals, rtol=1e-3): all_match = False - print(f" Trial {i+1} DIFFERS! ref_sum={ref_sum:.6f}, got={s:.6f}") + print(f" Trial {i + 1} DIFFERS! ref_sum={ref_sum:.6f}, got={s:.6f}") if all_match: print(" ✓ DETERMINISTIC KERNEL (fixed BS): DETERMINISTIC (as expected)") @@ -151,10 +151,10 @@ def worker(world_size, rank, port): # [a], [a, x], [a, x, x], ... # ========================================================================= if rank == 0: - print(f"\n{'='*70}") + print(f"\n{'=' * 70}") print("TEST 2: Deterministic kernel (different batch size)") print("Batches: [a], [a,x], [a,x,x], ...") - print(f"{'='*70}") + print(f"{'=' * 70}") dist.barrier() results_allreduce_only = {trial: [] for trial in range(num_trials)} diff --git a/test/registered/kernels/ops/communication/test_amd_nccl_allreduce_determinism.py b/test/registered/kernels/ops/communication/test_amd_nccl_allreduce_determinism.py index a7985eb5f..5af113d75 100644 --- a/test/registered/kernels/ops/communication/test_amd_nccl_allreduce_determinism.py +++ b/test/registered/kernels/ops/communication/test_amd_nccl_allreduce_determinism.py @@ -63,9 +63,9 @@ def worker(world_size, rank, port): # TEST 1: Default all-reduce (same batch size) - should be DETERMINISTIC # ========================================================================= if rank == 0: - print(f"\n{'='*70}") + print(f"\n{'=' * 70}") print("TEST 1: Default NCCL all_reduce (same batch size)") - print(f"{'='*70}") + print(f"{'=' * 70}") dist.barrier() results_allreduce_only = [] @@ -84,7 +84,7 @@ def worker(world_size, rank, port): if rank == 0: print( - f" Trial {trial+1:2d}: sum={checksum:.6f}, first5={first_vals.tolist()}" + f" Trial {trial + 1:2d}: sum={checksum:.6f}, first5={first_vals.tolist()}" ) # Check determinism @@ -94,7 +94,7 @@ def worker(world_size, rank, port): for i, (s, vals) in enumerate(results_allreduce_only[1:], 1): if abs(ref_sum - s) > 1e-3 or not torch.allclose(ref_vals, vals, rtol=1e-3): all_match = False - print(f" Trial {i+1} DIFFERS! ref_sum={ref_sum:.6f}, got={s:.6f}") + print(f" Trial {i + 1} DIFFERS! ref_sum={ref_sum:.6f}, got={s:.6f}") if all_match: print(" ✓ DEFAULT ALL_REDUCE (fixed BS): DETERMINISTIC (as expected)") @@ -108,10 +108,10 @@ def worker(world_size, rank, port): # [a], [a, x], [a, x, x], ... # ========================================================================= if rank == 0: - print(f"\n{'='*70}") + print(f"\n{'=' * 70}") print("TEST 2: Default NCCL all_reduce (different batch size)") print("Batches: [a], [a,x], [a,x,x], ...") - print(f"{'='*70}") + print(f"{'=' * 70}") dist.barrier() results_allreduce_only = {trial: [] for trial in range(num_trials)} diff --git a/test/registered/kernels/ops/gemm/test_minimax_fused_qkv_index_gemm.py b/test/registered/kernels/ops/gemm/test_minimax_fused_qkv_index_gemm.py index 207f97636..806829c72 100644 --- a/test/registered/kernels/ops/gemm/test_minimax_fused_qkv_index_gemm.py +++ b/test/registered/kernels/ops/gemm/test_minimax_fused_qkv_index_gemm.py @@ -58,9 +58,9 @@ def test_fused_equals_separate(T, N1, N2, K): fused = mxfp8_linear(x, w, sp) assert fused.shape == ref.shape - assert torch.equal( - fused, ref - ), f"max abs diff {(fused.float() - ref.float()).abs().max().item()}" + assert torch.equal(fused, ref), ( + f"max abs diff {(fused.float() - ref.float()).abs().max().item()}" + ) if __name__ == "__main__": diff --git a/test/registered/kernels/ops/kv_canary/test_kernel_config.py b/test/registered/kernels/ops/kv_canary/test_kernel_config.py index e3d7b6d6c..bbaad5a9c 100644 --- a/test/registered/kernels/ops/kv_canary/test_kernel_config.py +++ b/test/registered/kernels/ops/kv_canary/test_kernel_config.py @@ -135,15 +135,15 @@ def test_verify_byte_equal_across_repeated_launches_10x() -> None: snapshot_bufs.append(cuda_buf.clone()) for i in range(1, num_launches): - assert torch.equal( - snapshot_rings[0], snapshot_rings[i] - ), f"violation_ring differs between launch 0 and {i}" - assert torch.equal( - snapshot_write_indices[0], snapshot_write_indices[i] - ), f"violation_write_index differs between launch 0 and {i}" - assert torch.equal( - snapshot_bufs[0], snapshot_bufs[i] - ), f"canary_buf differs between launch 0 and {i}" + assert torch.equal(snapshot_rings[0], snapshot_rings[i]), ( + f"violation_ring differs between launch 0 and {i}" + ) + assert torch.equal(snapshot_write_indices[0], snapshot_write_indices[i]), ( + f"violation_write_index differs between launch 0 and {i}" + ) + assert torch.equal(snapshot_bufs[0], snapshot_bufs[i]), ( + f"canary_buf differs between launch 0 and {i}" + ) def test_write_byte_equal_across_repeated_launches_10x() -> None: @@ -190,15 +190,15 @@ def test_write_byte_equal_across_repeated_launches_10x() -> None: snapshot_counters.append(cuda_log.slot_run_counter.clone()) for i in range(1, num_launches): - assert torch.equal( - snapshot_bufs[0], snapshot_bufs[i] - ), f"canary_buf differs between launch 0 and {i}" - assert torch.equal( - snapshot_rings[0], snapshot_rings[i] - ), f"violation_ring differs between launch 0 and {i}" - assert torch.equal( - snapshot_counters[0], snapshot_counters[i] - ), f"slot_run_counter differs between launch 0 and {i}" + assert torch.equal(snapshot_bufs[0], snapshot_bufs[i]), ( + f"canary_buf differs between launch 0 and {i}" + ) + assert torch.equal(snapshot_rings[0], snapshot_rings[i]), ( + f"violation_ring differs between launch 0 and {i}" + ) + assert torch.equal(snapshot_counters[0], snapshot_counters[i]), ( + f"slot_run_counter differs between launch 0 and {i}" + ) def test_plan_byte_equal_across_repeated_launches_10x() -> None: @@ -254,18 +254,18 @@ def test_plan_byte_equal_across_repeated_launches_10x() -> None: snapshot_write_offsets.append(triton_w.write_offsets.clone()) for i in range(1, num_launches): - assert torch.equal( - snapshot_slots[0], snapshot_slots[i] - ), f"verify_slot_indices differs between launch 0 and {i}" - assert torch.equal( - snapshot_positions[0], snapshot_positions[i] - ), f"verify_expected_positions differs between launch 0 and {i}" - assert torch.equal( - snapshot_prevs[0], snapshot_prevs[i] - ), f"verify_prev_slot_indices differs between launch 0 and {i}" - assert torch.equal( - snapshot_write_offsets[0], snapshot_write_offsets[i] - ), f"write_offsets differs between launch 0 and {i}" + assert torch.equal(snapshot_slots[0], snapshot_slots[i]), ( + f"verify_slot_indices differs between launch 0 and {i}" + ) + assert torch.equal(snapshot_positions[0], snapshot_positions[i]), ( + f"verify_expected_positions differs between launch 0 and {i}" + ) + assert torch.equal(snapshot_prevs[0], snapshot_prevs[i]), ( + f"verify_prev_slot_indices differs between launch 0 and {i}" + ) + assert torch.equal(snapshot_write_offsets[0], snapshot_write_offsets[i]), ( + f"write_offsets differs between launch 0 and {i}" + ) def test_verify_multi_launch_100x_counter_linear() -> None: @@ -301,9 +301,9 @@ def test_verify_multi_launch_100x_counter_linear() -> None: torch.cuda.synchronize() - assert ( - int(cuda_log.kernel_run_counter[0].item()) == num_launches - ), f"kernel_run_counter expected {num_launches}, got {cuda_log.kernel_run_counter[0].item()}" + assert int(cuda_log.kernel_run_counter[0].item()) == num_launches, ( + f"kernel_run_counter expected {num_launches}, got {cuda_log.kernel_run_counter[0].item()}" + ) assert int(cuda_log.slot_run_counter[0].item()) == num_launches, ( f"slot_run_counter expected {num_launches} (1 active entry x 100 launches), " f"got {cuda_log.slot_run_counter[0].item()}" diff --git a/test/registered/kernels/ops/kv_canary/test_pipeline_e2e.py b/test/registered/kernels/ops/kv_canary/test_pipeline_e2e.py index 975495fcd..0178d2afd 100644 --- a/test/registered/kernels/ops/kv_canary/test_pipeline_e2e.py +++ b/test/registered/kernels/ops/kv_canary/test_pipeline_e2e.py @@ -529,9 +529,9 @@ def test_pipeline_pseudo_mode_on_token_mismatch_then_verify_clean() -> None: ) write_violations = int(log_real.write_index[0].item()) - assert ( - write_violations == n_tokens - ), f"expected {n_tokens} write violations, got {write_violations}" + assert write_violations == n_tokens, ( + f"expected {n_tokens} write violations, got {write_violations}" + ) def test_pipeline_empty_batch() -> None: @@ -779,9 +779,9 @@ def test_pipeline_token_mismatch_detected_via_pool() -> None: fail_bits = int( log_real.ring[row_idx, consts.VIOLATION_FIELD_FAIL_REASON_BITS].item() ) - assert fail_bits & int( - consts.FailReason.VERIFY_TOKEN_MISMATCH - ), f"row {row_idx}: VERIFY_TOKEN_MISMATCH bit missing in {fail_bits:#b}" + assert fail_bits & int(consts.FailReason.VERIFY_TOKEN_MISMATCH), ( + f"row {row_idx}: VERIFY_TOKEN_MISMATCH bit missing in {fail_bits:#b}" + ) stored = int(log_real.ring[row_idx, consts.VIOLATION_FIELD_STORED_TOKEN].item()) expected = int( log_real.ring[row_idx, consts.VIOLATION_FIELD_EXPECTED_TOKEN].item() diff --git a/test/registered/kernels/ops/kv_canary/test_plan_hand.py b/test/registered/kernels/ops/kv_canary/test_plan_hand.py index 487f26d80..89ebfafc8 100644 --- a/test/registered/kernels/ops/kv_canary/test_plan_hand.py +++ b/test/registered/kernels/ops/kv_canary/test_plan_hand.py @@ -282,9 +282,9 @@ class TestSeedSlot: write_req_capacity=write_req_capacity, ) actual_seed = int(w_plan.write_seed_slot_indices[0].item()) - assert ( - actual_seed == expected_seed - ), f"[{label}] permuted-LUT seed expected {expected_seed} got {actual_seed}" + assert actual_seed == expected_seed, ( + f"[{label}] permuted-LUT seed expected {expected_seed} got {actual_seed}" + ) def test_swa_window_head_prev_slot_is_real_predecessor(self) -> None: """SWA window with non-zero window_start: head entry's prev_slot != -1; it is the real predecessor.""" @@ -327,12 +327,12 @@ class TestSeedSlot: write_req_capacity=write_req_capacity, ) actual_prev = int(v_plan.verify_prev_slot_indices[0].item()) - assert ( - actual_prev != -1 - ), f"[{label}] SWA window head must have real predecessor, got -1" - assert ( - actual_prev == expected_prev - ), f"[{label}] expected prev={expected_prev} got {actual_prev}" + assert actual_prev != -1, ( + f"[{label}] SWA window head must have real predecessor, got -1" + ) + assert actual_prev == expected_prev, ( + f"[{label}] expected prev={expected_prev} got {actual_prev}" + ) class TestPadding: @@ -404,9 +404,9 @@ class TestPadding: write_req_capacity=write_req_capacity, ) actual_slots = v_plan.verify_slot_indices[:prefix].detach().cpu().tolist() - assert ( - actual_slots == expected_slots - ), f"[{label}] sparse-rtt slots expected {expected_slots} got {actual_slots}" + assert actual_slots == expected_slots, ( + f"[{label}] sparse-rtt slots expected {expected_slots} got {actual_slots}" + ) def test_padding_row_with_garbage_prefix_does_not_oob(self) -> None: """rpi==0 padding row with absurd prefix_lens must not OOB-read req_to_token (row is skipped).""" @@ -435,9 +435,9 @@ class TestPadding: write_req_capacity=write_req_capacity, ) assert int(v_plan.verify_num_valid[0].item()) == 8, label - assert ( - int(w_plan.write_seed_slot_indices[1].item()) == -1 - ), f"[{label}] padding row seed must be -1" + assert int(w_plan.write_seed_slot_indices[1].item()) == -1, ( + f"[{label}] padding row seed must be -1" + ) PlanInvariants.assert_all( verify_plan=v_plan, write_plan=w_plan, @@ -848,9 +848,9 @@ class TestMisc: tail_offsets = ( write_plan.write_offsets[n_active + 1 : 8].detach().cpu().tolist() ) - assert all( - v == 0 for v in tail_offsets - ), f"[{label}] stale write_offsets tail not cleared: {tail_offsets}" + assert all(v == 0 for v in tail_offsets), ( + f"[{label}] stale write_offsets tail not cleared: {tail_offsets}" + ) class TestVerifyContent: @@ -983,16 +983,16 @@ class TestByteEqual: ] for i, value in enumerate(expected_write_offsets): - assert ( - int(triton_w.write_offsets[i].item()) == value - ), f"write_offsets[{i}] expected {value} got {int(triton_w.write_offsets[i].item())}" - assert ( - int(triton_v.verify_num_valid[0].item()) == expected_verify_num_valid - ), f"verify_num_valid expected {expected_verify_num_valid}" + assert int(triton_w.write_offsets[i].item()) == value, ( + f"write_offsets[{i}] expected {value} got {int(triton_w.write_offsets[i].item())}" + ) + assert int(triton_v.verify_num_valid[0].item()) == expected_verify_num_valid, ( + f"verify_num_valid expected {expected_verify_num_valid}" + ) for i, expected_seed in enumerate(expected_seeds): - assert ( - int(triton_w.write_seed_slot_indices[i].item()) == expected_seed - ), f"write_seed_slot_indices[{i}] expected {expected_seed}" + assert int(triton_w.write_seed_slot_indices[i].item()) == expected_seed, ( + f"write_seed_slot_indices[{i}] expected {expected_seed}" + ) class TestBoundarySweep: diff --git a/test/registered/kernels/ops/kv_canary/test_verify_hand.py b/test/registered/kernels/ops/kv_canary/test_verify_hand.py index 266b7ae61..a1bea876a 100644 --- a/test/registered/kernels/ops/kv_canary/test_verify_hand.py +++ b/test/registered/kernels/ops/kv_canary/test_verify_hand.py @@ -380,9 +380,9 @@ class TestChain: buf_pair=buf_pair, plan_pair=plan_pair, assert_equal=False ) - assert ( - _n_violations(cuda_log) == 0 - ), f"unexpected violation at iteration token={token} position={position} slot={slot_idx}" + assert _n_violations(cuda_log) == 0, ( + f"unexpected violation at iteration token={token} position={position} slot={slot_idx}" + ) def test_prev_slot_padding_skips_chain_check_arbitrary_stored_hash(self) -> None: """prev_slot_idx == TOKEN_TO_KV_SLOT_PADDING → chain check is skipped, regardless of stored chain hash.""" @@ -692,9 +692,9 @@ class TestViolationField: f"(bit_to_trigger={bit_to_trigger} injection_position={injection_position})" ) else: - assert ( - _n_violations(cuda_log) > ring_capacity - ), "write_index did not advance beyond ring_capacity after overflow" + assert _n_violations(cuda_log) > ring_capacity, ( + "write_index did not advance beyond ring_capacity after overflow" + ) def test_position_mismatch_sets_position_bit_only(self) -> None: """Plan.position != stored.position with chain hash correct → only POSITION bit set.""" @@ -706,12 +706,12 @@ class TestViolationField: cuda_log, _ = run_verify_diff(buf_pair=buf_pair, plan_pair=plan_pair) assert _n_violations(cuda_log) == 1 bits = _fail_bits(cuda_log) - assert ( - bits & consts.FailReason.VERIFY_POSITION_MISMATCH - ), f"expected POSITION bit, got {bits:#b}" - assert ( - bits & consts.FailReason.VERIFY_CHAIN_HASH_MISMATCH - ) == 0, f"chain hash bit unexpectedly set: {bits:#b}" + assert bits & consts.FailReason.VERIFY_POSITION_MISMATCH, ( + f"expected POSITION bit, got {bits:#b}" + ) + assert (bits & consts.FailReason.VERIFY_CHAIN_HASH_MISMATCH) == 0, ( + f"chain hash bit unexpectedly set: {bits:#b}" + ) class TestRealKvHash: @@ -990,9 +990,9 @@ class TestRealKvHash: assert _n_violations(cuda_log) >= 1 bits = _fail_bits(cuda_log) - assert ( - bits & consts.FailReason.VERIFY_REAL_KV_HASH_MISMATCH - ), f"expected REAL_KV_HASH bit, got {bits:#b}" + assert bits & consts.FailReason.VERIFY_REAL_KV_HASH_MISMATCH, ( + f"expected REAL_KV_HASH bit, got {bits:#b}" + ) def test_real_kv_off_does_not_deref_real_kv_sources(self) -> None: buf_pair = _buf_pair(num_slots=8) @@ -1738,9 +1738,9 @@ class TestViolationRing: plan_slot_set = set(slot_indices) for row in range(n_violations): kind = int(cuda_log.ring[row, consts.VIOLATION_FIELD_KERNEL_KIND].item()) - assert kind == int( - launch_tag - ), f"row {row} kind {kind} != {int(launch_tag)}" + assert kind == int(launch_tag), ( + f"row {row} kind {kind} != {int(launch_tag)}" + ) slot = int(cuda_log.ring[row, 1].item()) assert slot in plan_slot_set, f"row {row} slot {slot} not in plan" diff --git a/test/registered/kernels/ops/kv_canary/test_write_fuzz.py b/test/registered/kernels/ops/kv_canary/test_write_fuzz.py index cd78d6a35..33004395a 100644 --- a/test/registered/kernels/ops/kv_canary/test_write_fuzz.py +++ b/test/registered/kernels/ops/kv_canary/test_write_fuzz.py @@ -204,9 +204,9 @@ def _run_one(inputs: WriteFuzzInputs) -> None: kernel_kind=inputs.kernel_kind, assert_equal=False, ) - assert torch.equal( - inputs.cuda_canary_buf, inputs.ref_canary_buf - ), "CUDA vs ref canary_buf diverged" + assert torch.equal(inputs.cuda_canary_buf, inputs.ref_canary_buf), ( + "CUDA vs ref canary_buf diverged" + ) assert int(cuda_log.write_index[0].item()) == int(ref_log.write_index[0].item()) assert int(cuda_log.slot_run_counter[0].item()) == int( ref_log.slot_run_counter[0].item() diff --git a/test/registered/kernels/ops/kv_canary/test_write_hand.py b/test/registered/kernels/ops/kv_canary/test_write_hand.py index 065633aa1..79ba9f260 100644 --- a/test/registered/kernels/ops/kv_canary/test_write_hand.py +++ b/test/registered/kernels/ops/kv_canary/test_write_hand.py @@ -362,9 +362,9 @@ class TestSeedSlot: new_stored = read_slot_fields(canary_buf=self.buf_pair[0], slot_idx=new_slot) assert new_stored[0] == new_token assert new_stored[1] == new_position - assert new_stored[2] == to_signed_int64( - expected_running - ), f"new slot prev_hash {new_stored[2]} != expected {to_signed_int64(expected_running)}" + assert new_stored[2] == to_signed_int64(expected_running), ( + f"new slot prev_hash {new_stored[2]} != expected {to_signed_int64(expected_running)}" + ) class TestChain: @@ -446,9 +446,9 @@ class TestChain: stored_prev_signed, stored_real_kv_hash = read_slot_fields( canary_buf=cuda_buf, slot_idx=slot_idx )[2:] - assert stored_prev_signed == to_signed_int64( - running - ), f"slot {slot_idx}: stored prev_hash != recomputed chain step" + assert stored_prev_signed == to_signed_int64(running), ( + f"slot {slot_idx}: stored prev_hash != recomputed chain step" + ) running = splitmix64_mix3(running, token, position) @@ -711,9 +711,9 @@ class TestSlotHandling: for slot in range(cuda_buf.shape[0]): if slot in (5, 7): continue - assert torch.equal( - after[slot], cuda_buf_before_slot_view[slot] - ), f"slot {slot} should not have been written" + assert torch.equal(after[slot], cuda_buf_before_slot_view[slot]), ( + f"slot {slot} should not have been written" + ) def test_shrink_active_reqs_does_not_write_stale_slots(self) -> None: """Run write with bs=3 plan after a bs=8 run on same buffer: stale slots from bs=8 stay intact.""" @@ -749,9 +749,9 @@ class TestSlotHandling: after = cuda_buf.view(torch.int64) for slot in big_slots: - assert torch.equal( - after[slot], untouched_snapshot[slot] - ), f"slot {slot} from earlier bs=8 run was overwritten by bs=3 run" + assert torch.equal(after[slot], untouched_snapshot[slot]), ( + f"slot {slot} from earlier bs=8 run was overwritten by bs=3 run" + ) class TestRealKvHash: @@ -927,9 +927,9 @@ class TestRealKvHash: _, _, _, stored_real_kv_hash = read_slot_fields( canary_buf=self.buf_pair[0], slot_idx=0 ) - assert stored_real_kv_hash == to_signed_int64( - expected_hash - ), f"stored_real_kv_hash={stored_real_kv_hash:#x} expected={to_signed_int64(expected_hash):#x}" + assert stored_real_kv_hash == to_signed_int64(expected_hash), ( + f"stored_real_kv_hash={stored_real_kv_hash:#x} expected={to_signed_int64(expected_hash):#x}" + ) def test_paged_real_kv_hash_consistent_across_slots(self) -> None: """page=16: writing two slots inside same page yields independent real_kv_hash per slot.""" @@ -1008,9 +1008,9 @@ class TestRealKvHash: fields_b = _run_with(sources_b) assert fields_a[3] != 0 assert fields_b[3] != 0 - assert ( - fields_a[3] != fields_b[3] - ), "reversing source order must change real_kv_hash (fold is ordered)" + assert fields_a[3] != fields_b[3], ( + "reversing source order must change real_kv_hash (fold is ordered)" + ) class TestRunCounter: @@ -1201,9 +1201,9 @@ class TestPseudoMode: ) assert int(cuda_log.write_index[0].item()) >= 1 bits = int(cuda_log.ring[0, consts.VIOLATION_FIELD_FAIL_REASON_BITS].item()) - assert ( - bits & consts.FailReason.WRITE_TOKEN_MISMATCH - ), f"expected WRITE_TOKEN_MISMATCH bit, got {bits:#b}" + assert bits & consts.FailReason.WRITE_TOKEN_MISMATCH, ( + f"expected WRITE_TOKEN_MISMATCH bit, got {bits:#b}" + ) def test_pseudo_mode_off_skips_token_check(self) -> None: """enable_write_verify_inputs=False makes the caller pass no expected-input tensors.""" diff --git a/test/registered/kernels/ops/kvcache/test_hicache.py b/test/registered/kernels/ops/kvcache/test_hicache.py index ec69f336b..358f6d3b3 100644 --- a/test/registered/kernels/ops/kvcache/test_hicache.py +++ b/test/registered/kernels/ops/kvcache/test_hicache.py @@ -97,9 +97,9 @@ def _run_transfer_roundtrip_mha(layout: str, element_dim: int) -> None: device_pool=device_pool, layout=layout, ) - assert ( - host_pool.can_use_jit - ), f"Expected JIT HiCache kernel for MHA dim={element_dim}" + assert host_pool.can_use_jit, ( + f"Expected JIT HiCache kernel for MHA dim={element_dim}" + ) for layer_id in range(NUM_LAYERS): _copy_tensor_with_offset(device_pool.k_buffer[layer_id], layer_id) @@ -191,9 +191,9 @@ def _run_transfer_roundtrip_mla(layout: str, element_dim: int) -> None: device_pool=device_pool, layout=layout, ) - assert ( - host_pool.can_use_jit - ), f"Expected JIT HiCache kernel for MLA dim={element_dim}" + assert host_pool.can_use_jit, ( + f"Expected JIT HiCache kernel for MLA dim={element_dim}" + ) for layer_id in range(NUM_LAYERS): _copy_tensor_with_offset(device_pool.kv_buffer[layer_id], layer_id) diff --git a/test/registered/kernels/ops/layernorm/test_gemma4_fused_routing.py b/test/registered/kernels/ops/layernorm/test_gemma4_fused_routing.py index bb94c4784..bdb6a2189 100644 --- a/test/registered/kernels/ops/layernorm/test_gemma4_fused_routing.py +++ b/test/registered/kernels/ops/layernorm/test_gemma4_fused_routing.py @@ -68,9 +68,9 @@ def test_matches_reference(fused_routing, dtype, T, E, K): # Tie-break order may differ; require the same top-K set and weight sum. ref_set = ref_i.sort(dim=-1).values out_set = out_i.sort(dim=-1).values - assert torch.equal( - out_set, ref_set - ), "fused routing picked a different top-K set than reference" + assert torch.equal(out_set, ref_set), ( + "fused routing picked a different top-K set than reference" + ) torch.testing.assert_close( out_w.sum(dim=-1).to(torch.float32), ref_w.sum(dim=-1).to(torch.float32), diff --git a/test/registered/kernels/ops/layernorm/test_rmsnorm_hf.py b/test/registered/kernels/ops/layernorm/test_rmsnorm_hf.py index 2a380f363..c676de355 100644 --- a/test/registered/kernels/ops/layernorm/test_rmsnorm_hf.py +++ b/test/registered/kernels/ops/layernorm/test_rmsnorm_hf.py @@ -98,9 +98,9 @@ def test_rmsnorm_hf_matches_hf_not_sgl(dtype: torch.dtype) -> None: assert (sgl_ref - hf_ref).abs().max() > 0, "inputs don't exercise the difference" diff_hf = (out - hf_ref).abs().max().item() diff_sgl = (out - sgl_ref).abs().max().item() - assert ( - diff_hf < diff_sgl - ), f"kernel closer to SGL than HF (hf={diff_hf}, sgl={diff_sgl})" + assert diff_hf < diff_sgl, ( + f"kernel closer to SGL than HF (hf={diff_hf}, sgl={diff_sgl})" + ) def test_rmsnorm_hf_empty_input() -> None: diff --git a/test/registered/kernels/ops/moe/test_minimax_quant_scatter.py b/test/registered/kernels/ops/moe/test_minimax_quant_scatter.py index 7007fd989..6a7c709c1 100644 --- a/test/registered/kernels/ops/moe/test_minimax_quant_scatter.py +++ b/test/registered/kernels/ops/moe/test_minimax_quant_scatter.py @@ -94,9 +94,9 @@ def test_quant_scatter_matches_quant_plus_fill(num_tokens, topk, hidden, group): assert torch.equal( gi_new[e, m].view(torch.uint8), gi_ref[e, m].view(torch.uint8) ), f"fp8 mismatch token={t} slot={j} expert={e}" - assert torch.equal( - gs_new[e, :, m], gs_ref[e, :, m] - ), f"scale mismatch token={t} slot={j} expert={e}" + assert torch.equal(gs_new[e, :, m], gs_ref[e, :, m]), ( + f"scale mismatch token={t} slot={j} expert={e}" + ) def test_standard_deepgemm_preprocess_quantizes_with_ue8m0_scale(): diff --git a/test/registered/kernels/ops/moe/test_moe_lora_align_block_size.py b/test/registered/kernels/ops/moe/test_moe_lora_align_block_size.py index b0962d36d..35cc31427 100644 --- a/test/registered/kernels/ops/moe/test_moe_lora_align_block_size.py +++ b/test/registered/kernels/ops/moe/test_moe_lora_align_block_size.py @@ -159,9 +159,9 @@ def test_moe_lora_align_block_size( # Check that all tokens in this block truly belong to 'lora_idx' actual_owners = token_ownership[original_token_indices] - assert torch.all( - actual_owners == lora_idx - ), f"Kernel put tokens from LoRA {actual_owners} into block for LoRA {lora_idx}" + assert torch.all(actual_owners == lora_idx), ( + f"Kernel put tokens from LoRA {actual_owners} into block for LoRA {lora_idx}" + ) if __name__ == "__main__": diff --git a/test/registered/kernels/ops/moe/test_moe_topk_sigmoid.py b/test/registered/kernels/ops/moe/test_moe_topk_sigmoid.py index 785ef7ec3..247e7d8f0 100644 --- a/test/registered/kernels/ops/moe/test_moe_topk_sigmoid.py +++ b/test/registered/kernels/ops/moe/test_moe_topk_sigmoid.py @@ -319,12 +319,14 @@ def test_topk_sigmoid_vs_ref(num_tokens, num_experts, topk, dtype, renormalize): ref_w.sort(dim=-1)[0], atol=1e-3, rtol=1e-3, - ), f"Weight mismatch (dtype={dtype}, n_exp={num_experts}, topk={topk}, renorm={renormalize})" + ), ( + f"Weight mismatch (dtype={dtype}, n_exp={num_experts}, topk={topk}, renorm={renormalize})" + ) # Exact index match is only reliable for float32 (fp16/bf16 tie-breaking may differ) if dtype == torch.float32: - assert torch.equal( - topk_i, ref_i - ), f"Index mismatch (dtype={dtype}, n_exp={num_experts}, topk={topk})" + assert torch.equal(topk_i, ref_i), ( + f"Index mismatch (dtype={dtype}, n_exp={num_experts}, topk={topk})" + ) # --------------------------------------------------------------------------- @@ -351,12 +353,12 @@ def test_topk_sigmoid_with_correction_bias(num_tokens, num_experts, topk, renorm ref_w, ref_i = topk_sigmoid_ref(gating, topk, renormalize, correction_bias=bias) - assert torch.allclose( - topk_w, ref_w, atol=1e-3, rtol=1e-3 - ), f"Weight mismatch with bias (n_exp={num_experts}, topk={topk}, renorm={renormalize})" - assert torch.equal( - topk_i, ref_i - ), f"Index mismatch with bias (n_exp={num_experts}, topk={topk})" + assert torch.allclose(topk_w, ref_w, atol=1e-3, rtol=1e-3), ( + f"Weight mismatch with bias (n_exp={num_experts}, topk={topk}, renorm={renormalize})" + ) + assert torch.equal(topk_i, ref_i), ( + f"Index mismatch with bias (n_exp={num_experts}, topk={topk})" + ) # --------------------------------------------------------------------------- @@ -394,12 +396,12 @@ def test_topk_sigmoid_with_fused_shared_experts( gating, topk + 1, renormalize, correction_bias=bias, num_fused_shared_experts=1 ) - assert torch.allclose( - topk_w, ref_w, atol=1e-3, rtol=1e-3 - ), f"Weight mismatch with bias (n_exp={num_experts}, topk={topk}, renorm={renormalize})" - assert torch.equal( - topk_i, ref_i - ), f"Index mismatch with bias (n_exp={num_experts}, topk={topk})" + assert torch.allclose(topk_w, ref_w, atol=1e-3, rtol=1e-3), ( + f"Weight mismatch with bias (n_exp={num_experts}, topk={topk}, renorm={renormalize})" + ) + assert torch.equal(topk_i, ref_i), ( + f"Index mismatch with bias (n_exp={num_experts}, topk={topk})" + ) # --------------------------------------------------------------------------- @@ -484,12 +486,12 @@ def test_topk_sigmoid_vs_aot(num_tokens, num_experts, topk, dtype, renormalize): topk_i_aot = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda") topk_sigmoid_aot(topk_w_aot, topk_i_aot, gating, renormalize=renormalize) - assert torch.allclose( - topk_w_jit, topk_w_aot, atol=1e-3, rtol=1e-3 - ), f"JIT vs AOT weight mismatch (dtype={dtype}, n_exp={num_experts}, topk={topk})" - assert torch.equal( - topk_i_jit, topk_i_aot - ), f"JIT vs AOT index mismatch (dtype={dtype}, n_exp={num_experts}, topk={topk})" + assert torch.allclose(topk_w_jit, topk_w_aot, atol=1e-3, rtol=1e-3), ( + f"JIT vs AOT weight mismatch (dtype={dtype}, n_exp={num_experts}, topk={topk})" + ) + assert torch.equal(topk_i_jit, topk_i_aot), ( + f"JIT vs AOT index mismatch (dtype={dtype}, n_exp={num_experts}, topk={topk})" + ) if __name__ == "__main__": diff --git a/test/registered/kernels/ops/moe/test_shuffle_rows_with_scales.py b/test/registered/kernels/ops/moe/test_shuffle_rows_with_scales.py index 0b6d46a6d..767a32c00 100644 --- a/test/registered/kernels/ops/moe/test_shuffle_rows_with_scales.py +++ b/test/registered/kernels/ops/moe/test_shuffle_rows_with_scales.py @@ -67,9 +67,9 @@ def _inputs(k, num_src_rows, num_dst_rows, seed): def _assert_same_bytes(got, ref, what): - assert torch.equal( - got.view(torch.int8), ref.view(torch.int8) - ), f"{what} bytes differ" + assert torch.equal(got.view(torch.int8), ref.view(torch.int8)), ( + f"{what} bytes differ" + ) @pytest.mark.parametrize("k,num_src_rows,num_dst_rows", CASES) diff --git a/test/registered/kernels/ops/quantization/test_per_token_group_quant_8bit_v2.py b/test/registered/kernels/ops/quantization/test_per_token_group_quant_8bit_v2.py index be19861fc..9a7d749b3 100644 --- a/test/registered/kernels/ops/quantization/test_per_token_group_quant_8bit_v2.py +++ b/test/registered/kernels/ops/quantization/test_per_token_group_quant_8bit_v2.py @@ -179,9 +179,9 @@ def test_v2_jit_masked_matches_aot(num_experts, hidden, tokens_pad): ) torch.cuda.synchronize() - assert torch.equal( - x_q.view(torch.int8), q_ref.view(torch.int8) - ), "masked fp8 differ" + assert torch.equal(x_q.view(torch.int8), q_ref.view(torch.int8)), ( + "masked fp8 differ" + ) assert torch.equal(x_s, s_ref), "masked scales differ" diff --git a/test/registered/kernels/test_jit_cache.py b/test/registered/kernels/test_jit_cache.py index 7fb2e7fe0..e547db11f 100644 --- a/test/registered/kernels/test_jit_cache.py +++ b/test/registered/kernels/test_jit_cache.py @@ -187,9 +187,9 @@ def test_no_unordered_container_reaches_the_key(monkeypatch): _build_key() def walk(value, path="parts"): - assert not isinstance( - value, (set, frozenset, dict) - ), f"unordered container at {path}: {type(value).__name__}" + assert not isinstance(value, (set, frozenset, dict)), ( + f"unordered container at {path}: {type(value).__name__}" + ) if isinstance(value, (list, tuple)): for index, item in enumerate(value): walk(item, f"{path}[{index}]") diff --git a/test/registered/kernels/test_kda_replayssm_ring_ragged.py b/test/registered/kernels/test_kda_replayssm_ring_ragged.py index b428b263e..daef77a70 100644 --- a/test/registered/kernels/test_kda_replayssm_ring_ragged.py +++ b/test/registered/kernels/test_kda_replayssm_ring_ragged.py @@ -114,9 +114,9 @@ def _run_case(bs, gamma, HV, H, K, V, lens, acc, L, pad_last=False): base = inter[slots[j], int(acc[j]) - 1] fold = ckpt[slots[j]] rel = ((fold - base).abs().max() / base.abs().max().clamp_min(1e-6)).item() - assert ( - rel < 1e-3 - ), f"row={j} len={int(lens[j])} acc={int(acc[j])}: rel={rel:.3e}" + assert rel < 1e-3, ( + f"row={j} len={int(lens[j])} acc={int(acc[j])}: rel={rel:.3e}" + ) @pytest.mark.parametrize("bs,gamma,HV,H,K,V", SHAPES, ids=SHAPE_IDS) diff --git a/test/registered/kernels/test_kernel_inventory.py b/test/registered/kernels/test_kernel_inventory.py index 99069e51f..481b1209a 100644 --- a/test/registered/kernels/test_kernel_inventory.py +++ b/test/registered/kernels/test_kernel_inventory.py @@ -128,9 +128,9 @@ def test_registered_kernel_test_groups_are_known(): registered_root = REPO_ROOT / "test" / "registered" / "kernels" for kind in ("ops", "benchmark"): unknown = _directory_names(registered_root / kind) - declared_groups - assert ( - not unknown - ), f"Unknown {kind} kernel group directories: {sorted(unknown)}" + assert not unknown, ( + f"Unknown {kind} kernel group directories: {sorted(unknown)}" + ) def test_internal_registry_target_attributes_are_declared(): @@ -172,7 +172,9 @@ def test_jit_source_declarations_exist(): function_name = ( call.func.id if isinstance(call.func, ast.Name) - else call.func.attr if isinstance(call.func, ast.Attribute) else None + else call.func.attr + if isinstance(call.func, ast.Attribute) + else None ) if function_name != "load_jit": continue diff --git a/test/registered/kernels/test_lean_attention.py b/test/registered/kernels/test_lean_attention.py index 73e5124d4..0f444c0b9 100644 --- a/test/registered/kernels/test_lean_attention.py +++ b/test/registered/kernels/test_lean_attention.py @@ -209,9 +209,9 @@ def _run_pair_paged( kv_group_num = H_Q // H_KV sm = 1.0 / (D**0.5) tot = B * S - assert ( - tot % page_size == 0 - ), "test setup: total tokens must be a multiple of page_size" + assert tot % page_size == 0, ( + "test setup: total tokens must be a multiple of page_size" + ) num_pages = tot // page_size # 4-D paged KV buffers [num_pages, page_size, head, dim] (the shared-pool layout). diff --git a/test/registered/kernels/test_quick_allreduce_bf16_range.py b/test/registered/kernels/test_quick_allreduce_bf16_range.py index 347a40693..1e58bd016 100644 --- a/test/registered/kernels/test_quick_allreduce_bf16_range.py +++ b/test/registered/kernels/test_quick_allreduce_bf16_range.py @@ -63,9 +63,9 @@ def _run_bf16_range_test(rank: int, world_size: int, port: int) -> None: dist.barrier() out = quick_all_reduce.quick_all_reduce(inp) torch.cuda.synchronize() - assert ( - torch.isfinite(out).all().item() - ), f"{quant_mode=} {case_name=} produced non-finite output" + assert torch.isfinite(out).all().item(), ( + f"{quant_mode=} {case_name=} produced non-finite output" + ) if quant_mode == "FP" or case_name in ("low", "ordinary"): torch.testing.assert_close( out, diff --git a/test/registered/kv_canary/test_self_e2e_pp_baseline.py b/test/registered/kv_canary/test_self_e2e_pp_baseline.py index 1e004befd..2656bc7dc 100644 --- a/test/registered/kv_canary/test_self_e2e_pp_baseline.py +++ b/test/registered/kv_canary/test_self_e2e_pp_baseline.py @@ -11,7 +11,6 @@ register_amd_ci(est_time=243, stage="extra-a", runner_config="2-gpu-large-amd") class TestPPBaselineSwa(CanaryPPFixture): - kv_canary_mode = CanaryMode.LOG def test_no_violation(self) -> None: diff --git a/test/registered/kv_canary/test_self_e2e_pp_perturb.py b/test/registered/kv_canary/test_self_e2e_pp_perturb.py index a97881631..8e4c32f1a 100644 --- a/test/registered/kv_canary/test_self_e2e_pp_perturb.py +++ b/test/registered/kv_canary/test_self_e2e_pp_perturb.py @@ -13,7 +13,6 @@ register_amd_ci(est_time=298, stage="extra-a", runner_config="2-gpu-large-amd") class TestPPPerturbSwaSwa(CanaryPPFixture): - kv_canary_mode = CanaryMode.LOG target_group: ClassVar[TargetGroupKind] = TargetGroupKind.SWA extra_server_args = ("--kv-canary-real-data", "partial") diff --git a/test/registered/kv_canary/test_self_unit_endpoint.py b/test/registered/kv_canary/test_self_unit_endpoint.py index ad94a8879..3143226d6 100644 --- a/test/registered/kv_canary/test_self_unit_endpoint.py +++ b/test/registered/kv_canary/test_self_unit_endpoint.py @@ -70,14 +70,17 @@ class TestSelfUnitEndpoint(CustomTestCase): def test_launch_sweep_only_calls_verify(self): """Verify sweep launch invokes only the verify kernel.""" calls: list[str] = [] - with patch.object( - endpoint_module, - "launch_canary_verify_kernel", - lambda **kwargs: calls.append("verify"), - ), patch.object( - endpoint_module, - "launch_canary_write_kernel", - lambda **kwargs: calls.append("write"), + with ( + patch.object( + endpoint_module, + "launch_canary_verify_kernel", + lambda **kwargs: calls.append("verify"), + ), + patch.object( + endpoint_module, + "launch_canary_write_kernel", + lambda **kwargs: calls.append("write"), + ), ): ep = _make_endpoint( device=self.device, kernel_kind=CanaryLaunchTag.SWEEP_K_FULL @@ -93,14 +96,21 @@ class TestSelfUnitEndpoint(CustomTestCase): def test_launch_per_forward_passes_kernel_kind(self): """Verify per-forward launch passes the endpoint kernel kind.""" captured: list[tuple[str, CanaryLaunchTag]] = [] - with patch.object( - endpoint_module, - "launch_canary_verify_kernel", - lambda **kwargs: captured.append(("verify", kwargs["context"].kernel_kind)), - ), patch.object( - endpoint_module, - "launch_canary_write_kernel", - lambda **kwargs: captured.append(("write", kwargs["context"].kernel_kind)), + with ( + patch.object( + endpoint_module, + "launch_canary_verify_kernel", + lambda **kwargs: captured.append( + ("verify", kwargs["context"].kernel_kind) + ), + ), + patch.object( + endpoint_module, + "launch_canary_write_kernel", + lambda **kwargs: captured.append( + ("write", kwargs["context"].kernel_kind) + ), + ), ): ep = _make_endpoint( device=self.device, kernel_kind=CanaryLaunchTag.TAIL_V_SWA @@ -124,16 +134,19 @@ class TestSelfUnitEndpoint(CustomTestCase): def test_endpoint_shares_violation_log_across_launches(self): """Verify endpoints can reuse the same violation log.""" captured_rings: list[int] = [] - with patch.object( - endpoint_module, - "launch_canary_verify_kernel", - lambda **kwargs: captured_rings.append( - kwargs["context"].violation_ring.data_ptr() + with ( + patch.object( + endpoint_module, + "launch_canary_verify_kernel", + lambda **kwargs: captured_rings.append( + kwargs["context"].violation_ring.data_ptr() + ), + ), + patch.object( + endpoint_module, + "launch_canary_write_kernel", + lambda **kwargs: None, ), - ), patch.object( - endpoint_module, - "launch_canary_write_kernel", - lambda **kwargs: None, ): shared_log = ViolationLog.allocate(ring_capacity=2, device=self.device) ep_a = _make_endpoint( @@ -160,12 +173,15 @@ class TestSelfUnitEndpoint(CustomTestCase): def test_swa_endpoint_pre_translates_out_cache_loc(self): """Verify SWA endpoints translate cache locations before write launch.""" captured: list[torch.Tensor] = [] - with patch.object( - endpoint_module, "launch_canary_verify_kernel", lambda **kwargs: None - ), patch.object( - endpoint_module, - "launch_canary_write_kernel", - lambda **kwargs: captured.append(kwargs["out_cache_loc"]), + with ( + patch.object( + endpoint_module, "launch_canary_verify_kernel", lambda **kwargs: None + ), + patch.object( + endpoint_module, + "launch_canary_write_kernel", + lambda **kwargs: captured.append(kwargs["out_cache_loc"]), + ), ): # LUT maps full slot i → swa slot (i + 100) so we can verify the gather happened. lut = torch.arange(8, dtype=torch.int64, device=self.device) + 100 @@ -214,12 +230,15 @@ class TestSelfUnitEndpoint(CustomTestCase): def test_swa_endpoint_trailing_sentinel_row_yields_skip(self): """Verify SWA sentinel cache rows become write-skip markers.""" captured: list[torch.Tensor] = [] - with patch.object( - endpoint_module, "launch_canary_verify_kernel", lambda **kwargs: None - ), patch.object( - endpoint_module, - "launch_canary_write_kernel", - lambda **kwargs: captured.append(kwargs["out_cache_loc"]), + with ( + patch.object( + endpoint_module, "launch_canary_verify_kernel", lambda **kwargs: None + ), + patch.object( + endpoint_module, + "launch_canary_write_kernel", + lambda **kwargs: captured.append(kwargs["out_cache_loc"]), + ), ): # 8 in-window rows + 1 trailing sentinel row at index 8. lut = torch.arange(8, dtype=torch.int64, device=self.device) diff --git a/test/registered/kv_canary/test_self_unit_perturb.py b/test/registered/kv_canary/test_self_unit_perturb.py index 402dab919..99b4ce19d 100644 --- a/test/registered/kv_canary/test_self_unit_perturb.py +++ b/test/registered/kv_canary/test_self_unit_perturb.py @@ -159,22 +159,27 @@ class TestPerturbManager(CustomTestCase): forward_batch = make_forward_batch(device, bs=1, seq_lens_list=(1,)) calls: list[str] = [] - with patch.object( - manager, - "perturb_real_kv_post_forward", - lambda batch: calls.append("real_kv_post_forward"), - ), patch.object( - manager, - "perturb_req_to_token", - lambda batch: calls.append("req_to_token"), - ), patch.object( - manager, - "perturb_real_kv_used", - lambda batch: calls.append("real_kv_used"), - ), patch.object( - manager, - "perturb_real_kv_unused_cache", - lambda batch: calls.append("real_kv_unused_cache"), + with ( + patch.object( + manager, + "perturb_real_kv_post_forward", + lambda batch: calls.append("real_kv_post_forward"), + ), + patch.object( + manager, + "perturb_req_to_token", + lambda batch: calls.append("req_to_token"), + ), + patch.object( + manager, + "perturb_real_kv_used", + lambda batch: calls.append("real_kv_used"), + ), + patch.object( + manager, + "perturb_real_kv_unused_cache", + lambda batch: calls.append("real_kv_unused_cache"), + ), ): manager.perturb_post_forward(maybe_inaccurate_forward_batch=forward_batch) @@ -394,10 +399,13 @@ class TestRealKvUsedPerturb(CustomTestCase): pool_snapshot = pool.req_to_token.clone() source_snapshot = source.tensor.clone() - with patch.object(torch, "rand", return_value=torch.tensor(0.0)), patch.object( - real_kv_unused_cache_module, - "_pick_sweep_slot_for_group", - return_value=3, + with ( + patch.object(torch, "rand", return_value=torch.tensor(0.0)), + patch.object( + real_kv_unused_cache_module, + "_pick_sweep_slot_for_group", + return_value=3, + ), ): manager.perturb(maybe_inaccurate_forward_batch=forward_batch) @@ -436,10 +444,13 @@ class TestRealKvUnusedCachePerturb(CustomTestCase): manager.attach_radix_cache(make_radix_cache([[], [3]], device=device)) snapshot = source.tensor.clone() - with patch.object(torch, "rand", return_value=torch.tensor(0.0)), patch.object( - torch, - "randint", - return_value=torch.tensor(0), + with ( + patch.object(torch, "rand", return_value=torch.tensor(0.0)), + patch.object( + torch, + "randint", + return_value=torch.tensor(0), + ), ): manager.perturb_real_kv_unused_cache(None) diff --git a/test/registered/kv_canary/test_self_unit_runner_per_forward.py b/test/registered/kv_canary/test_self_unit_runner_per_forward.py index fa72180d9..a6e47bbd2 100644 --- a/test/registered/kv_canary/test_self_unit_runner_per_forward.py +++ b/test/registered/kv_canary/test_self_unit_runner_per_forward.py @@ -28,21 +28,25 @@ class TestManagerPerForward(CanaryManagerTestCase): def test_per_forward_orchestrates_plan_head_tail(self) -> None: """Verify per-forward execution launches plan, head/tail verify kernels, and write kernels in order.""" calls: list[object] = [] - with patch.object( - kernel_launcher_module, - "launch_canary_plan_kernels", - lambda **kwargs: calls.append("plan"), - ), patch.object( - endpoint_module, - "launch_canary_verify_kernel", - lambda **kwargs: calls.append( - ("verify", kwargs["context"].kernel_kind.name) + with ( + patch.object( + kernel_launcher_module, + "launch_canary_plan_kernels", + lambda **kwargs: calls.append("plan"), ), - ), patch.object( - endpoint_module, - "launch_canary_write_kernel", - lambda **kwargs: calls.append( - ("write", kwargs["context"].kernel_kind.name) + patch.object( + endpoint_module, + "launch_canary_verify_kernel", + lambda **kwargs: calls.append( + ("verify", kwargs["context"].kernel_kind.name) + ), + ), + patch.object( + endpoint_module, + "launch_canary_write_kernel", + lambda **kwargs: calls.append( + ("write", kwargs["context"].kernel_kind.name) + ), ), ): manager = make_manager(device=self.device) diff --git a/test/registered/kv_canary/test_self_unit_runner_swa_divergence.py b/test/registered/kv_canary/test_self_unit_runner_swa_divergence.py index f56a330ad..38fb39253 100644 --- a/test/registered/kv_canary/test_self_unit_runner_swa_divergence.py +++ b/test/registered/kv_canary/test_self_unit_runner_swa_divergence.py @@ -447,16 +447,18 @@ class TestSwaDivergenceLogFindAll(CustomTestCase): class TestCanaryManagerSwaDivergenceWiring(CanaryManagerTestCase): def test_swa_divergence_report_is_none_when_env_disabled(self) -> None: - with envs.SGLANG_KV_CANARY_SWA_DIVERGENCE_STATS_INTERVAL.override( - 0 - ), envs.SGLANG_KV_CANARY_PERTURB_TARGET_GROUP.override("full"): + with ( + envs.SGLANG_KV_CANARY_SWA_DIVERGENCE_STATS_INTERVAL.override(0), + envs.SGLANG_KV_CANARY_PERTURB_TARGET_GROUP.override("full"), + ): manager = make_manager(device=self.device) self.assertIsNone(manager._swa_divergence_report) def test_swa_divergence_report_present_when_env_enabled(self) -> None: - with envs.SGLANG_KV_CANARY_SWA_DIVERGENCE_STATS_INTERVAL.override( - 20 - ), envs.SGLANG_KV_CANARY_PERTURB_TARGET_GROUP.override("full"): + with ( + envs.SGLANG_KV_CANARY_SWA_DIVERGENCE_STATS_INTERVAL.override(20), + envs.SGLANG_KV_CANARY_PERTURB_TARGET_GROUP.override("full"), + ): manager = make_manager(device=self.device) self.assertIsNotNone(manager._swa_divergence_report) self.assertIsInstance(manager._swa_divergence_report, SwaDivergenceReporter) diff --git a/test/registered/layers/mamba/test_mamba2_mixer.py b/test/registered/layers/mamba/test_mamba2_mixer.py index 4a127d1b5..173cd4e1a 100644 --- a/test/registered/layers/mamba/test_mamba2_mixer.py +++ b/test/registered/layers/mamba/test_mamba2_mixer.py @@ -42,9 +42,9 @@ def test_mixer2_gated_norm_multi_gpu( if device not in ["cuda", "xpu"]: pytest.skip("Test only supports CUDA and XPU devices") - assert ( - get_device_count() >= NUM_GPUS - ), f"This test requires at least {NUM_GPUS} GPUs, but only {get_device_count()} available" + assert get_device_count() >= NUM_GPUS, ( + f"This test requires at least {NUM_GPUS} GPUs, but only {get_device_count()} available" + ) hidden_size, n_groups = hidden_size_n_groups num_processes = NUM_GPUS diff --git a/test/registered/layers/mamba/test_mamba_ssm_ssd.py b/test/registered/layers/mamba/test_mamba_ssm_ssd.py index 9df125a86..d19eaf391 100644 --- a/test/registered/layers/mamba/test_mamba_ssm_ssd.py +++ b/test/registered/layers/mamba/test_mamba_ssm_ssd.py @@ -181,7 +181,6 @@ def generate_continuous_batched_examples( IND_E = None for spec in example_lens_by_batch: - # get the (maybe partial) example seen in this cont batch dt2, X2, B2, C2 = get_continuous_batch(spec) @@ -356,7 +355,6 @@ def test_mamba_chunk_scan_cont_batch(d_head, n_heads, seq_len_chunk_size_cases, ) in generate_continuous_batched_examples( cases, num_examples, seqlen, last_taken, exhausted, n_heads, d_head, itype ): - chunk_indices, chunk_offsets = ( Mamba2Metadata._query_start_loc_to_chunk_indices_offsets( cu_seqlens, chunk_size, cu_seqlens[-1] @@ -383,7 +381,6 @@ def test_mamba_chunk_scan_cont_batch(d_head, n_heads, seq_len_chunk_size_cases, # just test the last in sequence for i in range(num_examples): - # just test one dim and dstate Y_eg = Y[0, cu_seqlens[i] : cu_seqlens[i + 1], 0, 0] Y_min_eg = Y_min[i][:, 0, 0] diff --git a/test/registered/lora/test_lora_deepseek_v3_base_logprob_diff.py b/test/registered/lora/test_lora_deepseek_v3_base_logprob_diff.py index c7a93597e..64e1c14b1 100644 --- a/test/registered/lora/test_lora_deepseek_v3_base_logprob_diff.py +++ b/test/registered/lora/test_lora_deepseek_v3_base_logprob_diff.py @@ -73,7 +73,6 @@ def get_prompt_logprobs(engine, input_ids, lora_path): class TestLoRADeepSeekV3BaseLogprobDiff(CustomTestCase): - def test_lora_deepseek_v3_base_logprob_accuracy(self): adapter_path = snapshot_download( LORA_HF_REPO, diff --git a/test/registered/lora/test_lora_gpt_oss_20b_logprob_diff.py b/test/registered/lora/test_lora_gpt_oss_20b_logprob_diff.py index 4482dab05..e305befa9 100644 --- a/test/registered/lora/test_lora_gpt_oss_20b_logprob_diff.py +++ b/test/registered/lora/test_lora_gpt_oss_20b_logprob_diff.py @@ -67,7 +67,6 @@ def get_prompt_logprobs(engine, input_ids, lora_path): class TestLoRAGptOss20BLogprobDiff(CustomTestCase): - def test_lora_gpt_oss_20b_logprob_accuracy(self): adapter_path = snapshot_download( LORA_HF_REPO, diff --git a/test/registered/lora/test_lora_kimi_k25_logprob_diff.py b/test/registered/lora/test_lora_kimi_k25_logprob_diff.py index 891997b7f..8ba824dfe 100644 --- a/test/registered/lora/test_lora_kimi_k25_logprob_diff.py +++ b/test/registered/lora/test_lora_kimi_k25_logprob_diff.py @@ -67,7 +67,6 @@ def get_prompt_logprobs(engine, input_ids, lora_path): class TestLoRAKimiK25LogprobDiff(CustomTestCase): - def test_lora_kimi_k25_logprob_accuracy(self): adapter_path = snapshot_download( LORA_HF_REPO, diff --git a/test/registered/lora/test_lora_moe_vllm_sgl_logprob_diff.py b/test/registered/lora/test_lora_moe_vllm_sgl_logprob_diff.py index a35e69dd9..2986f63d8 100644 --- a/test/registered/lora/test_lora_moe_vllm_sgl_logprob_diff.py +++ b/test/registered/lora/test_lora_moe_vllm_sgl_logprob_diff.py @@ -284,7 +284,6 @@ REFERENCE_STATS = { class TestMoELoraRegression(unittest.TestCase): - def test_sglang_moe_parity_strict(self): with SRTRunner( @@ -299,7 +298,6 @@ class TestMoELoraRegression(unittest.TestCase): attention_backend="flashinfer", mem_fraction_static=0.80, ) as srt_runner: - srt_outputs = srt_runner.forward( MOE_LORA_TEST_PROMPTS, max_new_tokens=10, diff --git a/test/registered/lora/test_lora_nemotron_3_super_120b_a12b_logprob_diff.py b/test/registered/lora/test_lora_nemotron_3_super_120b_a12b_logprob_diff.py index fcdc0bf0c..9a19d2826 100644 --- a/test/registered/lora/test_lora_nemotron_3_super_120b_a12b_logprob_diff.py +++ b/test/registered/lora/test_lora_nemotron_3_super_120b_a12b_logprob_diff.py @@ -73,7 +73,6 @@ def get_prompt_logprobs(engine, input_ids, lora_path): class TestLoRANemotron3Super120B_A12B_LogprobDiff(CustomTestCase): - def test_lora_nemotron_3_super_120b_a12b_logprob_accuracy(self): adapter_path = snapshot_download( LORA_HF_REPO, diff --git a/test/registered/lora/test_lora_overlap_loading.py b/test/registered/lora/test_lora_overlap_loading.py index 4d5d9f360..2d73ef920 100644 --- a/test/registered/lora/test_lora_overlap_loading.py +++ b/test/registered/lora/test_lora_overlap_loading.py @@ -42,7 +42,6 @@ class TestLoRAOverlapLoading(CustomTestCase): class TestLoRAOverlapLoaderUnitTests(CustomTestCase): - mock_lora_manager: MagicMock mock_stream: MagicMock mock_stream_context: MagicMock @@ -96,8 +95,8 @@ class TestLoRAOverlapLoaderUnitTests(CustomTestCase): self._create_mock_event(query_return=False), ] self.mock_device_module.Event.side_effect = events - self.mock_lora_manager.validate_lora_batch.side_effect = ( - lambda lora_ids: len(lora_ids) <= 1 + self.mock_lora_manager.validate_lora_batch.side_effect = lambda lora_ids: ( + len(lora_ids) <= 1 ) self.assertTrue( @@ -124,8 +123,8 @@ class TestLoRAOverlapLoaderUnitTests(CustomTestCase): self._create_mock_event(query_return=False), ] self.mock_device_module.Event.side_effect = events - self.mock_lora_manager.validate_lora_batch.side_effect = ( - lambda lora_ids: len(lora_ids) <= 2 + self.mock_lora_manager.validate_lora_batch.side_effect = lambda lora_ids: ( + len(lora_ids) <= 2 ) self.assertTrue(loader._try_start_overlap_load("lora_A", running_loras=set())) diff --git a/test/registered/lora/test_lora_qwen3_30b_a3b_instruct_2507_logprob_diff.py b/test/registered/lora/test_lora_qwen3_30b_a3b_instruct_2507_logprob_diff.py index ec3269688..259bf1c37 100644 --- a/test/registered/lora/test_lora_qwen3_30b_a3b_instruct_2507_logprob_diff.py +++ b/test/registered/lora/test_lora_qwen3_30b_a3b_instruct_2507_logprob_diff.py @@ -67,7 +67,6 @@ def get_prompt_logprobs(engine, input_ids, lora_path): class TestLoRAQwen3_30B_A3B_Instruct_2507_LogprobDiff(CustomTestCase): - def test_lora_qwen3_30b_a3b_instruct_2507_logprob_accuracy(self): adapter_path = snapshot_download( LORA_HF_REPO, diff --git a/test/registered/lora/test_lora_qwen3_5_35b_a3b_logprob_diff.py b/test/registered/lora/test_lora_qwen3_5_35b_a3b_logprob_diff.py index 486ad2bf6..5a68ebdc6 100644 --- a/test/registered/lora/test_lora_qwen3_5_35b_a3b_logprob_diff.py +++ b/test/registered/lora/test_lora_qwen3_5_35b_a3b_logprob_diff.py @@ -77,7 +77,6 @@ def get_prompt_logprobs(engine, input_ids, lora_path): class TestLoRAQwen3_5_35B_A3B_LogprobDiff(CustomTestCase): - def test_lora_qwen3_5_35b_a3b_logprob_accuracy(self): adapter_path = snapshot_download( LORA_HF_REPO, diff --git a/test/registered/lora/test_lora_qwen3_5_4b_logprob_diff.py b/test/registered/lora/test_lora_qwen3_5_4b_logprob_diff.py index 912222039..948fc3e17 100644 --- a/test/registered/lora/test_lora_qwen3_5_4b_logprob_diff.py +++ b/test/registered/lora/test_lora_qwen3_5_4b_logprob_diff.py @@ -69,7 +69,6 @@ def get_prompt_logprobs(engine, input_ids, lora_path): class TestLoRAQwen3_5_4BLogprobDiff(CustomTestCase): - def test_lora_qwen3_5_4b_logprob_accuracy(self): adapter_path = snapshot_download( LORA_HF_REPO, diff --git a/test/registered/lora/test_lora_qwen3_8b_logprob_diff.py b/test/registered/lora/test_lora_qwen3_8b_logprob_diff.py index a11386c34..b56ec7f09 100644 --- a/test/registered/lora/test_lora_qwen3_8b_logprob_diff.py +++ b/test/registered/lora/test_lora_qwen3_8b_logprob_diff.py @@ -103,7 +103,6 @@ def _build_qwen3_mock(): class TestLoRAQwen3_8BLogprobDiff(CustomTestCase): - def test_auto_detect_lora_target_modules(self): """Verify auto_detect_lora_target_modules returns the expected module set for a Qwen3-8B-like (dense) architecture. Catches silent renames diff --git a/test/registered/lora/test_lora_qwen3_vl_30b_a3b_instruct_logprob_diff.py b/test/registered/lora/test_lora_qwen3_vl_30b_a3b_instruct_logprob_diff.py index a558ca646..aa0faad92 100644 --- a/test/registered/lora/test_lora_qwen3_vl_30b_a3b_instruct_logprob_diff.py +++ b/test/registered/lora/test_lora_qwen3_vl_30b_a3b_instruct_logprob_diff.py @@ -71,7 +71,6 @@ def get_prompt_logprobs(engine, input_ids, lora_path): class TestLoRAQwen3VL_30B_A3B_Instruct_LogprobDiff(CustomTestCase): - def test_lora_qwen3_vl_30b_a3b_instruct_logprob_accuracy(self): adapter_path = snapshot_download( LORA_HF_REPO, diff --git a/test/registered/lora/test_lora_tp.py b/test/registered/lora/test_lora_tp.py index 3ed859d11..f22a876dc 100644 --- a/test/registered/lora/test_lora_tp.py +++ b/test/registered/lora/test_lora_tp.py @@ -42,7 +42,6 @@ register_amd_ci( class TestLoRATP(CustomTestCase): - def _run_tp_on_model_cases( self, model_cases: List[LoRAModelCase], diff --git a/test/registered/model_loading/test_load_weights_from_remote_instance.py b/test/registered/model_loading/test_load_weights_from_remote_instance.py index 2fded2ab5..245654044 100644 --- a/test/registered/model_loading/test_load_weights_from_remote_instance.py +++ b/test/registered/model_loading/test_load_weights_from_remote_instance.py @@ -350,7 +350,6 @@ def test_load_weights_from_remote_instance( class TestLoadWeightsFromRemoteInstance(CustomTestCase): - def test_load_weights_from_remote_instance(self): assert torch.cuda.device_count() >= 2, "At least 2 GPUs are required" diff --git a/test/registered/model_loading/test_load_weights_from_remote_instance_npu.py b/test/registered/model_loading/test_load_weights_from_remote_instance_npu.py index cf6f3d7aa..b3034a9d1 100644 --- a/test/registered/model_loading/test_load_weights_from_remote_instance_npu.py +++ b/test/registered/model_loading/test_load_weights_from_remote_instance_npu.py @@ -358,7 +358,6 @@ def test_load_weights_from_remote_instance( class TestLoadWeightsFromRemoteInstance(CustomTestCase): - def test_load_weights_from_remote_instance(self): assert torch.npu.device_count() >= 2, "At least 2 GPUs are required" diff --git a/test/registered/models_e2e/test_dummy_grok_models.py b/test/registered/models_e2e/test_dummy_grok_models.py index e98413e08..f252d6ad3 100644 --- a/test/registered/models_e2e/test_dummy_grok_models.py +++ b/test/registered/models_e2e/test_dummy_grok_models.py @@ -12,7 +12,6 @@ register_cuda_ci( class TestDummyGrok1(CustomTestCase): - def test_dummy_grok_1(self): _, output_throughput, _ = run_bench_one_batch( None, diff --git a/test/registered/models_e2e/test_generation_models.py b/test/registered/models_e2e/test_generation_models.py index 4c1bd79b3..5d78c3d04 100644 --- a/test/registered/models_e2e/test_generation_models.py +++ b/test/registered/models_e2e/test_generation_models.py @@ -129,7 +129,6 @@ TORCH_DTYPES = [torch.float16] class TestGenerationModels(CustomTestCase): - @classmethod def setUpClass(cls): mp.set_start_method("spawn", force=True) diff --git a/test/registered/models_e2e/test_minimax_m25_basic.py b/test/registered/models_e2e/test_minimax_m25_basic.py index 0230ba7bd..ce9e65597 100644 --- a/test/registered/models_e2e/test_minimax_m25_basic.py +++ b/test/registered/models_e2e/test_minimax_m25_basic.py @@ -64,7 +64,7 @@ class TestMiniMaxM25Basic(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_gsm8k (minimax-m25)\n" f'{metrics["accuracy"]=:.3f}\n' + f'### test_gsm8k (minimax-m25)\n{metrics["accuracy"]=:.3f}\n' ) self.assertGreater(metrics["accuracy"], 0.900) @@ -76,7 +76,7 @@ class TestMiniMaxM25Basic(CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_bs_1_speed (minimax-m25)\n" f"{speed=:.2f} token/s\n" + f"### test_bs_1_speed (minimax-m25)\n{speed=:.2f} token/s\n" ) self.assertGreater(speed, 90) diff --git a/test/registered/moe/test_hpc_ops_moe.py b/test/registered/moe/test_hpc_ops_moe.py index 483d3a978..aca7cace2 100644 --- a/test/registered/moe/test_hpc_ops_moe.py +++ b/test/registered/moe/test_hpc_ops_moe.py @@ -60,7 +60,6 @@ def _quant_blockwise(w: torch.Tensor, block: int = 128): "requires HPC-Ops (install from source: https://github.com/Tencent/hpc-ops) and an SM90 (Hopper) GPU", ) class TestHpcOpsMoeBlockwise(CustomTestCase): - @classmethod def setUpClass(cls): set_global_server_args_for_scheduler(ServerArgs(model_path="dummy")) diff --git a/test/registered/npu/accuracy/deepseek_v3_2/test_npu_deepseek_v3_2_8p_aime25.py b/test/registered/npu/accuracy/deepseek_v3_2/test_npu_deepseek_v3_2_8p_aime25.py index 00e5a080b..b6ac070b7 100644 --- a/test/registered/npu/accuracy/deepseek_v3_2/test_npu_deepseek_v3_2_8p_aime25.py +++ b/test/registered/npu/accuracy/deepseek_v3_2/test_npu_deepseek_v3_2_8p_aime25.py @@ -32,7 +32,6 @@ OTHER_ARGS = [ class TestNPUDeepSeek_V3_2_8P_AIME2025(TestNpuAccuracyTestCaseBase): - model = DEEPSEEK_V3_2_EXP_W8A8_WEIGHTS_PATH other_args = OTHER_ARGS accuracy = 0.936 diff --git a/test/registered/npu/accuracy/glm4_7_flash/test_npu_glm4_7_flash_1p_aime25.py b/test/registered/npu/accuracy/glm4_7_flash/test_npu_glm4_7_flash_1p_aime25.py index 9394ca6b1..f3c135924 100644 --- a/test/registered/npu/accuracy/glm4_7_flash/test_npu_glm4_7_flash_1p_aime25.py +++ b/test/registered/npu/accuracy/glm4_7_flash/test_npu_glm4_7_flash_1p_aime25.py @@ -57,7 +57,6 @@ OTHER_ARGS = [ class TestNPUGLM_4_7_FLASH_1P_AIME25(TestNpuAccuracyTestCaseBase): - model = GLM_4_7_FLASH_MODEL_PATH envs = ENVS other_args = OTHER_ARGS diff --git a/test/registered/npu/accuracy/glm4_7_flash/test_npu_glm4_7_flash_1p_gsm8k.py b/test/registered/npu/accuracy/glm4_7_flash/test_npu_glm4_7_flash_1p_gsm8k.py index 1551d0558..bb716c8a1 100644 --- a/test/registered/npu/accuracy/glm4_7_flash/test_npu_glm4_7_flash_1p_gsm8k.py +++ b/test/registered/npu/accuracy/glm4_7_flash/test_npu_glm4_7_flash_1p_gsm8k.py @@ -54,7 +54,6 @@ OTHER_ARGS = [ class TestNPUGlm4_7Flash_1P_GSM8K(TestNpuAccuracyTestCaseBase): - model = GLM_4_7_FLASH_MODEL_PATH envs = ENVS other_args = OTHER_ARGS diff --git a/test/registered/npu/accuracy/glm5_top64_pruned/test_npu_glm5_top64_pruned_bf16_8p_gsm8k.py b/test/registered/npu/accuracy/glm5_top64_pruned/test_npu_glm5_top64_pruned_bf16_8p_gsm8k.py index ccfe108af..3e08a489b 100644 --- a/test/registered/npu/accuracy/glm5_top64_pruned/test_npu_glm5_top64_pruned_bf16_8p_gsm8k.py +++ b/test/registered/npu/accuracy/glm5_top64_pruned/test_npu_glm5_top64_pruned_bf16_8p_gsm8k.py @@ -46,7 +46,6 @@ OTHER_ARGS = [ class TestNPUGLM5_Top64_Pruned_GSM8K(TestNpuAccuracyTestCaseBase): - model = GLM5_TOP64_PRUNED_GSM8K_MODEL_PATH envs = ENVS other_args = OTHER_ARGS diff --git a/test/registered/npu/accuracy/kimi_k2_6/test_npu_kimi_k2_6_w4a8_16p_in64k_out1k_100ms_aime25.py b/test/registered/npu/accuracy/kimi_k2_6/test_npu_kimi_k2_6_w4a8_16p_in64k_out1k_100ms_aime25.py index 780c8b2f9..0490b38be 100644 --- a/test/registered/npu/accuracy/kimi_k2_6/test_npu_kimi_k2_6_w4a8_16p_in64k_out1k_100ms_aime25.py +++ b/test/registered/npu/accuracy/kimi_k2_6/test_npu_kimi_k2_6_w4a8_16p_in64k_out1k_100ms_aime25.py @@ -92,7 +92,6 @@ MODEL_CONFIG = { class TestNPUKimiK2_6_W4A8_16P_AIME2025(TestNpuAccuracyMultiNodePdMixTestCaseBase): - model_config = MODEL_CONFIG accuracy = 0.961 datasets = ["aime25"] diff --git a/test/registered/npu/basic_function/HiCache/test_npu_hicache_mha.py b/test/registered/npu/basic_function/HiCache/test_npu_hicache_mha.py index 4647549ff..a7e5f05e9 100644 --- a/test/registered/npu/basic_function/HiCache/test_npu_hicache_mha.py +++ b/test/registered/npu/basic_function/HiCache/test_npu_hicache_mha.py @@ -26,7 +26,6 @@ TEST_MODEL_MATRIX = { class TestAscendMhaHicache(CustomTestCase): - @classmethod def setUpClass(cls): cls.models = TEST_MODEL_MATRIX.keys() diff --git a/test/registered/npu/basic_function/HiCache/test_npu_hicache_mla.py b/test/registered/npu/basic_function/HiCache/test_npu_hicache_mla.py index 08214ff48..659c9dab8 100644 --- a/test/registered/npu/basic_function/HiCache/test_npu_hicache_mla.py +++ b/test/registered/npu/basic_function/HiCache/test_npu_hicache_mla.py @@ -26,7 +26,6 @@ TEST_MODEL_MATRIX = { class TestAscendMlaHicache(CustomTestCase): - @classmethod def setUpClass(cls): cls.models = TEST_MODEL_MATRIX.keys() diff --git a/test/registered/npu/basic_function/dllm/test_npu_llada2_mini.py b/test/registered/npu/basic_function/dllm/test_npu_llada2_mini.py index e34ee4b93..12c136ef4 100644 --- a/test/registered/npu/basic_function/dllm/test_npu_llada2_mini.py +++ b/test/registered/npu/basic_function/dllm/test_npu_llada2_mini.py @@ -40,8 +40,7 @@ class TestLLaDA2Mini(GSM8KAscendMixin, CustomTestCase): if is_in_ci(): write_github_step_summary( - f"### test_bs_1_speed (llada2-mini) with tp1\n" - f"{speed=:.2f} token/s\n" + f"### test_bs_1_speed (llada2-mini) with tp1\n{speed=:.2f} token/s\n" ) self.assertGreater(speed, 130) diff --git a/test/registered/npu/basic_function/offloading/test_npu_offload_modes.py b/test/registered/npu/basic_function/offloading/test_npu_offload_modes.py index fd29cdb57..b93bc7614 100644 --- a/test/registered/npu/basic_function/offloading/test_npu_offload_modes.py +++ b/test/registered/npu/basic_function/offloading/test_npu_offload_modes.py @@ -21,7 +21,6 @@ TEST_MODEL_MATRIX = { class TestAscendOffloadModes(CustomTestCase): - @classmethod def setUpClass(cls): cls.models = TEST_MODEL_MATRIX diff --git a/test/registered/npu/basic_function/optimization_debug/test_npu_compile_graph_tp1_bf16.py b/test/registered/npu/basic_function/optimization_debug/test_npu_compile_graph_tp1_bf16.py index 05cad7f0e..a164b6763 100644 --- a/test/registered/npu/basic_function/optimization_debug/test_npu_compile_graph_tp1_bf16.py +++ b/test/registered/npu/basic_function/optimization_debug/test_npu_compile_graph_tp1_bf16.py @@ -29,7 +29,6 @@ os.environ["ASCEND_USE_FIA"] = "true" class TestAscendTp1Bf16(CustomTestCase): - @classmethod def setUpClass(cls): cls.models = TEST_MODEL_MATRIX.keys() diff --git a/test/registered/npu/basic_function/optimization_debug/test_npu_graph_tp1_bf16.py b/test/registered/npu/basic_function/optimization_debug/test_npu_graph_tp1_bf16.py index f5cccafde..d174da645 100644 --- a/test/registered/npu/basic_function/optimization_debug/test_npu_graph_tp1_bf16.py +++ b/test/registered/npu/basic_function/optimization_debug/test_npu_graph_tp1_bf16.py @@ -24,7 +24,6 @@ TEST_MODEL_MATRIX = { class TestAscendGraphTp1Bf16(CustomTestCase): - @classmethod def setUpClass(cls): cls.models = TEST_MODEL_MATRIX.keys() diff --git a/test/registered/npu/basic_function/optimization_debug/test_npu_graph_tp2_bf16.py b/test/registered/npu/basic_function/optimization_debug/test_npu_graph_tp2_bf16.py index 7a333b9a1..2cd449eeb 100644 --- a/test/registered/npu/basic_function/optimization_debug/test_npu_graph_tp2_bf16.py +++ b/test/registered/npu/basic_function/optimization_debug/test_npu_graph_tp2_bf16.py @@ -24,7 +24,6 @@ TEST_MODEL_MATRIX = { class TestAscendGraphTp2Bf16(CustomTestCase): - @classmethod def setUpClass(cls): cls.models = TEST_MODEL_MATRIX.keys() diff --git a/test/registered/npu/basic_function/parallel_strategy/expert_parallelism/test_npu_deepep.py b/test/registered/npu/basic_function/parallel_strategy/expert_parallelism/test_npu_deepep.py index e7b158622..cb38a6add 100644 --- a/test/registered/npu/basic_function/parallel_strategy/expert_parallelism/test_npu_deepep.py +++ b/test/registered/npu/basic_function/parallel_strategy/expert_parallelism/test_npu_deepep.py @@ -26,7 +26,6 @@ TEST_MODEL_MATRIX = { class TestAscendDeepEP(CustomTestCase): - @classmethod def setUpClass(cls): cls.models = TEST_MODEL_MATRIX.keys() diff --git a/test/registered/npu/basic_function/parallel_strategy/expert_parallelism/test_npu_eplb_min_rebalancing_utilization_threshold.py b/test/registered/npu/basic_function/parallel_strategy/expert_parallelism/test_npu_eplb_min_rebalancing_utilization_threshold.py index dcba511fe..efad306bc 100644 --- a/test/registered/npu/basic_function/parallel_strategy/expert_parallelism/test_npu_eplb_min_rebalancing_utilization_threshold.py +++ b/test/registered/npu/basic_function/parallel_strategy/expert_parallelism/test_npu_eplb_min_rebalancing_utilization_threshold.py @@ -113,7 +113,7 @@ class TestEplbMinRebalancingUtilizationThresholdBase(CustomTestCase): self.assertGreaterEqual( metrics["score"], self.accuracy, - f'Accuracy of {self.model} is {str(metrics["score"])}, is lower than {self.accuracy}', + f"Accuracy of {self.model} is {str(metrics['score'])}, is lower than {self.accuracy}", ) """ diff --git a/test/registered/npu/basic_function/parameter/test_npu_fim_completion.py b/test/registered/npu/basic_function/parameter/test_npu_fim_completion.py index c335320a2..d1dbcef1e 100644 --- a/test/registered/npu/basic_function/parameter/test_npu_fim_completion.py +++ b/test/registered/npu/basic_function/parameter/test_npu_fim_completion.py @@ -80,9 +80,9 @@ class TestFimCompletion(CustomTestCase): assert response.id assert response.created assert response.object == "text_completion" - assert ( - response.usage.prompt_tokens == num_prompt_tokens - ), f"{response.usage.prompt_tokens} vs {num_prompt_tokens}" + assert response.usage.prompt_tokens == num_prompt_tokens, ( + f"{response.usage.prompt_tokens} vs {num_prompt_tokens}" + ) assert response.usage.completion_tokens > 0 assert response.usage.total_tokens > 0 diff --git a/test/registered/npu/basic_function/quant/test_npu_autoround_dense.py b/test/registered/npu/basic_function/quant/test_npu_autoround_dense.py index 97a8249ec..e91d8caed 100644 --- a/test/registered/npu/basic_function/quant/test_npu_autoround_dense.py +++ b/test/registered/npu/basic_function/quant/test_npu_autoround_dense.py @@ -28,7 +28,6 @@ TEST_MODEL_MATRIX = { class TestAscendAutoRoundDense(CustomTestCase): - @classmethod def setUpClass(cls): cls.models = TEST_MODEL_MATRIX.keys() diff --git a/test/registered/npu/basic_function/quant/test_npu_autoround_moe.py b/test/registered/npu/basic_function/quant/test_npu_autoround_moe.py index 1c831a938..386f36125 100644 --- a/test/registered/npu/basic_function/quant/test_npu_autoround_moe.py +++ b/test/registered/npu/basic_function/quant/test_npu_autoround_moe.py @@ -29,7 +29,6 @@ TEST_MODEL_MATRIX = { class TestAscendAutoRoundMoE(CustomTestCase): - @classmethod def setUpClass(cls): cls.models = TEST_MODEL_MATRIX.keys() diff --git a/test/registered/npu/basic_function/quant/test_npu_gguf.py b/test/registered/npu/basic_function/quant/test_npu_gguf.py index a81d367f6..09557d096 100644 --- a/test/registered/npu/basic_function/quant/test_npu_gguf.py +++ b/test/registered/npu/basic_function/quant/test_npu_gguf.py @@ -26,7 +26,6 @@ TEST_MODEL_MATRIX = { class TestAscendGGUF(CustomTestCase): - @classmethod def setUpClass(cls): cls.models = TEST_MODEL_MATRIX.keys() diff --git a/test/registered/npu/basic_function/quant/test_npu_gguf_moe.py b/test/registered/npu/basic_function/quant/test_npu_gguf_moe.py index 9b7ffa31c..c6b3cb10d 100644 --- a/test/registered/npu/basic_function/quant/test_npu_gguf_moe.py +++ b/test/registered/npu/basic_function/quant/test_npu_gguf_moe.py @@ -28,7 +28,6 @@ TEST_MODEL_MATRIX = { class TestAscendGGUFMoE(CustomTestCase): - @classmethod def setUpClass(cls): cls.models = TEST_MODEL_MATRIX.keys() diff --git a/test/registered/npu/basic_function/quant/test_npu_gptq_moe.py b/test/registered/npu/basic_function/quant/test_npu_gptq_moe.py index bf9d4a296..2512e84d9 100644 --- a/test/registered/npu/basic_function/quant/test_npu_gptq_moe.py +++ b/test/registered/npu/basic_function/quant/test_npu_gptq_moe.py @@ -30,7 +30,6 @@ TEST_MODEL_MATRIX = { class TestAscendGPTQMoEInt4(CustomTestCase): - @classmethod def setUpClass(cls): cls.models = TEST_MODEL_MATRIX.keys() diff --git a/test/registered/npu/basic_function/quant/test_npu_w4a4_quantization.py b/test/registered/npu/basic_function/quant/test_npu_w4a4_quantization.py index f3743bcb0..42c6db12e 100644 --- a/test/registered/npu/basic_function/quant/test_npu_w4a4_quantization.py +++ b/test/registered/npu/basic_function/quant/test_npu_w4a4_quantization.py @@ -18,7 +18,6 @@ register_npu_ci(est_time=400, suite="nightly-4-npu-a3", nightly=True) class TestAscendW4A4(GSM8KAscendMixin, CustomTestCase): - model = ECO_TECH_QWEN3_32B_W4A4_LAOS_WEIGHTS_PATH other_args = [ "--trust-remote-code", diff --git a/test/registered/npu/basic_function/runtime_opts/test_npu_mla_fia_w8a8int8.py b/test/registered/npu/basic_function/runtime_opts/test_npu_mla_fia_w8a8int8.py index 0fb4236be..664fa851d 100644 --- a/test/registered/npu/basic_function/runtime_opts/test_npu_mla_fia_w8a8int8.py +++ b/test/registered/npu/basic_function/runtime_opts/test_npu_mla_fia_w8a8int8.py @@ -27,7 +27,6 @@ TEST_MODEL_MATRIX = { class TestAscendMlaW8A8Int8(CustomTestCase): - @classmethod def setUpClass(cls): cls.models = TEST_MODEL_MATRIX.keys() diff --git a/test/registered/npu/basic_function/runtime_opts/test_npu_mla_w8a8int8.py b/test/registered/npu/basic_function/runtime_opts/test_npu_mla_w8a8int8.py index 50dcd88fa..5de0dcd6a 100644 --- a/test/registered/npu/basic_function/runtime_opts/test_npu_mla_w8a8int8.py +++ b/test/registered/npu/basic_function/runtime_opts/test_npu_mla_w8a8int8.py @@ -26,7 +26,6 @@ TEST_MODEL_MATRIX = { class TestAscendMlaW8A8Int8(CustomTestCase): - @classmethod def setUpClass(cls): cls.models = TEST_MODEL_MATRIX.keys() diff --git a/test/registered/npu/basic_function/runtime_opts/test_npu_tp1_bf16.py b/test/registered/npu/basic_function/runtime_opts/test_npu_tp1_bf16.py index 1556c62d7..5a77c9fe2 100644 --- a/test/registered/npu/basic_function/runtime_opts/test_npu_tp1_bf16.py +++ b/test/registered/npu/basic_function/runtime_opts/test_npu_tp1_bf16.py @@ -26,7 +26,6 @@ TEST_MODEL_MATRIX = { class TestAscendTp1Bf16(CustomTestCase): - @classmethod def setUpClass(cls): cls.models = TEST_MODEL_MATRIX.keys() diff --git a/test/registered/npu/basic_function/runtime_opts/test_npu_tp2_bf16.py b/test/registered/npu/basic_function/runtime_opts/test_npu_tp2_bf16.py index 5a88ad3b0..4f900d4c7 100644 --- a/test/registered/npu/basic_function/runtime_opts/test_npu_tp2_bf16.py +++ b/test/registered/npu/basic_function/runtime_opts/test_npu_tp2_bf16.py @@ -26,7 +26,6 @@ TEST_MODEL_MATRIX = { class TestAscendTp2Bf16(CustomTestCase): - @classmethod def setUpClass(cls): cls.models = TEST_MODEL_MATRIX.keys() diff --git a/test/registered/npu/basic_function/runtime_opts/test_npu_tp2_fia_bf16.py b/test/registered/npu/basic_function/runtime_opts/test_npu_tp2_fia_bf16.py index bf2ff3361..8a4e87dc4 100644 --- a/test/registered/npu/basic_function/runtime_opts/test_npu_tp2_fia_bf16.py +++ b/test/registered/npu/basic_function/runtime_opts/test_npu_tp2_fia_bf16.py @@ -27,7 +27,6 @@ TEST_MODEL_MATRIX = { class TestAscendTp2Bf16(CustomTestCase): - @classmethod def setUpClass(cls): cls.models = TEST_MODEL_MATRIX.keys() diff --git a/test/registered/npu/basic_function/runtime_opts/test_npu_tp4_bf16.py b/test/registered/npu/basic_function/runtime_opts/test_npu_tp4_bf16.py index f9842d10b..638549082 100644 --- a/test/registered/npu/basic_function/runtime_opts/test_npu_tp4_bf16.py +++ b/test/registered/npu/basic_function/runtime_opts/test_npu_tp4_bf16.py @@ -25,7 +25,6 @@ TEST_MODEL_MATRIX = { class TestAscendTp4Bf16(CustomTestCase): - @classmethod def setUpClass(cls): cls.models = TEST_MODEL_MATRIX.keys() diff --git a/test/registered/npu/basic_function/speculative_inference/test_npu_speculative_moe_a2a_backend.py b/test/registered/npu/basic_function/speculative_inference/test_npu_speculative_moe_a2a_backend.py index 0604897bd..f4a1e1b3f 100644 --- a/test/registered/npu/basic_function/speculative_inference/test_npu_speculative_moe_a2a_backend.py +++ b/test/registered/npu/basic_function/speculative_inference/test_npu_speculative_moe_a2a_backend.py @@ -21,7 +21,6 @@ register_npu_ci(est_time=400, suite="nightly-16-npu-a3", nightly=True) class TestAscendDistTimeout(CustomTestCase): - @classmethod def setUpClass(cls): cls.model = DEEPSEEK_R1_0528_W8A8_WEIGHTS_PATH diff --git a/test/registered/npu/embedding_models/test_npu_bge_large_en_v1_5.py b/test/registered/npu/embedding_models/test_npu_bge_large_en_v1_5.py index 5da11e7a4..d5a6e4e39 100644 --- a/test/registered/npu/embedding_models/test_npu_bge_large_en_v1_5.py +++ b/test/registered/npu/embedding_models/test_npu_bge_large_en_v1_5.py @@ -29,7 +29,6 @@ TORCH_DTYPES = [torch.float16] class TestEmbeddingModels(CustomTestCase): - @classmethod def setUpClass(cls): mp.set_start_method("spawn", force=True) @@ -94,9 +93,9 @@ class TestEmbeddingModels(CustomTestCase): print("similarity diff", abs(similarity - 1)) if len(prompts[i]) <= 1000: - assert torch.all( - abs(similarity - 1) < prefill_tolerance - ), "embeddings are not all close" + assert torch.all(abs(similarity - 1) < prefill_tolerance), ( + "embeddings are not all close" + ) def test_prefill_logits(self): models_to_test = MODELS diff --git a/test/registered/npu/interface/test_npu_api_abort_request.py b/test/registered/npu/interface/test_npu_api_abort_request.py index 7e62c4d6d..084215773 100644 --- a/test/registered/npu/interface/test_npu_api_abort_request.py +++ b/test/registered/npu/interface/test_npu_api_abort_request.py @@ -74,5 +74,4 @@ class TestNpuApi(CustomTestCase): if __name__ == "__main__": - unittest.main() diff --git a/test/registered/npu/interface/test_npu_openai_function_calling.py b/test/registered/npu/interface/test_npu_openai_function_calling.py index daea9486b..e41d18dbb 100644 --- a/test/registered/npu/interface/test_npu_openai_function_calling.py +++ b/test/registered/npu/interface/test_npu_openai_function_calling.py @@ -119,9 +119,9 @@ class TestOpenAIServerFunctionCalling(CustomTestCase): tool_calls = response.choices[0].message.tool_calls - assert ( - isinstance(tool_calls, list) and len(tool_calls) > 0 - ), "tool_calls should be a non-empty list" + assert isinstance(tool_calls, list) and len(tool_calls) > 0, ( + "tool_calls should be a non-empty list" + ) function_name = tool_calls[0].function.name assert function_name == "add", "Function name should be 'add'" @@ -199,9 +199,9 @@ class TestOpenAIServerFunctionCalling(CustomTestCase): tools=tools, ) - assert ( - "8" in final_response.choices[0].message.content - ), "tool_call response should have the sum 8 in the content" + assert "8" in final_response.choices[0].message.content, ( + "tool_call response should have the sum 8 in the content" + ) def test_function_calling_streaming_simple(self): """ diff --git a/test/registered/npu/rerank_models/test_npu_bge_reranker_v2_m3.py b/test/registered/npu/rerank_models/test_npu_bge_reranker_v2_m3.py index 8b552741a..d395bae96 100644 --- a/test/registered/npu/rerank_models/test_npu_bge_reranker_v2_m3.py +++ b/test/registered/npu/rerank_models/test_npu_bge_reranker_v2_m3.py @@ -63,9 +63,9 @@ class TestBgeReranker(CustomTestCase): for i in range(len(srt_scores)): score_difference = abs(hf_scores[i] - srt_scores[i]) - assert ( - score_difference < score_tolerance - ), "cross encoder scores are not all close" + assert score_difference < score_tolerance, ( + "cross encoder scores are not all close" + ) def preprocess_prompts(self, prompt): processed_prompts = [] diff --git a/test/registered/npu/reward_models/test_npu_gemma_2_27b_v0_2.py b/test/registered/npu/reward_models/test_npu_gemma_2_27b_v0_2.py index 086a82342..b30999e35 100644 --- a/test/registered/npu/reward_models/test_npu_gemma_2_27b_v0_2.py +++ b/test/registered/npu/reward_models/test_npu_gemma_2_27b_v0_2.py @@ -34,7 +34,6 @@ CONVS = [ class TestRewardModels(CustomTestCase): - @classmethod def setUpClass(cls): mp.set_start_method("spawn", force=True) @@ -70,9 +69,9 @@ class TestRewardModels(CustomTestCase): logger.info(f"{hf_scores=}") logger.info(f"{srt_scores=}") - assert torch.all( - abs(hf_scores - srt_scores) < tolerance - ), "reward scores are not all close" + assert torch.all(abs(hf_scores - srt_scores) < tolerance), ( + "reward scores are not all close" + ) def test_reward_scores(self): for model, tp_size, tolerance in MODELS: diff --git a/test/registered/npu/reward_models/test_npu_llama_3_1_8b_v0_2.py b/test/registered/npu/reward_models/test_npu_llama_3_1_8b_v0_2.py index c702e8d0b..14118d683 100644 --- a/test/registered/npu/reward_models/test_npu_llama_3_1_8b_v0_2.py +++ b/test/registered/npu/reward_models/test_npu_llama_3_1_8b_v0_2.py @@ -70,9 +70,9 @@ class TestLlama(CustomTestCase): print(f"{hf_scores=}") print(f"{srt_scores=}") - assert torch.all( - abs(hf_scores - srt_scores) < tolerance - ), "reward scores are not all close" + assert torch.all(abs(hf_scores - srt_scores) < tolerance), ( + "reward scores are not all close" + ) def test_reward_scores(self): for model, tp_size, tolerance in MODELS: diff --git a/test/registered/npu/test_npu_memory_consumption.py b/test/registered/npu/test_npu_memory_consumption.py index b9dd8b0f6..51229b61e 100644 --- a/test/registered/npu/test_npu_memory_consumption.py +++ b/test/registered/npu/test_npu_memory_consumption.py @@ -32,7 +32,6 @@ DEFAULT_URL_FOR_TEST = f"http://127.0.0.1:{DEFAULT_PORT_FOR_SRT_TEST_RUNNER + 10 class TestMemoryConsumptionAscend(CustomTestCase): - def test_memory_consumption(self): model = "/root/.cache/modelscope/hub/models/Qwen/Qwen3-30B-A3B-w8a8" diff --git a/test/registered/openai_server/basic/test_openai_server.py b/test/registered/openai_server/basic/test_openai_server.py index a3c517602..6b4f3f46c 100644 --- a/test/registered/openai_server/basic/test_openai_server.py +++ b/test/registered/openai_server/basic/test_openai_server.py @@ -193,9 +193,9 @@ class TestOpenAIServer(CustomTestCase, AnthropicMessagesMixin): if logprobs: assert response.choices[0].logprobs - assert isinstance( - response.choices[0].logprobs.tokens[0], str - ), f"{response=}" + assert isinstance(response.choices[0].logprobs.tokens[0], str), ( + f"{response=}" + ) assert isinstance(response.choices[0].logprobs.top_logprobs[1], dict) ret_num_top_logprobs = len(response.choices[0].logprobs.top_logprobs[1]) @@ -209,9 +209,9 @@ class TestOpenAIServer(CustomTestCase, AnthropicMessagesMixin): assert response.id assert response.created - assert ( - response.usage.prompt_tokens == num_prompt_tokens - ), f"{response.usage.prompt_tokens} vs {num_prompt_tokens}" + assert response.usage.prompt_tokens == num_prompt_tokens, ( + f"{response.usage.prompt_tokens} vs {num_prompt_tokens}" + ) assert response.usage.completion_tokens > 0 assert response.usage.total_tokens > 0 @@ -264,9 +264,9 @@ class TestOpenAIServer(CustomTestCase, AnthropicMessagesMixin): if logprobs: assert response.choices[0].logprobs, f"no logprobs in response" - assert isinstance( - response.choices[0].logprobs.tokens[0], str - ), f"{response.choices[0].logprobs.tokens[0]} is not a string" + assert isinstance(response.choices[0].logprobs.tokens[0], str), ( + f"{response.choices[0].logprobs.tokens[0]} is not a string" + ) if not (is_first and echo): assert isinstance( response.choices[0].logprobs.top_logprobs[0], dict @@ -280,17 +280,17 @@ class TestOpenAIServer(CustomTestCase, AnthropicMessagesMixin): if is_first: if echo: - assert response.choices[0].text.startswith( - prompt - ), f"{response.choices[0].text} and all args {echo} {logprobs} {token_input} {is_first}" + assert response.choices[0].text.startswith(prompt), ( + f"{response.choices[0].text} and all args {echo} {logprobs} {token_input} {is_first}" + ) is_firsts[index] = False assert response.id, f"no id in response" assert response.created, f"no created in response" for index in [i for i in range(parallel_sample_num * num_choices)]: - assert not is_firsts.get( - index, True - ), f"index {index} is not found in the response" + assert not is_firsts.get(index, True), ( + f"index {index} is not found in the response" + ) def run_chat_completion(self, logprobs, parallel_sample_num): client = openai.Client(api_key=self.api_key, base_url=self.base_url) @@ -317,9 +317,9 @@ class TestOpenAIServer(CustomTestCase, AnthropicMessagesMixin): ret_num_top_logprobs = len( response.choices[0].logprobs.content[0].top_logprobs ) - assert ( - ret_num_top_logprobs == logprobs - ), f"{ret_num_top_logprobs} vs {logprobs}" + assert ret_num_top_logprobs == logprobs, ( + f"{ret_num_top_logprobs} vs {logprobs}" + ) assert len(response.choices) == parallel_sample_num assert response.choices[0].message.role == "assistant" @@ -366,9 +366,9 @@ class TestOpenAIServer(CustomTestCase, AnthropicMessagesMixin): data = response.choices[0].delta if is_firsts.get(index, True): - assert ( - data.role == "assistant" - ), f"data.role was not 'assistant' for first chunk" + assert data.role == "assistant", ( + f"data.role was not 'assistant' for first chunk" + ) is_firsts[index] = False continue @@ -383,9 +383,9 @@ class TestOpenAIServer(CustomTestCase, AnthropicMessagesMixin): ret_num_top_logprobs = len( response.choices[0].logprobs.content[0].top_logprobs ) - assert ( - ret_num_top_logprobs == logprobs - ), f"{ret_num_top_logprobs} vs {logprobs}" + assert ret_num_top_logprobs == logprobs, ( + f"{ret_num_top_logprobs} vs {logprobs}" + ) assert ( isinstance(data.content, str) @@ -397,18 +397,18 @@ class TestOpenAIServer(CustomTestCase, AnthropicMessagesMixin): assert response.created for index in [i for i in range(parallel_sample_num)]: - assert not is_firsts.get( - index, True - ), f"index {index} is not found in the response" + assert not is_firsts.get(index, True), ( + f"index {index} is not found in the response" + ) # Verify that each choice gets exactly one finish_reason chunk for index in range(parallel_sample_num): - assert ( - index in finish_reason_counts - ), f"No finish_reason found for index {index}" - assert ( - finish_reason_counts[index] == 1 - ), f"Expected 1 finish_reason chunk for index {index}, got {finish_reason_counts[index]}" + assert index in finish_reason_counts, ( + f"No finish_reason found for index {index}" + ) + assert finish_reason_counts[index] == 1, ( + f"Expected 1 finish_reason chunk for index {index}, got {finish_reason_counts[index]}" + ) def test_completion(self): for echo in [False, True]: diff --git a/test/registered/openai_server/features/test_openai_server_hidden_states.py b/test/registered/openai_server/features/test_openai_server_hidden_states.py index 76ae3f8bf..31b5fc2a5 100644 --- a/test/registered/openai_server/features/test_openai_server_hidden_states.py +++ b/test/registered/openai_server/features/test_openai_server_hidden_states.py @@ -25,7 +25,6 @@ register_amd_ci( class BaseTestOpenAIServerWithHiddenStates(ABC): - @classmethod def setUpClass(cls): cls.return_hidden_states = [False, True] @@ -55,18 +54,14 @@ class BaseTestOpenAIServerWithHiddenStates(ABC): def test_chat_completion(self): for return_hidden_states in self.return_hidden_states: - for ( - parallel_sample_num - ) in ( + for parallel_sample_num in ( self.parallel_sample_nums ): # parallel sample num 2 breaks in the adapter with a 400 for EAGLE self.run_chat_completion(parallel_sample_num, return_hidden_states) def test_chat_completion_stream(self): for return_hidden_states in self.return_hidden_states: - for ( - parallel_sample_num - ) in ( + for parallel_sample_num in ( self.parallel_sample_nums ): # parallel sample num > 1 breaks in the adapter with a 400 for EAGLE self.run_chat_completion_stream( @@ -144,13 +139,13 @@ class BaseTestOpenAIServerWithHiddenStates(ABC): hidden_states_list.append(choice.hidden_states) if return_hidden_states: - assert ( - len(hidden_states_list) == parallel_sample_num * num_choices - ), f"Expected {parallel_sample_num * num_choices} hidden states, got {len(hidden_states_list)}" + assert len(hidden_states_list) == parallel_sample_num * num_choices, ( + f"Expected {parallel_sample_num * num_choices} hidden states, got {len(hidden_states_list)}" + ) else: - assert ( - hidden_states_list == [] - ), "hidden_states were returned and should not have been" + assert hidden_states_list == [], ( + "hidden_states were returned and should not have been" + ) def run_chat_completion(self, parallel_sample_num, return_hidden_states): client = openai.Client(api_key=self.api_key, base_url=self.base_url) @@ -201,13 +196,13 @@ class BaseTestOpenAIServerWithHiddenStates(ABC): hidden_states_list.append(choice.delta.hidden_states) if return_hidden_states: - assert ( - len(hidden_states_list) == parallel_sample_num - ), f"Expected {parallel_sample_num} hidden states, got {len(hidden_states_list)}" + assert len(hidden_states_list) == parallel_sample_num, ( + f"Expected {parallel_sample_num} hidden states, got {len(hidden_states_list)}" + ) else: - assert ( - hidden_states_list == [] - ), "hidden_states were returned and should not have been" + assert hidden_states_list == [], ( + "hidden_states were returned and should not have been" + ) class TestOpenAIServerWithHiddenStatesEnabled( diff --git a/test/registered/openai_server/function_call/test_openai_function_calling.py b/test/registered/openai_server/function_call/test_openai_function_calling.py index 8ad8139ad..ed45d9e6e 100644 --- a/test/registered/openai_server/function_call/test_openai_function_calling.py +++ b/test/registered/openai_server/function_call/test_openai_function_calling.py @@ -110,9 +110,9 @@ class TestOpenAIServerFunctionCalling(CustomTestCase): tool_calls = response.choices[0].message.tool_calls - assert ( - isinstance(tool_calls, list) and len(tool_calls) > 0 - ), "tool_calls should be a non-empty list" + assert isinstance(tool_calls, list) and len(tool_calls) > 0, ( + "tool_calls should be a non-empty list" + ) function_name = tool_calls[0].function.name assert function_name == "add", "Function name should be 'add'" @@ -191,9 +191,9 @@ class TestOpenAIServerFunctionCalling(CustomTestCase): tools=tools, ) - assert ( - "8" in final_response.choices[0].message.content - ), "tool_call response should have the sum 8 in the content" + assert "8" in final_response.choices[0].message.content, ( + "tool_call response should have the sum 8 in the content" + ) def test_function_calling_streaming_simple(self): """ diff --git a/test/registered/ops/test_aiter_allgather_amd.py b/test/registered/ops/test_aiter_allgather_amd.py index 01595c95d..0103ae615 100644 --- a/test/registered/ops/test_aiter_allgather_amd.py +++ b/test/registered/ops/test_aiter_allgather_amd.py @@ -12,7 +12,6 @@ register_amd_ci(est_time=180, suite="stage-c-test-large-8-gpu-amd") class TestAiterAllGatherAmd(unittest.TestCase): - @staticmethod def _gpu_count(): return torch.cuda.device_count() if torch.cuda.is_available() else 0 diff --git a/test/registered/ops/test_aiter_allreduce_fusion_amd.py b/test/registered/ops/test_aiter_allreduce_fusion_amd.py index 7b99f4b6c..926e4f338 100755 --- a/test/registered/ops/test_aiter_allreduce_fusion_amd.py +++ b/test/registered/ops/test_aiter_allreduce_fusion_amd.py @@ -158,7 +158,6 @@ def _run_residual_accuracy_check(): class TestAiterAllreduceFusionAmd(unittest.TestCase): - @staticmethod def _gpu_count(): return torch.cuda.device_count() if torch.cuda.is_available() else 0 diff --git a/test/registered/perf/test_bench_one_batch_2gpu.py b/test/registered/perf/test_bench_one_batch_2gpu.py index 89505b3d4..8be302f5c 100644 --- a/test/registered/perf/test_bench_one_batch_2gpu.py +++ b/test/registered/perf/test_bench_one_batch_2gpu.py @@ -16,7 +16,6 @@ register_amd_ci(est_time=630, suite="stage-b-test-2-gpu-large-amd") class TestBenchOneBatch2GPU(CustomTestCase): - def test_moe_tp2_bs1(self): output_throughput = run_bench_offline_throughput( DEFAULT_MOE_MODEL_NAME_FOR_TEST, diff --git a/test/registered/prefill_only/test_embedding_models.py b/test/registered/prefill_only/test_embedding_models.py index a0a3d6c7b..e3f0c4824 100644 --- a/test/registered/prefill_only/test_embedding_models.py +++ b/test/registered/prefill_only/test_embedding_models.py @@ -57,7 +57,6 @@ TORCH_DTYPES = [torch.float16] class TestEmbeddingModels(CustomTestCase): - @classmethod def setUpClass(cls): mp.set_start_method("spawn", force=True) @@ -122,9 +121,9 @@ class TestEmbeddingModels(CustomTestCase): print("similarity diff", abs(similarity - 1)) if len(prompts[i]) <= 1000: - assert torch.all( - abs(similarity - 1) < prefill_tolerance - ), "embeddings are not all close" + assert torch.all(abs(similarity - 1) < prefill_tolerance), ( + "embeddings are not all close" + ) def test_prefill_logits(self): models_to_test = MODELS diff --git a/test/registered/prefill_only/test_reward_models.py b/test/registered/prefill_only/test_reward_models.py index b839a434e..90f8ead06 100644 --- a/test/registered/prefill_only/test_reward_models.py +++ b/test/registered/prefill_only/test_reward_models.py @@ -52,7 +52,6 @@ CONVS = [ class TestRewardModels(CustomTestCase): - @classmethod def setUpClass(cls): mp.set_start_method("spawn", force=True) @@ -87,9 +86,9 @@ class TestRewardModels(CustomTestCase): print(f"{hf_scores=}") print(f"{srt_scores=}") - assert torch.all( - abs(hf_scores - srt_scores) < tolerance - ), "reward scores are not all close" + assert torch.all(abs(hf_scores - srt_scores) < tolerance), ( + "reward scores are not all close" + ) def test_reward_scores(self): for model, tp_size, tolerance in MODELS: diff --git a/test/registered/profiling/test_diffusion_stage_profiler_sync.py b/test/registered/profiling/test_diffusion_stage_profiler_sync.py index 38993033a..82ba6fe29 100644 --- a/test/registered/profiling/test_diffusion_stage_profiler_sync.py +++ b/test/registered/profiling/test_diffusion_stage_profiler_sync.py @@ -40,9 +40,9 @@ def test_stage_entry_sync_excludes_previous_stage_tail(monkeypatch): torch.ones(8, device="cuda").sum().cpu() producer_ms, consumer_ms = metrics.stages["producer"], metrics.stages["consumer"] - assert ( - producer_ms > 250 - ), f"queued work not attributed to producer: {metrics.stages}" + assert producer_ms > 250, ( + f"queued work not attributed to producer: {metrics.stages}" + ) assert consumer_ms < 100, f"producer tail leaked into consumer: {metrics.stages}" diff --git a/test/registered/profiling/test_profile_v2.py b/test/registered/profiling/test_profile_v2.py index 17b194f2d..ab21fefad 100644 --- a/test/registered/profiling/test_profile_v2.py +++ b/test/registered/profiling/test_profile_v2.py @@ -26,7 +26,6 @@ register_cuda_ci( class TestStartProfile(CustomTestCase): - @classmethod def setUpClass(cls): cls.output_dir = tempfile.mkdtemp() diff --git a/test/registered/profiling/test_start_profile.py b/test/registered/profiling/test_start_profile.py index 10b8b31d8..15d2092de 100644 --- a/test/registered/profiling/test_start_profile.py +++ b/test/registered/profiling/test_start_profile.py @@ -45,7 +45,6 @@ def _is_nsys_available(): class TestStartProfile(CustomTestCase): - @classmethod def setUpClass(cls): envs.SGLANG_TORCH_PROFILER_DIR.set(OUTPUT_DIR) diff --git a/test/registered/quant/test_block_int8.py b/test/registered/quant/test_block_int8.py index dc456b2f4..6ebd41af2 100644 --- a/test/registered/quant/test_block_int8.py +++ b/test/registered/quant/test_block_int8.py @@ -22,9 +22,9 @@ def native_per_token_group_quant_int8(x, group_size, eps=1e-10, dtype=torch.int8 quantized tensor along with the scaling factor used for quantization. Note that only `torch.float8_e4m3fn` is supported for now. """ - assert ( - x.shape[-1] % group_size == 0 - ), "the last dimension of `x` cannot be divisible by `group_size`" + assert x.shape[-1] % group_size == 0, ( + "the last dimension of `x` cannot be divisible by `group_size`" + ) assert x.is_contiguous(), "`x` is not contiguous" iinfo = torch.iinfo(dtype) diff --git a/test/registered/quant/test_fp8_utils.py b/test/registered/quant/test_fp8_utils.py index 77f0c4db5..c256c5d1d 100644 --- a/test/registered/quant/test_fp8_utils.py +++ b/test/registered/quant/test_fp8_utils.py @@ -37,12 +37,12 @@ class TestInverseTransformScaleUe8m0(CustomTestCase): sf_packed_recreated = transform_scale_ue8m0(sf_fp32_recreated, mn=mn) - assert torch.all( - sf_packed_original == sf_packed_recreated - ), f"{sf_packed_original=} {sf_packed_recreated}" - assert torch.all( - sf_fp32_original == sf_fp32_recreated - ), f"{sf_fp32_original=} {sf_fp32_recreated}" + assert torch.all(sf_packed_original == sf_packed_recreated), ( + f"{sf_packed_original=} {sf_packed_recreated}" + ) + assert torch.all(sf_fp32_original == sf_fp32_recreated), ( + f"{sf_fp32_original=} {sf_fp32_recreated}" + ) class TestApplyFp8LinearScaleDispatch(CustomTestCase): @@ -97,14 +97,16 @@ class TestApplyFp8LinearScaleDispatch(CustomTestCase): "is_sm120": False, } capabilities[capability] = True - with patch.object( - fp8_utils, - "get_platform", - return_value=SimpleNamespace(**capabilities), - ), patch.object( - fp8_utils, "fp8_scaled_mm", side_effect=fake_fp8_scaled_mm - ), patch.object( - fp8_utils, "get_exec", return_value=exec_config + with ( + patch.object( + fp8_utils, + "get_platform", + return_value=SimpleNamespace(**capabilities), + ), + patch.object( + fp8_utils, "fp8_scaled_mm", side_effect=fake_fp8_scaled_mm + ), + patch.object(fp8_utils, "get_exec", return_value=exec_config), ): fp8_utils.apply_fp8_linear( input, @@ -157,15 +159,18 @@ class TestApplyFp8LinearScaleDispatch(CustomTestCase): (mat_a.shape[0], mat_b.shape[1]), dtype=out_dtype, device=mat_a.device ) - with patch.object( - fp8_utils, - "get_platform", - return_value=SimpleNamespace( - is_sm90=False, - is_sm100=False, - is_sm120=False, + with ( + patch.object( + fp8_utils, + "get_platform", + return_value=SimpleNamespace( + is_sm90=False, + is_sm100=False, + is_sm120=False, + ), ), - ), patch.object(fp8_utils, "fp8_scaled_mm", side_effect=fake_fp8_scaled_mm): + patch.object(fp8_utils, "fp8_scaled_mm", side_effect=fake_fp8_scaled_mm), + ): fp8_utils.apply_fp8_linear( input, weight, diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_hicache_pp_kl.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_hicache_pp_kl.py index dd460ad0c..4433dc30b 100644 --- a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_hicache_pp_kl.py +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_hicache_pp_kl.py @@ -28,9 +28,9 @@ def _assert_pp_decode_cached_tokens(result, history_len, output_len, label): expected = history_len + output_len actual = result["meta_info"]["cached_tokens"] lower = max(0, expected - 1) - assert ( - lower <= actual <= expected - ), f"{label}: expected cached_tokens in [{lower}, {expected}], got {actual}" + assert lower <= actual <= expected, ( + f"{label}: expected cached_tokens in [{lower}, {expected}], got {actual}" + ) class TestUnifiedQwen3HiCachePP(UnifiedRadixTreeTestMixin, CustomTestCase): diff --git a/test/registered/rl/test_lora_load_from_tensor.py b/test/registered/rl/test_lora_load_from_tensor.py index c37aab68b..586560e23 100644 --- a/test/registered/rl/test_lora_load_from_tensor.py +++ b/test/registered/rl/test_lora_load_from_tensor.py @@ -63,7 +63,7 @@ class TestLoRALoadFromTensor(CustomTestCase): # This should trigger LRU eviction when we exceed the limit TEST_LORA_COUNT = 10 for i in range(TEST_LORA_COUNT): - print(f"[Test]Loading LoRA adapter {i+1}/10: self_cognition_Alice_{i}") + print(f"[Test]Loading LoRA adapter {i + 1}/10: self_cognition_Alice_{i}") result = test_engine.load_lora_adapter_from_tensors( lora_name=f"self_cognition_Alice_{i}", tensors=self.lora_tensors, @@ -74,7 +74,7 @@ class TestLoRALoadFromTensor(CustomTestCase): f"Failed to load LoRA adapter {i}: {result.error_message}", ) print( - f"[Test]Successfully loaded LoRA {i+1}, current loaded adapters: {list(result.loaded_adapters.keys())}" + f"[Test]Successfully loaded LoRA {i + 1}, current loaded adapters: {list(result.loaded_adapters.keys())}" ) EXPECTED_LORA_ADAPTERS = [ diff --git a/test/registered/rl/test_return_routed_experts.py b/test/registered/rl/test_return_routed_experts.py index d8ae2be83..d06e607a5 100644 --- a/test/registered/rl/test_return_routed_experts.py +++ b/test/registered/rl/test_return_routed_experts.py @@ -160,14 +160,14 @@ class TestReturnRoutedExperts(CustomTestCase): captured_baseline_experts, captured_reference_experts ) logger.info( - f"Total mismatches report: {num_mismatches} out of {num_baseline_topks} ({num_mismatches/num_baseline_topks:.4%})" + f"Total mismatches report: {num_mismatches} out of {num_baseline_topks} ({num_mismatches / num_baseline_topks:.4%})" ) print( - f"Total mismatches report: {num_mismatches} out of {num_baseline_topks} ({num_mismatches/num_baseline_topks:.4%})" + f"Total mismatches report: {num_mismatches} out of {num_baseline_topks} ({num_mismatches / num_baseline_topks:.4%})" + ) + assert num_mismatches / num_baseline_topks < 0.10, ( + f"Too many mismatches: {num_mismatches} out of {num_baseline_topks} ({num_mismatches / num_baseline_topks:.4%})" ) - assert ( - num_mismatches / num_baseline_topks < 0.10 - ), f"Too many mismatches: {num_mismatches} out of {num_baseline_topks} ({num_mismatches/num_baseline_topks:.4%})" @classmethod def _collect_results( diff --git a/test/registered/rl/test_update_weights_from_distributed.py b/test/registered/rl/test_update_weights_from_distributed.py index 1b3e2adac..f1cdf8b6d 100644 --- a/test/registered/rl/test_update_weights_from_distributed.py +++ b/test/registered/rl/test_update_weights_from_distributed.py @@ -636,9 +636,9 @@ def test_update_weights_from_distributed( f"sgl_dp_2_instruct_params rank {i}", ) - assert len(params["hf_instruct"]) == len( - params["hf_base"] - ), "hf_instruct_params and hf_base_params have different lengths" + assert len(params["hf_instruct"]) == len(params["hf_base"]), ( + "hf_instruct_params and hf_base_params have different lengths" + ) # Check if the weights of lm_head are tied with embed_tokens. params_to_check = [ @@ -688,18 +688,18 @@ def test_update_weights_from_distributed( # On local H100, it's 1 / 2 time_limit = 3 if model_name == DEFAULT_SMALL_MODEL_NAME_FOR_TEST else 6 - assert ( - params["broadcast_time"] < time_limit - ), f"broadcast_time exceeds time limit {time_limit}s" + assert params["broadcast_time"] < time_limit, ( + f"broadcast_time exceeds time limit {time_limit}s" + ) - assert ( - params["update_sgl_dp_1_time"] < time_limit - ), f"update_sgl_dp_one_time exceeds time limit {time_limit}s" + assert params["update_sgl_dp_1_time"] < time_limit, ( + f"update_sgl_dp_one_time exceeds time limit {time_limit}s" + ) if dp_size == 2: - assert ( - params["update_sgl_dp_2_time"] < time_limit - ), f"update_sgl_dp_two_time exceeds time limit {time_limit}s" + assert params["update_sgl_dp_2_time"] < time_limit, ( + f"update_sgl_dp_two_time exceeds time limit {time_limit}s" + ) # Delete the context and close the parameter queue. del context @@ -710,7 +710,6 @@ def test_update_weights_from_distributed( class TestUpdateWeightsFromDistributed(CustomTestCase): - def test_update_weights_from_distributed(self): assert torch.cuda.device_count() >= 2, "At least 2 GPUs are required" diff --git a/test/registered/rl/test_update_weights_from_tensor.py b/test/registered/rl/test_update_weights_from_tensor.py index 4512244a7..58d924601 100644 --- a/test/registered/rl/test_update_weights_from_tensor.py +++ b/test/registered/rl/test_update_weights_from_tensor.py @@ -52,9 +52,9 @@ def test_update_weights_from_tensor(tp_size): torch.cuda.ipc_collect() torch.cuda.empty_cache() memory_after = torch.cuda.memory_allocated() - assert ( - memory_after <= memory_before + 1024 - ), f"Memory leak detected: {memory_after - memory_before} bytes" + assert memory_after <= memory_before + 1024, ( + f"Memory leak detected: {memory_after - memory_before} bytes" + ) class TestUpdateWeightsFromTensor(CustomTestCase): @@ -283,9 +283,9 @@ class TestServerUpdateWeightsFromTensorNonBlocking(CustomTestCase): def _check_param(engine, param_name, expect_values): actual_values = torch.tensor(engine.get_weights_by_name(param_name))[0, :5] - assert torch.allclose( - actual_values, torch.tensor(expect_values), atol=0.002 - ), f"{actual_values=}" + assert torch.allclose(actual_values, torch.tensor(expect_values), atol=0.002), ( + f"{actual_values=}" + ) if __name__ == "__main__": diff --git a/test/registered/rust/test_rust_extension.py b/test/registered/rust/test_rust_extension.py index 2b93cf714..60ec85345 100644 --- a/test/registered/rust/test_rust_extension.py +++ b/test/registered/rust/test_rust_extension.py @@ -59,7 +59,8 @@ features = ["python"] [lib] name = "demo_extension" crate-type = ["cdylib"] -""".strip() + "\n", +""".strip() + + "\n", encoding="utf-8", ) (crate / "lib.rs").write_text("fn input() {}\n", encoding="utf-8") diff --git a/test/registered/scheduler/test_min_free_slots_delayer.py b/test/registered/scheduler/test_min_free_slots_delayer.py index 44f2f8c49..960469471 100644 --- a/test/registered/scheduler/test_min_free_slots_delayer.py +++ b/test/registered/scheduler/test_min_free_slots_delayer.py @@ -10,7 +10,6 @@ register_cpu_ci(est_time=5, suite="base-a-test-cpu") class TestResolveMinFreeSlots(unittest.TestCase): - def test_unset_non_dflash_disables(self): self.assertIsNone(resolve_min_free_slots(None, 512, is_dflash_family=False)) @@ -50,7 +49,6 @@ class TestResolveMinFreeSlots(unittest.TestCase): class TestMinFreeSlotsDelayer(unittest.TestCase): - def test_delays_below_threshold(self): delayer = MinFreeSlotsDelayer(min_free_slots=4) self.assertTrue(delayer.should_delay(running_bs=100, num_allocatable_reqs=2)) diff --git a/test/registered/scheduler/test_prefill_delayer.py b/test/registered/scheduler/test_prefill_delayer.py index b5f89c3d8..2737f30f1 100644 --- a/test/registered/scheduler/test_prefill_delayer.py +++ b/test/registered/scheduler/test_prefill_delayer.py @@ -126,9 +126,9 @@ def _run_negotiate_test(rank, test_cases): ) # On a release after a real wait, seconds must be observed too. if case.expected_wait_forward_passes > 0: - assert ( - result.wait_seconds > 0.0 - ), f"Case {case.name} rank {rank}: wait_seconds not surfaced" + assert result.wait_seconds > 0.0, ( + f"Case {case.name} rank {rank}: wait_seconds not surfaced" + ) override.restore() diff --git a/test/registered/scripted_runtime/test_scripted_runtime_core.py b/test/registered/scripted_runtime/test_scripted_runtime_core.py index b30521973..5e731868d 100644 --- a/test/registered/scripted_runtime/test_scripted_runtime_core.py +++ b/test/registered/scripted_runtime/test_scripted_runtime_core.py @@ -55,9 +55,9 @@ class TestScriptedRuntimeCore(ScriptedTestCase): ) assert r.rid == "explicit-rid-test", f"explicit rid not honored: {r.rid!r}" yield - assert ( - t.find_req_by_rid("explicit-rid-test") is not None - ), "explicit rid not visible to the scheduler after one step" + assert t.find_req_by_rid("explicit-rid-test") is not None, ( + "explicit rid not visible to the scheduler after one step" + ) def test_find_req_by_rid_hit_and_miss(self): self.server.execute_script(self._script_find_req_by_rid_hit_and_miss) @@ -67,12 +67,12 @@ class TestScriptedRuntimeCore(ScriptedTestCase): r = t.start_req(prompt_len=_SHORT_PROMPT_LEN, max_new_tokens=4) yield found = t.find_req_by_rid(r.rid) - assert ( - found is not None and found.rid == r.rid - ), f"find_req_by_rid missed the live rid {r.rid!r}" - assert ( - t.find_req_by_rid("no-such-rid") is None - ), "find_req_by_rid returned a req for an unknown rid" + assert found is not None and found.rid == r.rid, ( + f"find_req_by_rid missed the live rid {r.rid!r}" + ) + assert t.find_req_by_rid("no-such-rid") is None, ( + "find_req_by_rid returned a req for an unknown rid" + ) yield from run_until_finished(r) def test_is_finished_reflects_completion(self): @@ -93,9 +93,9 @@ class TestScriptedRuntimeCore(ScriptedTestCase): def _script_req_handle_req_property(t: ScriptedContext): r = t.start_req(prompt_len=_SHORT_PROMPT_LEN, max_new_tokens=4) yield - assert ( - r.req is not None and r.req.rid == r.rid - ), f"handle.req did not resolve to the live req for {r.rid!r}" + assert r.req is not None and r.req.rid == r.rid, ( + f"handle.req did not resolve to the live req for {r.rid!r}" + ) bogus = ScriptedReqHandle(rid="no-such-rid", context=t) assert bogus.req is None, "handle.req returned a req for an unknown rid" yield from run_until_finished(r) @@ -121,12 +121,12 @@ class TestScriptedRuntimeCore(ScriptedTestCase): r = t.start_req(prompt_len=_LONG_PROMPT_LEN, max_new_tokens=2) yield from advance_to_nth_chunk(r, 1) assert r.is_chunking, "handle.is_chunking False during multi-chunk prefill" - assert t.is_chunking( - r.rid - ), "context.is_chunking False during multi-chunk prefill" - assert not t.is_chunking( - "no-such-rid" - ), "context.is_chunking True for an unknown rid" + assert t.is_chunking(r.rid), ( + "context.is_chunking False during multi-chunk prefill" + ) + assert not t.is_chunking("no-such-rid"), ( + "context.is_chunking True for an unknown rid" + ) yield from run_until_finished(r) assert not r.is_chunking, "handle.is_chunking still True after finish" assert not t.is_chunking(r.rid), "context.is_chunking still True after finish" @@ -146,9 +146,9 @@ class TestScriptedRuntimeCore(ScriptedTestCase): yield req = r.req - assert ( - req is not None and req in t.scheduler.waiting_queue - ), f"pause(retract) did not park the req in waiting_queue; found {req!r}" + assert req is not None and req in t.scheduler.waiting_queue, ( + f"pause(retract) did not park the req in waiting_queue; found {req!r}" + ) frozen = len(req.output_ids) for _ in range(3): @@ -180,9 +180,9 @@ class TestScriptedRuntimeCore(ScriptedTestCase): yield req = r.req - assert ( - req is not None and req not in t.scheduler.waiting_queue - ), f"pause(in_place) should not retract the req to waiting_queue; found {req!r}" + assert req is not None and req not in t.scheduler.waiting_queue, ( + f"pause(in_place) should not retract the req to waiting_queue; found {req!r}" + ) for _ in range(3): yield @@ -242,9 +242,9 @@ class TestScriptedRuntimeCore(ScriptedTestCase): t.flush_cache() yield - assert ( - not t.get_all_node_hit_counts() - ), "flush_cache did not clear the radix tree" + assert not t.get_all_node_hit_counts(), ( + "flush_cache did not clear the radix tree" + ) def test_get_all_node_hit_counts_increments_on_cache_hit(self): self.server.execute_script(self._script_hit_counts_increment_on_cache_hit) @@ -283,17 +283,17 @@ class TestScriptedRuntimeCore(ScriptedTestCase): ) yield from advance_to_decode_step(r, 1) lock_refs = t.get_all_node_lock_refs() - assert ( - lock_refs and max(lock_refs.values()) >= 1 - ), f"expected a locked radix node while the req runs; got {lock_refs}" + assert lock_refs and max(lock_refs.values()) >= 1, ( + f"expected a locked radix node while the req runs; got {lock_refs}" + ) yield from run_until_finished(r) for _ in range(3): yield released = t.get_all_node_lock_refs() - assert released and all( - ref == 0 for ref in released.values() - ), f"radix nodes still locked after the req finished: {released}" + assert released and all(ref == 0 for ref in released.values()), ( + f"radix nodes still locked after the req finished: {released}" + ) def test_start_req_ignore_eos_runs_full_length(self): self.server.execute_script(self._script_ignore_eos_runs_full_length) @@ -304,9 +304,9 @@ class TestScriptedRuntimeCore(ScriptedTestCase): yield from run_until_finished(r) req = r.req assert req is not None, "finished req vanished before its output could be read" - assert ( - len(req.output_ids) == 6 - ), f"ignore_eos must decode the full length; got {list(req.output_ids)!r}" + assert len(req.output_ids) == 6, ( + f"ignore_eos must decode the full length; got {list(req.output_ids)!r}" + ) def test_start_req_priority_is_propagated(self): self.server.execute_script(self._script_priority_is_propagated) @@ -392,9 +392,9 @@ class TestScriptedRuntimeCore(ScriptedTestCase): yield from run_until_finished(r) for _ in range(5): yield - assert ( - t.list_active_reqs() == [] - ), "active reqs should drain to empty after finish" + assert t.list_active_reqs() == [], ( + "active reqs should drain to empty after finish" + ) def test_kv_pages_held_during_run_released_after(self): self.server.execute_script(self._script_kv_pages_set_then_released) @@ -438,9 +438,9 @@ class TestScriptedRuntimeCore(ScriptedTestCase): t.flush_cache() yield after_free = t.engine_stats()["kv_pool_free"] - assert ( - after_free >= during_free - ), f"kv_pool_free should recover after finish; during={during_free} after={after_free}" + assert after_free >= during_free, ( + f"kv_pool_free should recover after finish; during={during_free} after={after_free}" + ) def test_lock_refs_held_during_run_released_after(self): self.server.execute_script(self._script_lock_refs_held_then_released) @@ -455,13 +455,13 @@ class TestScriptedRuntimeCore(ScriptedTestCase): ignore_eos=True, ) yield from advance_to_decode_step(r, 1) - assert ( - r.lock_refs >= 1 - ), f"radix lock_ref must be held mid-run; got {r.lock_refs}" + assert r.lock_refs >= 1, ( + f"radix lock_ref must be held mid-run; got {r.lock_refs}" + ) yield from run_until_finished(r) - assert ( - r.lock_refs == 0 - ), f"lock_refs must be released after finish; got {r.lock_refs}" + assert r.lock_refs == 0, ( + f"lock_refs must be released after finish; got {r.lock_refs}" + ) def test_batch_composition_shape_and_disjoint(self): self.server.execute_script(self._script_batch_composition) @@ -477,9 +477,9 @@ class TestScriptedRuntimeCore(ScriptedTestCase): "chunked", "running", }, f"unexpected batch_composition keys: {comp!r}" - assert ( - r.rid in comp["chunked"] - ), f"chunked req must be in 'chunked'; got {comp!r}" + assert r.rid in comp["chunked"], ( + f"chunked req must be in 'chunked'; got {comp!r}" + ) prefill, decode, chunked = ( set(comp["prefill"]), set(comp["decode"]), @@ -491,9 +491,9 @@ class TestScriptedRuntimeCore(ScriptedTestCase): and decode.isdisjoint(chunked) ), f"prefill/decode/chunked subsets must be disjoint; got {comp!r}" yield from run_until_finished(r) - assert ( - t.batch_composition()["chunked"] == [] - ), "no chunked req should remain after the req finishes" + assert t.batch_composition()["chunked"] == [], ( + "no chunked req should remain after the req finishes" + ) def test_chunks_done_zero_for_unchunked_prompt(self): self.server.execute_script(self._script_chunks_done_zero) @@ -502,9 +502,9 @@ class TestScriptedRuntimeCore(ScriptedTestCase): def _script_chunks_done_zero(t: ScriptedContext): r = t.start_req(prompt_len=_SHORT_PROMPT_LEN, max_new_tokens=2, ignore_eos=True) yield from run_until_finished(r) - assert ( - r.chunks_done == 0 - ), f"prompt <= chunk must not chunk; got {r.chunks_done}" + assert r.chunks_done == 0, ( + f"prompt <= chunk must not chunk; got {r.chunks_done}" + ) def test_chunks_done_counts_two_chunks(self): self.server.execute_script(self._script_chunks_done_two) @@ -513,9 +513,9 @@ class TestScriptedRuntimeCore(ScriptedTestCase): def _script_chunks_done_two(t: ScriptedContext): r = t.start_req(prompt_len=_CHUNK_SIZE + 2, max_new_tokens=2, ignore_eos=True) yield from run_until_finished(r) - assert ( - r.chunks_done == 2 - ), f"chunk_size+2 prompt -> 2 chunks; got {r.chunks_done}" + assert r.chunks_done == 2, ( + f"chunk_size+2 prompt -> 2 chunks; got {r.chunks_done}" + ) def test_chunks_done_scales_with_prompt(self): self.server.execute_script(self._script_chunks_done_five) @@ -524,9 +524,9 @@ class TestScriptedRuntimeCore(ScriptedTestCase): def _script_chunks_done_five(t: ScriptedContext): r = t.start_req(prompt_len=5 * _CHUNK_SIZE, max_new_tokens=2, ignore_eos=True) yield from run_until_finished(r) - assert ( - r.chunks_done == 5 - ), f"5*chunk_size prompt -> 5 chunks; got {r.chunks_done}" + assert r.chunks_done == 5, ( + f"5*chunk_size prompt -> 5 chunks; got {r.chunks_done}" + ) def test_is_idle_reflects_engine_activity(self): self.server.execute_script(self._script_is_idle_reflects_activity) @@ -552,22 +552,22 @@ class TestScriptedRuntimeCore(ScriptedTestCase): @staticmethod def _script_status_transitions(t: ScriptedContext): - assert ( - t.status("no-such-rid") == "unknown" - ), "status of a never-seen rid must be 'unknown'" + assert t.status("no-such-rid") == "unknown", ( + "status of a never-seen rid must be 'unknown'" + ) r = t.start_req( prompt_len=_SHORT_PROMPT_LEN, max_new_tokens=_DECODE_MAX_NEW_TOKENS, ignore_eos=True, ) yield from advance_to_decode_step(r, 1) - assert ( - r.status == "running" - ), f"decoding req status should be running; got {r.status!r}" + assert r.status == "running", ( + f"decoding req status should be running; got {r.status!r}" + ) yield from run_until_finished(r) - assert ( - r.status == "finished" - ), f"completed req status should be finished; got {r.status!r}" + assert r.status == "finished", ( + f"completed req status should be finished; got {r.status!r}" + ) def test_last_batch_forward_mode_extend_then_decode(self): self.server.execute_script(self._script_last_batch_forward_mode) @@ -585,9 +585,9 @@ class TestScriptedRuntimeCore(ScriptedTestCase): f"got {t.last_batch_forward_mode!r}" ) yield from advance_to_decode_step(r, 1) - assert ( - t.last_batch_forward_mode == "DECODE" - ), f"decode batch mode should be DECODE; got {t.last_batch_forward_mode!r}" + assert t.last_batch_forward_mode == "DECODE", ( + f"decode batch mode should be DECODE; got {t.last_batch_forward_mode!r}" + ) yield from run_until_finished(r) def test_remaining_prompt_tokens_shrinks_to_zero(self): @@ -603,9 +603,9 @@ class TestScriptedRuntimeCore(ScriptedTestCase): f"got {rem} (prompt={_LONG_PROMPT_LEN})" ) yield from run_until_finished(r) - assert ( - r.remaining_prompt_tokens == 0 - ), f"finished req should have 0 remaining; got {r.remaining_prompt_tokens}" + assert r.remaining_prompt_tokens == 0, ( + f"finished req should have 0 remaining; got {r.remaining_prompt_tokens}" + ) def test_evict_radix_full_clears_tree_and_rejects_prefix(self): self.server.execute_script(self._script_evict_radix) @@ -620,18 +620,18 @@ class TestScriptedRuntimeCore(ScriptedTestCase): t.evict_radix(prefix_tokens=None) yield - assert ( - not t.get_all_node_hit_counts() - ), "evict_radix(prefix_tokens=None) did not clear the radix tree" + assert not t.get_all_node_hit_counts(), ( + "evict_radix(prefix_tokens=None) did not clear the radix tree" + ) rejected = False try: t.evict_radix(prefix_tokens=[1, 2, 3]) except AssertionError: rejected = True - assert ( - rejected - ), "evict_radix must reject a non-None prefix (only full evict supported)" + assert rejected, ( + "evict_radix must reject a non-None prefix (only full evict supported)" + ) def test_warmup_radix_populates_prefix(self): self.server.execute_script(self._script_warmup_radix_populates_prefix) @@ -644,9 +644,9 @@ class TestScriptedRuntimeCore(ScriptedTestCase): r = t.start_req(prompt_len=2 * _CHUNK_SIZE + 1, max_new_tokens=1) yield from run_until_finished(r) - assert ( - r.req is not None - ), "finished req vanished before cached_tokens could be read" + assert r.req is not None, ( + "finished req vanished before cached_tokens could be read" + ) assert r.req.cached_tokens > 0, ( f"req with the warmed prefix should hit the radix cache; " f"got cached_tokens={r.req.cached_tokens}" @@ -660,24 +660,24 @@ class TestScriptedRuntimeCore(ScriptedTestCase): stats = t.engine_stats() page = stats["page_size"] baseline = stats["kv_pool_free"] - assert ( - baseline > 4 * page - ), f"need KV headroom to test pressure; baseline={baseline}" + assert baseline > 4 * page, ( + f"need KV headroom to test pressure; baseline={baseline}" + ) t.exhaust_kv(leave_pages=2) pressured = t.engine_stats()["kv_pool_free"] - assert ( - pressured < baseline - ), f"exhaust_kv must reduce free KV; pressured={pressured} baseline={baseline}" - assert ( - pressured <= 3 * page - ), f"exhaust_kv(leave_pages=2) left too much free KV; got {pressured} (page={page})" + assert pressured < baseline, ( + f"exhaust_kv must reduce free KV; pressured={pressured} baseline={baseline}" + ) + assert pressured <= 3 * page, ( + f"exhaust_kv(leave_pages=2) left too much free KV; got {pressured} (page={page})" + ) t._release_exhausted_pools() restored = t.engine_stats()["kv_pool_free"] - assert ( - restored == baseline - ), f"release must restore the full pool; restored={restored} baseline={baseline}" + assert restored == baseline, ( + f"release must restore the full pool; restored={restored} baseline={baseline}" + ) yield def test_exhaust_row_pool_leaves_requested_free_rows(self): @@ -692,12 +692,12 @@ class TestScriptedRuntimeCore(ScriptedTestCase): yield from exhaust_row_pool(t, leave_rows=target_free) free_after = t.engine_stats()["req_pool_free"] - assert ( - free_after <= target_free - ), f"exhaust_row_pool should leave <= {target_free} free rows; got {free_after}" - assert ( - free_after < avail - ), f"exhaust_row_pool did not consume any rows; avail={avail} after={free_after}" + assert free_after <= target_free, ( + f"exhaust_row_pool should leave <= {target_free} free rows; got {free_after}" + ) + assert free_after < avail, ( + f"exhaust_row_pool did not consume any rows; avail={avail} after={free_after}" + ) def test_forward_ct_advances_once_per_yield(self): self.server.execute_script(self._script_forward_ct_advances_once_per_yield) @@ -710,9 +710,9 @@ class TestScriptedRuntimeCore(ScriptedTestCase): r = t.start_req( prompt_len=_SHORT_PROMPT_LEN, max_new_tokens=128, ignore_eos=True ) - assert ( - sched.forward_ct == before_no_yield - ), f"forward_ct moved without a yield: {before_no_yield} -> {sched.forward_ct}" + assert sched.forward_ct == before_no_yield, ( + f"forward_ct moved without a yield: {before_no_yield} -> {sched.forward_ct}" + ) yield from advance_to_decode_step(r, 1) @@ -781,7 +781,6 @@ class TestScriptedRuntimeCore(ScriptedTestCase): class TestScriptedRuntimeSession(CustomTestCase): - def test_shutdown_is_idempotent(self): session = ScriptedHttpServer.start(**_ENGINE_KWARGS) session.shutdown() diff --git a/test/registered/sessions/test_session_control.py b/test/registered/sessions/test_session_control.py index 38119dd70..d255b6436 100644 --- a/test/registered/sessions/test_session_control.py +++ b/test/registered/sessions/test_session_control.py @@ -293,9 +293,9 @@ class TestSessionControl(CustomTestCase): print(logprobs_from_session) print("logprobs from normal queries:") print(logprobs_normal) - assert len(logprobs_from_session) == len( - logprobs_normal - ), "logprobs must have equal length" + assert len(logprobs_from_session) == len(logprobs_normal), ( + "logprobs must have equal length" + ) for a, b in zip(logprobs_from_session, logprobs_normal): assert abs(a - b) <= 0.15, f"logprobs {a} and {b} differ by more than 0.15" @@ -426,9 +426,9 @@ class TestSessionControl(CustomTestCase): output_no_session = response["text"] print("second request output without session:") print(output_no_session) - assert ( - second_output == output_no_session - ), f"second_output: {second_output}, output_no_session: {output_no_session}" + assert second_output == output_no_session, ( + f"second_output: {second_output}, output_no_session: {output_no_session}" + ) @unittest.skip("broken") def test_session_control_backtrack_with_abort(self): @@ -547,9 +547,9 @@ class TestSessionControl(CustomTestCase): print(outputs_from_session) print("====== outputs from normal queries: =======") print(outputs_normal) - assert ( - outputs_from_session == outputs_normal - ), f"outputs_from_session: {outputs_from_session}, outputs_normal: {outputs_normal}" + assert outputs_from_session == outputs_normal, ( + f"outputs_from_session: {outputs_from_session}, outputs_normal: {outputs_normal}" + ) @unittest.skipIf( is_hip(), @@ -781,9 +781,9 @@ class TestSessionControlVision(CustomTestCase): print(outputs_from_session) print("outputs from normal queries:") print(outputs_normal) - assert ( - outputs_from_session == outputs_normal - ), f"outputs_from_session: {outputs_from_session}, outputs_normal: {outputs_normal}" + assert outputs_from_session == outputs_normal, ( + f"outputs_from_session: {outputs_from_session}, outputs_normal: {outputs_normal}" + ) if __name__ == "__main__": diff --git a/test/registered/spec/dspark/test_dspark_block_accept_estimator.py b/test/registered/spec/dspark/test_dspark_block_accept_estimator.py index 021c7191d..ab11d2de1 100644 --- a/test/registered/spec/dspark/test_dspark_block_accept_estimator.py +++ b/test/registered/spec/dspark/test_dspark_block_accept_estimator.py @@ -37,7 +37,6 @@ def _make_recorder(tmp_dir: str) -> tuple[BlockAcceptEstimateRecorder, Path]: class _FakeDelayed: - def __init__(self): self._pending = None diff --git a/test/registered/spec/dspark/test_dspark_scheduler.py b/test/registered/spec/dspark/test_dspark_scheduler.py index 53715fa1d..a76131ca2 100644 --- a/test/registered/spec/dspark/test_dspark_scheduler.py +++ b/test/registered/spec/dspark/test_dspark_scheduler.py @@ -291,7 +291,6 @@ class TestScheduleVerifyLensTopk(CustomTestCase): class TestVerifyLenAnchorContract(CustomTestCase): - @_for_each_impl def test_explicit_zero_min_still_clamped_to_anchor(self, impl): survival = _survival_from_confidence( diff --git a/test/registered/tokenizer/test_skip_tokenizer_init.py b/test/registered/tokenizer/test_skip_tokenizer_init.py index 4bc255bab..f9861d8fe 100644 --- a/test/registered/tokenizer/test_skip_tokenizer_init.py +++ b/test/registered/tokenizer/test_skip_tokenizer_init.py @@ -101,7 +101,7 @@ class TestSkipTokenizerInit(CustomTestCase): self.assertEqual( len(item["meta_info"]["input_token_logprobs"]), num_input_logprobs, - f'{len(item["meta_info"]["input_token_logprobs"])} mismatch with {len(input_ids)}', + f"{len(item['meta_info']['input_token_logprobs'])} mismatch with {len(input_ids)}", ) self.assertEqual( len(item["meta_info"]["output_token_logprobs"]), @@ -245,9 +245,9 @@ class TestSkipTokenizerInitVLM(TestSkipTokenizerInit): def get_request_json(self, *args, **kwargs): ret = super().get_request_json(*args, **kwargs) ret["image_data"] = [self.image_url] - ret["logprob_start_len"] = ( - -1 - ) # Do not try to calculate logprobs of image embeddings. + ret[ + "logprob_start_len" + ] = -1 # Do not try to calculate logprobs of image embeddings. return ret def test_simple_decode_stream(self): diff --git a/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py b/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py index c3f63c0a3..2077e3b09 100644 --- a/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py +++ b/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py @@ -36,9 +36,12 @@ def _make_target_verify_batch(bs: int) -> ForwardBatch: def _filter(batch: ForwardBatch, *, lo: int, hi: int) -> ForwardBatch: # filter_batch reads attention_backend (get_server_args) and # moe_dense_tp_size (get_parallel) from the published config. - with get_context().override_server_args( - attention_backend="fa3", moe_dense_tp_size=None - ), get_parallel().override(attn_tp_size=1): + with ( + get_context().override_server_args( + attention_backend="fa3", moe_dense_tp_size=None + ), + get_parallel().override(attn_tp_size=1), + ): return TboForwardBatchPreparer.filter_batch( batch, start_token_index=lo, diff --git a/test/registered/unit/bench/test_mmmu_eval_utils.py b/test/registered/unit/bench/test_mmmu_eval_utils.py index d6de7fe04..950d973a1 100644 --- a/test/registered/unit/bench/test_mmmu_eval_utils.py +++ b/test/registered/unit/bench/test_mmmu_eval_utils.py @@ -123,12 +123,7 @@ class TestMMMUEvalUtils(CustomTestCase): self, ): response = ( - "The options are:\n" - "(A) red\n" - "(B) blue\n" - "(C) green\n" - "(D) yellow\n" - "Answer: B" + "The options are:\n(A) red\n(B) blue\n(C) green\n(D) yellow\nAnswer: B" ) pred_ans = self.eval_utils.parse_multi_choice_response( diff --git a/test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py b/test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py index 13d03bc6a..b92d46f54 100644 --- a/test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py +++ b/test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py @@ -176,7 +176,9 @@ class TestMultimodalPiecewiseCudaGraph(CustomTestCase): args.disable_radix_cache = False args.chunked_prefill_size = 2048 - with (override_platform(is_cuda=True),): + with ( + override_platform(is_cuda=True), + ): handle_model_capability_adjustments(args) self.assertTrue(resolution_result(args, "disable_radix_cache")) diff --git a/test/registered/unit/disaggregation/test_decode_queue_cleanup.py b/test/registered/unit/disaggregation/test_decode_queue_cleanup.py index 7c979edce..0aabdbeef 100644 --- a/test/registered/unit/disaggregation/test_decode_queue_cleanup.py +++ b/test/registered/unit/disaggregation/test_decode_queue_cleanup.py @@ -303,8 +303,7 @@ class TestDecodeQueueCleanup(CustomTestCase): tail = decode_req(8) queue.pending_reqs.append(tail) with patch( - "sglang.srt.disaggregation.decode." - "CommonKVReceiver.query_prefill_dp_ranks", + "sglang.srt.disaggregation.decode.CommonKVReceiver.query_prefill_dp_ranks", return_value={"8": 2}, ) as query: queue._resolve_pending_reqs() diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index 7337d86e4..15e6eda1a 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py @@ -445,12 +445,13 @@ class TestEagleDsaSeedTransfer(unittest.TestCase): # positions are passed through unremapped. ("other", False, False, False, unremapped), ): - with self.subTest(platform=platform), envs.SGLANG_DSA_FUSE_TOPK.override( - True - ), patch( - "sglang.srt.layers.attention.dsa.utils.is_cuda", return_value=cuda - ), patch( - "sglang.srt.layers.attention.dsa.utils.is_hip", return_value=hip + with ( + self.subTest(platform=platform), + envs.SGLANG_DSA_FUSE_TOPK.override(True), + patch( + "sglang.srt.layers.attention.dsa.utils.is_cuda", return_value=cuda + ), + patch("sglang.srt.layers.attention.dsa.utils.is_hip", return_value=hip), ): self.assertEqual( should_use_dsa_fused_topk(seed_dsa_topk_from_draft_extend=True), diff --git a/test/registered/unit/disaggregation/test_nixl_backend_basic.py b/test/registered/unit/disaggregation/test_nixl_backend_basic.py index b5f8c8e5c..c40cc8b1e 100644 --- a/test/registered/unit/disaggregation/test_nixl_backend_basic.py +++ b/test/registered/unit/disaggregation/test_nixl_backend_basic.py @@ -1032,8 +1032,8 @@ class TestNixlStaging(CustomTestCase): staging_total_size=4096, ) calls = [] - mgr.send_kvcache_staged = ( - lambda *args, **kwargs: calls.append((args, kwargs)) or "handle" + mgr.send_kvcache_staged = lambda *args, **kwargs: ( + calls.append((args, kwargs)) or "handle" ) handle, deferred = mgr._do_staging_transfer( diff --git a/test/registered/unit/distributed/test_parallel_state.py b/test/registered/unit/distributed/test_parallel_state.py index 8043b3d40..3308b97eb 100644 --- a/test/registered/unit/distributed/test_parallel_state.py +++ b/test/registered/unit/distributed/test_parallel_state.py @@ -123,7 +123,6 @@ def test_parallel_group_construction_tp8_attn_cp2(): patch("torch.distributed.get_rank", return_value=0), patch("torch.distributed.get_backend", return_value="nccl"), ): - # Mock init_model_parallel_group to capture the groups being created created_groups = {} @@ -144,7 +143,6 @@ def test_parallel_group_construction_tp8_attn_cp2(): ), patch.object(parallel_state, "get_world_group") as mock_world_group, ): - # Mock world group mock_world = Mock() mock_world.device_group = Mock() @@ -174,18 +172,18 @@ def test_parallel_group_construction_tp8_attn_cp2(): # Verify ATTN_CP groups attn_cp_groups = created_groups.get("attn_cp", []) - assert ( - len(attn_cp_groups) == 4 - ), f"Expected 4 ATTN_CP groups, got {len(attn_cp_groups)}" + assert len(attn_cp_groups) == 4, ( + f"Expected 4 ATTN_CP groups, got {len(attn_cp_groups)}" + ) expected_attn_cp = [ [0, 4], [1, 5], [2, 6], [3, 7], ] - assert ( - attn_cp_groups == expected_attn_cp - ), f"Wrong ATTN_CP groups: {attn_cp_groups}" + assert attn_cp_groups == expected_attn_cp, ( + f"Wrong ATTN_CP groups: {attn_cp_groups}" + ) print("TP=8, Attn CP=2 group construction verified") @@ -223,7 +221,6 @@ def test_parallel_group_construction_tp8_moe_ep4_cp2(): patch("torch.distributed.get_rank", return_value=0), patch("torch.distributed.get_backend", return_value="nccl"), ): - # Mock init_model_parallel_group to capture the groups being created created_groups = {} @@ -244,7 +241,6 @@ def test_parallel_group_construction_tp8_moe_ep4_cp2(): ), patch.object(parallel_state, "get_world_group") as mock_world_group, ): - # Mock world group mock_world = Mock() mock_world.device_group = Mock() @@ -275,31 +271,31 @@ def test_parallel_group_construction_tp8_moe_ep4_cp2(): # Verify MOE_EP groups moe_ep_groups = created_groups.get("moe_ep", []) - assert ( - len(moe_ep_groups) == 2 - ), f"Expected 2 MOE_EP groups, got {len(moe_ep_groups)}" + assert len(moe_ep_groups) == 2, ( + f"Expected 2 MOE_EP groups, got {len(moe_ep_groups)}" + ) expected_moe_ep = [ [0, 1, 2, 3], [4, 5, 6, 7], ] - assert ( - moe_ep_groups == expected_moe_ep - ), f"Wrong MOE_EP groups: {moe_ep_groups}" + assert moe_ep_groups == expected_moe_ep, ( + f"Wrong MOE_EP groups: {moe_ep_groups}" + ) # Verify MOE_DP groups moe_dp_groups = created_groups.get("moe_dp", []) - assert ( - len(moe_dp_groups) == 4 - ), f"Expected 4 MOE_DP groups, got {len(moe_dp_groups)}" + assert len(moe_dp_groups) == 4, ( + f"Expected 4 MOE_DP groups, got {len(moe_dp_groups)}" + ) expected_moe_dp = [ [0, 4], [1, 5], [2, 6], [3, 7], ] - assert ( - moe_dp_groups == expected_moe_dp - ), f"Wrong MOE_DP groups: {moe_dp_groups}" + assert moe_dp_groups == expected_moe_dp, ( + f"Wrong MOE_DP groups: {moe_dp_groups}" + ) print("TP=8, MoE EP=4, MoE CP=2 group construction verified") diff --git a/test/registered/unit/entrypoints/openai/test_protocol.py b/test/registered/unit/entrypoints/openai/test_protocol.py index dda22cd44..77561a6bd 100644 --- a/test/registered/unit/entrypoints/openai/test_protocol.py +++ b/test/registered/unit/entrypoints/openai/test_protocol.py @@ -369,8 +369,9 @@ class TestChatCompletionRequest(unittest.TestCase): {"reasoning": {"effort": 1.1}}, {"reasoning": {"effort": "1.5"}}, ): - with self.subTest(request_kwargs=request_kwargs), self.assertRaises( - ValidationError + with ( + self.subTest(request_kwargs=request_kwargs), + self.assertRaises(ValidationError), ): ChatCompletionRequest( model="test-model", messages=messages, **request_kwargs diff --git a/test/registered/unit/entrypoints/test_effective_state_surfaces.py b/test/registered/unit/entrypoints/test_effective_state_surfaces.py index c1434d8cc..5e3c2f864 100644 --- a/test/registered/unit/entrypoints/test_effective_state_surfaces.py +++ b/test/registered/unit/entrypoints/test_effective_state_surfaces.py @@ -137,9 +137,9 @@ def _expanded_write_keys(rel: str, tree: ast.AST, call: ast.Call, kw: ast.keywor if isinstance(key, ast.Constant): keys.add(key.value) continue - assert isinstance( - key, ast.Name - ), f"non-literal dict key in a writer expansion at {rel}:{call.lineno}" + assert isinstance(key, ast.Name), ( + f"non-literal dict key in a writer expansion at {rel}:{call.lineno}" + ) bound = loop_bound(key.id) assert bound, ( f"dict key {key.id!r} at {rel}:{call.lineno} is not bound by a " @@ -150,12 +150,12 @@ def _expanded_write_keys(rel: str, tree: ast.AST, call: ast.Call, kw: ast.keywor if isinstance(kw.value, ast.Dict): return dict_keys(kw.value) - assert isinstance( - kw.value, ast.Name - ), f"unresolvable writer expansion at {rel}:{call.lineno}" - assert ( - enclosing is not None - ), f"writer expansion outside any function at {rel}:{call.lineno}" + assert isinstance(kw.value, ast.Name), ( + f"unresolvable writer expansion at {rel}:{call.lineno}" + ) + assert enclosing is not None, ( + f"writer expansion outside any function at {rel}:{call.lineno}" + ) name = kw.value.id if enclosing.args.kwarg is not None and enclosing.args.kwarg.arg == name: return set() @@ -349,7 +349,7 @@ class TestEffectiveStateSurfaces(CustomTestCase): self.assertEqual( missing, {}, - "a serving surface cannot report what it is running: " f"{missing}", + f"a serving surface cannot report what it is running: {missing}", ) def test_every_surface_reports_the_manager_owned_identity(self): diff --git a/test/registered/unit/function_call/test_function_call_parser.py b/test/registered/unit/function_call/test_function_call_parser.py index bc965a2ce..ad7feaf24 100644 --- a/test/registered/unit/function_call/test_function_call_parser.py +++ b/test/registered/unit/function_call/test_function_call_parser.py @@ -373,7 +373,7 @@ class TestInklingDetector(unittest.TestCase): self.assertEqual(result.normal_text, "Here you go.") def test_empty_name_is_allowed_on_the_canonical_path(self): - source = "<|content_invoke_tool_json|>" '{"name":"","args":{}}<|end_message|>' + source = '<|content_invoke_tool_json|>{"name":"","args":{}}<|end_message|>' result = InklingDetector().detect_and_parse(source, self.tools) self.assertEqual(len(result.calls), 1) self.assertEqual(result.calls[0].name, "") @@ -1305,7 +1305,6 @@ class TestLlama32Detector(unittest.TestCase): class TestKimiK2Detector(unittest.TestCase): - def setUp(self): """Set up test tools and detector.""" self.tools = [ @@ -1374,7 +1373,6 @@ class TestKimiK2Detector(unittest.TestCase): result = self.detector.parse_streaming_increment(chunk, self.tools) for tool_call_chunk in result.calls: if tool_call_chunk.tool_index is not None: - while len(tool_calls) <= tool_call_chunk.tool_index: tool_calls.append({"name": "", "parameters": ""}) @@ -1423,7 +1421,6 @@ class TestKimiK2Detector(unittest.TestCase): result = self.detector.parse_streaming_increment(chunk, self.tools) for tool_call_chunk in result.calls: if tool_call_chunk.tool_index is not None: - while len(tool_calls) <= tool_call_chunk.tool_index: tool_calls.append({"name": "", "parameters": ""}) @@ -1458,7 +1455,6 @@ class TestKimiK2Detector(unittest.TestCase): result = self.detector.parse_streaming_increment(chunk, self.tools) for tool_call_chunk in result.calls: if tool_call_chunk.tool_index is not None: - while len(tool_calls) <= tool_call_chunk.tool_index: tool_calls.append({"name": "", "parameters": ""}) @@ -1531,7 +1527,7 @@ class TestDeepSeekV3Detector(unittest.TestCase): "function<|tool▁sep|>", "get_tour", "ist_att", - "ractions\n```" 'json\n{"', + 'ractions\n```json\n{"', 'city": "', 'Beijing"}\n', "```<|tool▁call▁end|>", @@ -1718,9 +1714,9 @@ class TestDeepSeekV32Detector(unittest.TestCase): if call.name: tool_calls_by_index[call.tool_index]["name"] = call.name if call.parameters: - tool_calls_by_index[call.tool_index][ - "parameters" - ] += call.parameters + tool_calls_by_index[call.tool_index]["parameters"] += ( + call.parameters + ) self.assertGreater(num_tool_call_chunks, 8) @@ -1773,9 +1769,9 @@ class TestDeepSeekV32Detector(unittest.TestCase): if call.name: tool_calls_by_index[call.tool_index]["name"] = call.name if call.parameters: - tool_calls_by_index[call.tool_index][ - "parameters" - ] += call.parameters + tool_calls_by_index[call.tool_index]["parameters"] += ( + call.parameters + ) self.assertGreater(num_tool_call_chunks, 8) self.assertEqual(len(tool_calls_by_index), 1) @@ -1866,9 +1862,9 @@ class TestDeepSeekV32Detector(unittest.TestCase): if call.name: tool_calls_by_index[call.tool_index]["name"] = call.name if call.parameters: - tool_calls_by_index[call.tool_index][ - "parameters" - ] += call.parameters + tool_calls_by_index[call.tool_index]["parameters"] += ( + call.parameters + ) # Verify that the no-parameter function was correctly parsed self.assertEqual( @@ -1926,9 +1922,9 @@ class TestDeepSeekV32Detector(unittest.TestCase): if call.name: tool_calls_by_index[call.tool_index]["name"] = call.name if call.parameters: - tool_calls_by_index[call.tool_index][ - "parameters" - ] += call.parameters + tool_calls_by_index[call.tool_index]["parameters"] += ( + call.parameters + ) # Should still parse correctly even with whitespace-only content self.assertEqual( @@ -2128,9 +2124,9 @@ class TestDeepSeekV4Detector(unittest.TestCase): if call.name: tool_calls_by_index[call.tool_index]["name"] = call.name if call.parameters: - tool_calls_by_index[call.tool_index][ - "parameters" - ] += call.parameters + tool_calls_by_index[call.tool_index]["parameters"] += ( + call.parameters + ) self.assertGreater(num_tool_call_chunks, 8) @@ -4500,9 +4496,7 @@ class TestLfm2Detector(unittest.TestCase): def test_reserved_kwarg_with_nested_quote_recovered(self): """A keyword-named parameter holding a nested-quote command needs the rename and requote rewrites to compose.""" - text = ( - "<|tool_call_start|>[search(from='sed -n '1,5p' f.py')]" "<|tool_call_end|>" - ) + text = "<|tool_call_start|>[search(from='sed -n '1,5p' f.py')]<|tool_call_end|>" result = self.detector.detect_and_parse(text, self.tools) self.assertEqual(len(result.calls), 1) @@ -4720,9 +4714,9 @@ function call<|role_sep|> if call.name: tool_calls_by_index[call.tool_index]["name"] = call.name if call.parameters: - tool_calls_by_index[call.tool_index][ - "parameters" - ] += call.parameters + tool_calls_by_index[call.tool_index]["parameters"] += ( + call.parameters + ) self.assertEqual(len(tool_calls_by_index), 1) self.assertEqual(tool_calls_by_index[0]["name"], "manage_user_memory") @@ -4762,9 +4756,9 @@ function call<|role_sep|> if call.name: tool_calls_by_index[call.tool_index]["name"] = call.name if call.parameters: - tool_calls_by_index[call.tool_index][ - "parameters" - ] += call.parameters + tool_calls_by_index[call.tool_index]["parameters"] += ( + call.parameters + ) self.assertEqual(accumulated_text, "I'll help you.") self.assertEqual(len(tool_calls_by_index), 1) @@ -4802,9 +4796,9 @@ function call<|role_sep|> if call.name: tool_calls_by_index[call.tool_index]["name"] = call.name if call.parameters: - tool_calls_by_index[call.tool_index][ - "parameters" - ] += call.parameters + tool_calls_by_index[call.tool_index]["parameters"] += ( + call.parameters + ) self.assertEqual(len(tool_calls_by_index), 1) self.assertEqual(tool_calls_by_index[0]["name"], "manage_user_memory") @@ -4841,9 +4835,9 @@ function call<|role_sep|> if call.name: tool_calls_by_index[call.tool_index]["name"] = call.name if call.parameters: - tool_calls_by_index[call.tool_index][ - "parameters" - ] += call.parameters + tool_calls_by_index[call.tool_index]["parameters"] += ( + call.parameters + ) self.assertEqual(len(tool_calls_by_index), 1) self.assertEqual(tool_calls_by_index[0]["name"], "get_weather") @@ -4878,9 +4872,9 @@ function call<|role_sep|> if call.name: tool_calls_by_index[call.tool_index]["name"] = call.name if call.parameters: - tool_calls_by_index[call.tool_index][ - "parameters" - ] += call.parameters + tool_calls_by_index[call.tool_index]["parameters"] += ( + call.parameters + ) # Should have name but incomplete parameters self.assertEqual(len(tool_calls_by_index), 1) @@ -4916,9 +4910,9 @@ function call<|role_sep|> if call.name: tool_calls_by_index[call.tool_index]["name"] = call.name if call.parameters: - tool_calls_by_index[call.tool_index][ - "parameters" - ] += call.parameters + tool_calls_by_index[call.tool_index]["parameters"] += ( + call.parameters + ) self.assertEqual(accumulated_text, "I'll remember that.") self.assertEqual(len(tool_calls_by_index), 1) @@ -4961,9 +4955,9 @@ function call<|role_sep|> if call.name: tool_calls_by_index[call.tool_index]["name"] = call.name if call.parameters: - tool_calls_by_index[call.tool_index][ - "parameters" - ] += call.parameters + tool_calls_by_index[call.tool_index]["parameters"] += ( + call.parameters + ) self.assertEqual(len(tool_calls_by_index), 1) self.assertEqual(tool_calls_by_index[0]["name"], "get_weather") @@ -5011,9 +5005,9 @@ function call<|role_sep|> if call.name: tool_calls_by_index[call.tool_index]["name"] = call.name if call.parameters: - tool_calls_by_index[call.tool_index][ - "parameters" - ] += call.parameters + tool_calls_by_index[call.tool_index]["parameters"] += ( + call.parameters + ) self.assertEqual(accumulated_text, "I'll help you.") self.assertEqual(len(tool_calls_by_index), 1) @@ -5265,9 +5259,9 @@ class TestQwen25Detector(unittest.TestCase): if call.name: tool_calls_by_index[call.tool_index]["name"] = call.name if call.parameters: - tool_calls_by_index[call.tool_index][ - "parameters" - ] += call.parameters + tool_calls_by_index[call.tool_index]["parameters"] += ( + call.parameters + ) return tool_calls_by_index def test_streaming_multiple_tool_calls(self): @@ -5515,9 +5509,9 @@ class TestGemma4Detector(unittest.TestCase): if call.name: tool_calls_by_index[call.tool_index]["name"] = call.name if call.parameters: - tool_calls_by_index[call.tool_index][ - "parameters" - ] += call.parameters + tool_calls_by_index[call.tool_index]["parameters"] += ( + call.parameters + ) return normal_text, tool_calls_by_index def test_streaming_multiple_tool_calls(self): @@ -5556,9 +5550,9 @@ class TestGemma4Detector(unittest.TestCase): if call.name: tool_calls_by_index[call.tool_index]["name"] = call.name if call.parameters: - tool_calls_by_index[call.tool_index][ - "parameters" - ] += call.parameters + tool_calls_by_index[call.tool_index]["parameters"] += ( + call.parameters + ) self.assertEqual(len(tool_calls_by_index), 2) self.assertEqual(tool_calls_by_index[0]["name"], "get_weather") @@ -5623,9 +5617,9 @@ class TestGemma4Detector(unittest.TestCase): if call.name: tool_calls_by_index[call.tool_index]["name"] = call.name if call.parameters: - tool_calls_by_index[call.tool_index][ - "parameters" - ] += call.parameters + tool_calls_by_index[call.tool_index]["parameters"] += ( + call.parameters + ) self.assertIn("Hello!", normal_text) self.assertIn("Let me also check", normal_text) self.assertEqual(len(tool_calls_by_index), 2) diff --git a/test/registered/unit/function_call/test_hunyuan_detector.py b/test/registered/unit/function_call/test_hunyuan_detector.py index 18799de52..b1a66b177 100644 --- a/test/registered/unit/function_call/test_hunyuan_detector.py +++ b/test/registered/unit/function_call/test_hunyuan_detector.py @@ -351,9 +351,7 @@ class TestHunyuanDetectorStreaming(CustomTestCase): def test_complete_tool_call_single_chunk(self): detector = self._new_detector() text = ( - "" - "get_current_date" - "" + "get_current_date" ) result = detector.parse_streaming_increment(text, self.tools) collected = _collect_streamed_tool_calls(result.calls) diff --git a/test/registered/unit/function_call/test_kimik3_structural_tag.py b/test/registered/unit/function_call/test_kimik3_structural_tag.py index 8b96ecdb5..9bef313ac 100644 --- a/test/registered/unit/function_call/test_kimik3_structural_tag.py +++ b/test/registered/unit/function_call/test_kimik3_structural_tag.py @@ -197,7 +197,7 @@ def test_strict_schema_rejects_invalid_parameters(arguments): def test_required_allows_response_prefix_but_requires_tools(): grammar = _grammar([_tool()], tool_choice="required") - response = "<|open|>response<|sep|>Checking." "<|close|>response<|sep|>" + response = "<|open|>response<|sep|>Checking.<|close|>response<|sep|>" assert _accepts(grammar, response + _tools_section(_valid_weather_call())) assert not _accepts(grammar, response) diff --git a/test/registered/unit/function_call/test_minicpm5_detector.py b/test/registered/unit/function_call/test_minicpm5_detector.py index 8be593668..3a96d5037 100644 --- a/test/registered/unit/function_call/test_minicpm5_detector.py +++ b/test/registered/unit/function_call/test_minicpm5_detector.py @@ -118,7 +118,7 @@ def test_detect_and_parse_cdata_multiline_v3(): def test_unknown_tool_block_preserved_v3(): detector = MiniCPM5Detector() tools = make_tools_weather() - text = '' '1' "\n" + text = '1\n' res = detector.detect_and_parse(text, tools) assert len(res.calls) == 0 assert "unknown" in res.normal_text @@ -168,7 +168,7 @@ def test_multiple_calls_interleaved_text_v3(): def test_incomplete_missing_function_end_v3(): detector = MiniCPM5Detector() tools = make_tools_weather() - text = '' '北京' + text = '北京' res = detector.detect_and_parse(text, tools) assert len(res.calls) == 0 assert "get_weather" in res.normal_text @@ -204,11 +204,7 @@ def test_duplicate_param_names_invalid_v3(): def test_case_sensitive_param_name_invalid_v3(): detector = MiniCPM5Detector() tools = make_tools_weather() - text = ( - '' - '北京' - "\n" - ) + text = '北京\n' res = detector.detect_and_parse(text, tools) assert len(res.calls) == 0 @@ -243,9 +239,7 @@ def test_streaming_increment_v3(): def test_streaming_split_bot_token(): detector = MiniCPM5Detector() tools = make_tools_weather() - text = ( - '' '北京' "" - ) + text = '北京' r1 = detector.parse_streaming_increment("<", tools) assert r1.normal_text == "" @@ -274,9 +268,7 @@ def test_streaming_multiple_complete_blocks_in_one_delta(): def test_malformed_xml_with_unescaped_ampersand_falls_back_to_regex(): detector = MiniCPM5Detector() tools = make_tools_weather() - text = ( - '' 'A & B' "" - ) + text = 'A & B' result = detector.detect_and_parse(text, tools) assert len(result.calls) == 1 diff --git a/test/registered/unit/function_call/test_muse_glimmer_detector.py b/test/registered/unit/function_call/test_muse_glimmer_detector.py index 8e75ceb1b..62ba8d0ae 100644 --- a/test/registered/unit/function_call/test_muse_glimmer_detector.py +++ b/test/registered/unit/function_call/test_muse_glimmer_detector.py @@ -400,8 +400,7 @@ class TestMuseGlimmerDetector(CustomTestCase): it never will. Goes red if the tool parser loses its stream-end flush (``parse_stream_end`` / detector ``finish``).""" raw = ( - " to=self<|message|>r<|eom|>" - "<|start|>assistant to=user<|message|>answer<|st" + " to=self<|message|>r<|eom|><|start|>assistant to=user<|message|>answer<|st" ) for chunk_size in (1, 7, 100): _, content, calls = self.pipeline_stream(raw, chunk_size) diff --git a/test/registered/unit/function_call/test_poolside_v1_detector.py b/test/registered/unit/function_call/test_poolside_v1_detector.py index 23998d298..54ce766b0 100644 --- a/test/registered/unit/function_call/test_poolside_v1_detector.py +++ b/test/registered/unit/function_call/test_poolside_v1_detector.py @@ -154,9 +154,7 @@ class TestPoolsideV1Detector(CustomTestCase): in-flight call, matching the old closing-tag-anchored regex behavior. Without the truncated-call filter in detect_and_parse, streaming-as- primitive surfaced a tool call with parameters="{}" on this input.""" - text = ( - "get_weather\nlocation\n" "San Fr" - ) + text = "get_weather\nlocation\nSan Fr" result = self.detector.detect_and_parse(text, self.tools) self.assertEqual( len(result.calls), 0, "truncated mid-arg_value must yield 0 calls" diff --git a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py index e58518785..9f561ad58 100644 --- a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py +++ b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py @@ -1181,10 +1181,8 @@ class TestMlxOverlapScheduler(unittest.TestCase): spec_algorithm=SpeculativeAlgorithm.NONE, device="cpu", ) - scheduler.get_next_batch_to_run = ( - lambda running_batch, last_batch: SimpleNamespace( - batch_to_run=batch, running_batch=running_batch - ) + scheduler.get_next_batch_to_run = lambda running_batch, last_batch: ( + SimpleNamespace(batch_to_run=batch, running_batch=running_batch) ) with self.assertRaises(_StopLoop): diff --git a/test/registered/unit/hardware_backend/mlx/test_attn_dp_request_capacity.py b/test/registered/unit/hardware_backend/mlx/test_attn_dp_request_capacity.py index 4ee1ede88..46f2a817b 100644 --- a/test/registered/unit/hardware_backend/mlx/test_attn_dp_request_capacity.py +++ b/test/registered/unit/hardware_backend/mlx/test_attn_dp_request_capacity.py @@ -127,7 +127,7 @@ class TestAttentionDpRequestCapacity(CustomTestCase): with self.assertRaisesRegex( RuntimeError, - "max_mamba_cache_size=15.*per-worker auxiliary-state cap=3.*" "at least 16", + "max_mamba_cache_size=15.*per-worker auxiliary-state cap=3.*at least 16", ): _initialize_stub(stub, hybrid=True) diff --git a/test/registered/unit/hardware_backend/mlx/test_fused_swiglu.py b/test/registered/unit/hardware_backend/mlx/test_fused_swiglu.py index 02592544b..9eee6b5df 100644 --- a/test/registered/unit/hardware_backend/mlx/test_fused_swiglu.py +++ b/test/registered/unit/hardware_backend/mlx/test_fused_swiglu.py @@ -251,9 +251,9 @@ def test_fused_matches_unfused_synthetic(): assert y_ref.shape == y_fused.shape # A broken kernel must not leak NaN/Inf into the downstream down_proj matmul. - assert bool( - mx.all(mx.isfinite(y_fused.astype(mx.float32))).item() - ), f"B={B} hi={hi}: non-finite fused output" + assert bool(mx.all(mx.isfinite(y_fused.astype(mx.float32))).item()), ( + f"B={B} hi={hi}: non-finite fused output" + ) # Same bf16 bound as the @requires_model kernel test. max_abs, rel = _max_rel_diff(y_ref, y_fused) assert rel < 2e-2, f"B={B} hi={hi}: max_abs={max_abs:.3e} rel={rel:.2%}" diff --git a/test/registered/unit/hardware_backend/mlx/test_max_running_requests.py b/test/registered/unit/hardware_backend/mlx/test_max_running_requests.py index e346ee22d..3a16d3a54 100644 --- a/test/registered/unit/hardware_backend/mlx/test_max_running_requests.py +++ b/test/registered/unit/hardware_backend/mlx/test_max_running_requests.py @@ -226,8 +226,10 @@ class TestMlxHybridInitializeAllocation(CustomTestCase): stub = _hybrid_stub_for_initialize( max_running_requests=4, max_mamba_cache_size=2 ) - with _arch(hybrid=True), _published(stub), self.assertRaisesRegex( - RuntimeError, "max_mamba_cache_size" + with ( + _arch(hybrid=True), + _published(stub), + self.assertRaisesRegex(RuntimeError, "max_mamba_cache_size"), ): stub.initialize() diff --git a/test/registered/unit/hardware_backend/mlx/test_metal_profiler.py b/test/registered/unit/hardware_backend/mlx/test_metal_profiler.py index fa50e3bc4..90bb7a29c 100644 --- a/test/registered/unit/hardware_backend/mlx/test_metal_profiler.py +++ b/test/registered/unit/hardware_backend/mlx/test_metal_profiler.py @@ -96,8 +96,9 @@ class TestMetalCaptureProfilerMLX(unittest.TestCase): with tempfile.TemporaryDirectory() as tmp: trace_path = Path(tmp) / "test.gputrace" - with patch.object(mx.metal, "start_capture"), patch.object( - mx.metal, "stop_capture" + with ( + patch.object(mx.metal, "start_capture"), + patch.object(mx.metal, "stop_capture"), ): profiler, result = MetalCaptureProfiler.start_mlx(trace_path) @@ -131,9 +132,10 @@ class TestMetalCaptureProfilerMLX(unittest.TestCase): with tempfile.TemporaryDirectory() as tmp: trace_path = Path(tmp) / "test.gputrace" - with patch.object(mx.metal, "start_capture"), patch.object( - mx.metal, "stop_capture" - ) as mock_stop: + with ( + patch.object(mx.metal, "start_capture"), + patch.object(mx.metal, "stop_capture") as mock_stop, + ): profiler, _ = MetalCaptureProfiler.start_mlx(trace_path) profiler.stop() mock_stop.assert_called_once() @@ -261,9 +263,12 @@ class TestSchedulerProfilerManagerMPS(unittest.TestCase): with tempfile.TemporaryDirectory() as tmp: mgr = self._make_manager(tmp) capture_ctx = MagicMock() - with mock_patch.object( - torch.mps.profiler, "metal_capture", return_value=capture_ctx - ), mock_patch("torch.distributed.barrier"): + with ( + mock_patch.object( + torch.mps.profiler, "metal_capture", return_value=capture_ctx + ), + mock_patch("torch.distributed.barrier"), + ): result = mgr._start_profile() self.assertTrue(result.success, result.message) self.assertTrue(mgr.profile_in_progress) diff --git a/test/registered/unit/hardware_backend/mlx/test_muse_glimmer_mlx_model.py b/test/registered/unit/hardware_backend/mlx/test_muse_glimmer_mlx_model.py index 35dd77bb0..120c7debc 100644 --- a/test/registered/unit/hardware_backend/mlx/test_muse_glimmer_mlx_model.py +++ b/test/registered/unit/hardware_backend/mlx/test_muse_glimmer_mlx_model.py @@ -98,14 +98,13 @@ def _raw_weights(args): p + "post_attn_norm.weight": mx.full((hid,), 0.20), p + "post_attention_layernorm.weight": mx.full((hid,), 0.30), p + "post_ffn_norm.weight": mx.full((hid,), 0.40), - p - + "mlp.gate_proj.weight": mx.random.normal( + p + "mlp.gate_proj.weight": mx.random.normal( (args.intermediate_size, hid) ), - p - + "mlp.up_proj.weight": mx.random.normal((args.intermediate_size, hid)), - p - + "mlp.down_proj.weight": mx.random.normal( + p + "mlp.up_proj.weight": mx.random.normal( + (args.intermediate_size, hid) + ), + p + "mlp.down_proj.weight": mx.random.normal( (hid, args.intermediate_size) ), } diff --git a/test/registered/unit/hardware_backend/mlx/test_quantization.py b/test/registered/unit/hardware_backend/mlx/test_quantization.py index 267fb6ccd..387d9aa56 100644 --- a/test/registered/unit/hardware_backend/mlx/test_quantization.py +++ b/test/registered/unit/hardware_backend/mlx/test_quantization.py @@ -119,8 +119,8 @@ class TestMlxQuantization(unittest.TestCase): self.assertGreater( reduction, 0.40, - f"expected >40% memory reduction with mlx_q4, got {reduction*100:.1f}% " - f"(fp16={mem_fp/1024**3:.2f} GB, q4={mem_q4/1024**3:.2f} GB)", + f"expected >40% memory reduction with mlx_q4, got {reduction * 100:.1f}% " + f"(fp16={mem_fp / 1024**3:.2f} GB, q4={mem_q4 / 1024**3:.2f} GB)", ) def test_mlx_q8_creates_quantized_linear_modules(self): diff --git a/test/registered/unit/layers/attention/test_dsa_mqa_logits_chunking.py b/test/registered/unit/layers/attention/test_dsa_mqa_logits_chunking.py index fdb2375d6..c6bc7b295 100644 --- a/test/registered/unit/layers/attention/test_dsa_mqa_logits_chunking.py +++ b/test/registered/unit/layers/attention/test_dsa_mqa_logits_chunking.py @@ -27,10 +27,13 @@ HUGE_MEM_BUDGET = 64 * 2**30 def _decide(num_q, num_k, mem_budget=HUGE_MEM_BUDGET, is_hip=True): # __new__ skips an __init__ that needs a model config and a device. indexer = dsa_indexer.Indexer.__new__(dsa_indexer.Indexer) - with mock.patch.object(dsa_indexer, "_is_hip", is_hip), mock.patch.object( - dsa_indexer.Indexer, - "_get_mqa_logits_budget_bytes", - return_value=mem_budget, + with ( + mock.patch.object(dsa_indexer, "_is_hip", is_hip), + mock.patch.object( + dsa_indexer.Indexer, + "_get_mqa_logits_budget_bytes", + return_value=mem_budget, + ), ): return indexer._should_chunk_mqa_logits(num_q, num_k, 0) diff --git a/test/registered/unit/layers/attention/test_kda_helion_dispatcher.py b/test/registered/unit/layers/attention/test_kda_helion_dispatcher.py index dc4cd87e6..bc0593f75 100644 --- a/test/registered/unit/layers/attention/test_kda_helion_dispatcher.py +++ b/test/registered/unit/layers/attention/test_kda_helion_dispatcher.py @@ -24,8 +24,7 @@ class TestHelionKDADispatcher(unittest.TestCase): return_value=True, ), patch( - "sglang.srt.layers.attention.linear.kernels.kda_helion." - "HelionKDAKernel", + "sglang.srt.layers.attention.linear.kernels.kda_helion.HelionKDAKernel", return_value=helion_kernel, ) as constructor, ): diff --git a/test/registered/unit/layers/attention/test_vision_strided_qkv.py b/test/registered/unit/layers/attention/test_vision_strided_qkv.py index 59d4453ed..d598686c0 100644 --- a/test/registered/unit/layers/attention/test_vision_strided_qkv.py +++ b/test/registered/unit/layers/attention/test_vision_strided_qkv.py @@ -67,9 +67,10 @@ def single_rank(monkeypatch, gloo_world): "_determine_attention_backend", lambda self, passed_backend: passed_backend, ) - with get_parallel().override( - tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0 - ), get_context().override_server_args(): + with ( + get_parallel().override(tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0), + get_context().override_server_args(), + ): yield diff --git a/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py b/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py index f06106530..6585db94b 100644 --- a/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py +++ b/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py @@ -106,8 +106,8 @@ class TestDeepGemmMegaMoeApi(CustomTestCase): deep_gemm = ModuleType("deep_gemm") deep_gemm.transform_sf_into_required_layout = MagicMock( - side_effect=lambda _sf, mn, k, recipe, num_groups, disable_ue8m0_cast: torch.zeros( - (num_groups, mn, max(1, k // 32)), dtype=torch.int32 + side_effect=lambda _sf, mn, k, recipe, num_groups, disable_ue8m0_cast: ( + torch.zeros((num_groups, mn, max(1, k // 32)), dtype=torch.int32) ) ) deep_gemm.transform_weights_for_mega_moe = MagicMock( diff --git a/test/registered/unit/layers/moe/test_moe_runner_extensions.py b/test/registered/unit/layers/moe/test_moe_runner_extensions.py index 83230a3bc..028adddc0 100644 --- a/test/registered/unit/layers/moe/test_moe_runner_extensions.py +++ b/test/registered/unit/layers/moe/test_moe_runner_extensions.py @@ -150,18 +150,20 @@ def test_fused_moe_uses_explicit_quant_method_for_full_lifecycle(monkeypatch) -> lambda config: SimpleNamespace(), ) - with get_context().override_server_args( - model_path="dummy" - ), get_flags().moe.override( - runner_backend=MoeRunnerBackend.AUTO, - a2a_backend=MoeA2ABackend.NONE, - ), get_parallel().override( - moe_ep_size=1, - moe_ep_rank=0, - moe_tp_size=1, - moe_tp_rank=0, - tp_size=1, - tp_rank=0, + with ( + get_context().override_server_args(model_path="dummy"), + get_flags().moe.override( + runner_backend=MoeRunnerBackend.AUTO, + a2a_backend=MoeA2ABackend.NONE, + ), + get_parallel().override( + moe_ep_size=1, + moe_ep_rank=0, + moe_tp_size=1, + moe_tp_rank=0, + tp_size=1, + tp_rank=0, + ), ): layer = FusedMoE( num_experts=2, @@ -206,10 +208,9 @@ def test_lora_uses_quant_method_contract_for_registered_backend( monkeypatch.setattr( runner_module, "MoeRunner", - lambda selected_backend, config, lora_enabled: created_runners.append( - (selected_backend, config, lora_enabled) - ) - or object(), + lambda selected_backend, config, lora_enabled: ( + created_runners.append((selected_backend, config, lora_enabled)) or object() + ), ) wrapper = FusedMoEWithLoRA(base_layer, lora_backend) @@ -348,18 +349,20 @@ def test_fused_moe_layer_runner_is_none_when_method_builds_no_runner( lambda config: SimpleNamespace(), ) - with get_context().override_server_args( - model_path="dummy" - ), get_flags().moe.override( - runner_backend=MoeRunnerBackend.AUTO, - a2a_backend=MoeA2ABackend.NONE, - ), get_parallel().override( - moe_ep_size=1, - moe_ep_rank=0, - moe_tp_size=1, - moe_tp_rank=0, - tp_size=1, - tp_rank=0, + with ( + get_context().override_server_args(model_path="dummy"), + get_flags().moe.override( + runner_backend=MoeRunnerBackend.AUTO, + a2a_backend=MoeA2ABackend.NONE, + ), + get_parallel().override( + moe_ep_size=1, + moe_ep_rank=0, + moe_tp_size=1, + moe_tp_rank=0, + tp_size=1, + tp_rank=0, + ), ): layer = FusedMoE( num_experts=2, diff --git a/test/registered/unit/layers/quantization/test_deepgemm_ue8m0_requant.py b/test/registered/unit/layers/quantization/test_deepgemm_ue8m0_requant.py index e96b68a84..452c74996 100644 --- a/test/registered/unit/layers/quantization/test_deepgemm_ue8m0_requant.py +++ b/test/registered/unit/layers/quantization/test_deepgemm_ue8m0_requant.py @@ -39,9 +39,10 @@ class TestDeepGemmUE8M0Requant(CustomTestCase): def test_helper_requants_supported_deepgemm_bf16_once(self): weight, weight_scale = _make_params() - with self._enabled_deepgemm_ue8m0(), patch.object( - fp8_utils, "requant_weight_ue8m0_inplace" - ) as requant: + with ( + self._enabled_deepgemm_ue8m0(), + patch.object(fp8_utils, "requant_weight_ue8m0_inplace") as requant, + ): fired = fp8_utils.requant_block_scale_ue8m0_for_deepgemm( weight, weight_scale, @@ -67,9 +68,10 @@ class TestDeepGemmUE8M0Requant(CustomTestCase): def test_helper_skips_non_bf16_output(self): weight, weight_scale = _make_params() - with self._enabled_deepgemm_ue8m0(), patch.object( - fp8_utils, "requant_weight_ue8m0_inplace" - ) as requant: + with ( + self._enabled_deepgemm_ue8m0(), + patch.object(fp8_utils, "requant_weight_ue8m0_inplace") as requant, + ): fired = fp8_utils.requant_block_scale_ue8m0_for_deepgemm( weight, weight_scale, @@ -86,9 +88,10 @@ class TestDeepGemmUE8M0Requant(CustomTestCase): def test_helper_skips_shape_deepgemm_will_not_run(self): weight, weight_scale = _make_params(n=96, k=128) - with self._enabled_deepgemm_ue8m0(), patch.object( - fp8_utils, "requant_weight_ue8m0_inplace" - ) as requant: + with ( + self._enabled_deepgemm_ue8m0(), + patch.object(fp8_utils, "requant_weight_ue8m0_inplace") as requant, + ): fired = fp8_utils.requant_block_scale_ue8m0_for_deepgemm( weight, weight_scale, @@ -105,9 +108,10 @@ class TestDeepGemmUE8M0Requant(CustomTestCase): def test_helper_skips_non_deepgemm_runner(self): weight, weight_scale = _make_params() - with self._enabled_deepgemm_ue8m0(), patch.object( - fp8_utils, "requant_weight_ue8m0_inplace" - ) as requant: + with ( + self._enabled_deepgemm_ue8m0(), + patch.object(fp8_utils, "requant_weight_ue8m0_inplace") as requant, + ): fired = fp8_utils.requant_block_scale_ue8m0_for_deepgemm( weight, weight_scale, @@ -125,9 +129,10 @@ class TestDeepGemmUE8M0Requant(CustomTestCase): weight, weight_scale = _make_params() unsupported_block_size = [128, 256] - with self._enabled_deepgemm_ue8m0(), patch.object( - fp8_utils, "requant_weight_ue8m0_inplace" - ) as requant: + with ( + self._enabled_deepgemm_ue8m0(), + patch.object(fp8_utils, "requant_weight_ue8m0_inplace") as requant, + ): fired = fp8_utils.requant_block_scale_ue8m0_for_deepgemm( weight, weight_scale, @@ -154,9 +159,10 @@ class TestDeepGemmUE8M0Requant(CustomTestCase): layer.weight, layer.weight_scale = _make_params() layer.orig_dtype = torch.bfloat16 - with self._enabled_deepgemm_ue8m0(), patch.object( - fp8_utils, "requant_weight_ue8m0_inplace" - ) as requant: + with ( + self._enabled_deepgemm_ue8m0(), + patch.object(fp8_utils, "requant_weight_ue8m0_inplace") as requant, + ): scheme.process_weights_after_loading(layer) scheme.process_weights_after_loading(layer) @@ -179,18 +185,22 @@ class TestDeepGemmUE8M0Requant(CustomTestCase): weight_scale.format_ue8m0 = True return True - with patch.multiple( - fp8_quant, - _is_cpu=False, - _is_fp8_fnuz=False, - _use_aiter=False, - ), patch.object( - method, "is_deepgemm_moe_runner_backend_enabled", return_value=True - ), patch.object( - fp8_quant, - "requant_block_scale_ue8m0_for_deepgemm", - side_effect=_mark_ue8m0, - ) as requant: + with ( + patch.multiple( + fp8_quant, + _is_cpu=False, + _is_fp8_fnuz=False, + _use_aiter=False, + ), + patch.object( + method, "is_deepgemm_moe_runner_backend_enabled", return_value=True + ), + patch.object( + fp8_quant, + "requant_block_scale_ue8m0_for_deepgemm", + side_effect=_mark_ue8m0, + ) as requant, + ): method.process_weights_after_loading_block_quant(layer) self.assertEqual( diff --git a/test/registered/unit/layers/quantization/test_flashinfer_trtllm_fp8_fallback.py b/test/registered/unit/layers/quantization/test_flashinfer_trtllm_fp8_fallback.py index 979bb3455..f64b80b62 100644 --- a/test/registered/unit/layers/quantization/test_flashinfer_trtllm_fp8_fallback.py +++ b/test/registered/unit/layers/quantization/test_flashinfer_trtllm_fp8_fallback.py @@ -54,17 +54,16 @@ class TestFlashinferTrtllmFp8Fallback(CustomTestCase): trtllm_spy = MagicMock(return_value=torch.zeros((M, N), dtype=dtype)) quant_spy = MagicMock(return_value=(MagicMock(), MagicMock())) - with patch.object( - fp8_utils, - "_get_flashinfer_groupwise_backend", - return_value="trtllm", - create=True, - ), patch.object( - fp8_utils, "gemm_fp8_nt_groupwise", trtllm_spy, create=True - ), patch.object( - fp8_utils, "triton_w8a8_block_fp8_linear", triton_spy - ), patch.object( - fp8_utils, "sglang_per_token_group_quant_fp8", quant_spy + with ( + patch.object( + fp8_utils, + "_get_flashinfer_groupwise_backend", + return_value="trtllm", + create=True, + ), + patch.object(fp8_utils, "gemm_fp8_nt_groupwise", trtllm_spy, create=True), + patch.object(fp8_utils, "triton_w8a8_block_fp8_linear", triton_spy), + patch.object(fp8_utils, "sglang_per_token_group_quant_fp8", quant_spy), ): fp8_utils.flashinfer_gemm_w8a8_block_fp8_linear_with_fallback( input_2d, weight, BLOCK_SIZE, weight_scale diff --git a/test/registered/unit/layers/quantization/test_mxfp4_flashinfer_activation_prep.py b/test/registered/unit/layers/quantization/test_mxfp4_flashinfer_activation_prep.py index f87b3eb20..5c63c85f9 100644 --- a/test/registered/unit/layers/quantization/test_mxfp4_flashinfer_activation_prep.py +++ b/test/registered/unit/layers/quantization/test_mxfp4_flashinfer_activation_prep.py @@ -25,16 +25,20 @@ class TestMxfp4FlashinferActivationPrep(CustomTestCase): x_quant = torch.empty(3, 64, dtype=torch.float8_e4m3fn) x_scale = torch.arange(6, dtype=torch.uint8).reshape(3, 2) - with patch( - "sglang.srt.layers.moe.route_quant_handoff.take", return_value=None - ) as take, patch( - "sglang.srt.layers.quantization.mxfp4._is_sm107_supported", - return_value=True, - ), patch( - "sglang.srt.layers.quantization.fp8_utils.flashinfer_mxfp8_quantize", - return_value=(x_quant, x_scale), - create=True, - ) as quantize: + with ( + patch( + "sglang.srt.layers.moe.route_quant_handoff.take", return_value=None + ) as take, + patch( + "sglang.srt.layers.quantization.mxfp4._is_sm107_supported", + return_value=True, + ), + patch( + "sglang.srt.layers.quantization.fp8_utils.flashinfer_mxfp8_quantize", + return_value=(x_quant, x_scale), + create=True, + ) as quantize, + ): actual_x, packed_topk, actual_quant, actual_scale = ( _prepare_flashinfer_mxfp8_activations(x, 64) ) @@ -51,19 +55,24 @@ class TestMxfp4FlashinferActivationPrep(CustomTestCase): x_quant = torch.empty(3, 64, dtype=torch.float8_e4m3fn) x_scale = torch.arange(6, dtype=torch.uint8).reshape(3, 2) - with patch( - "sglang.srt.layers.moe.route_quant_handoff.take", return_value=None - ) as take, patch( - "sglang.srt.layers.quantization.mxfp4._is_sm107_supported", - return_value=True, - ), patch.object( - per_token_group_quant_module, - "per_token_group_quant", - return_value=(x_quant, x_scale), - ) as quantize, patch( - "sglang.srt.layers.quantization.fp8_utils.flashinfer_mxfp8_quantize", - create=True, - ) as flashinfer_quantize: + with ( + patch( + "sglang.srt.layers.moe.route_quant_handoff.take", return_value=None + ) as take, + patch( + "sglang.srt.layers.quantization.mxfp4._is_sm107_supported", + return_value=True, + ), + patch.object( + per_token_group_quant_module, + "per_token_group_quant", + return_value=(x_quant, x_scale), + ) as quantize, + patch( + "sglang.srt.layers.quantization.fp8_utils.flashinfer_mxfp8_quantize", + create=True, + ) as flashinfer_quantize, + ): actual_x, packed_topk, actual_quant, actual_scale = ( _prepare_flashinfer_mxfp8_activations(x, 64) ) @@ -81,19 +90,22 @@ class TestMxfp4FlashinferActivationPrep(CustomTestCase): x_quant = torch.empty(3, 64, dtype=torch.float8_e4m3fn) x_scale = torch.arange(6, dtype=torch.uint8).reshape(3, 2) - with patch( - "sglang.srt.layers.moe.route_quant_handoff.take", return_value=None - ), patch( - "sglang.srt.layers.quantization.mxfp4._is_sm107_supported", - return_value=False, - ), patch.object( - per_token_group_quant_module, - "per_token_group_quant", - return_value=(x_quant, x_scale), - ) as quantize, patch( - "sglang.srt.layers.quantization.fp8_utils.flashinfer_mxfp8_quantize", - create=True, - ) as flashinfer_quantize: + with ( + patch("sglang.srt.layers.moe.route_quant_handoff.take", return_value=None), + patch( + "sglang.srt.layers.quantization.mxfp4._is_sm107_supported", + return_value=False, + ), + patch.object( + per_token_group_quant_module, + "per_token_group_quant", + return_value=(x_quant, x_scale), + ) as quantize, + patch( + "sglang.srt.layers.quantization.fp8_utils.flashinfer_mxfp8_quantize", + create=True, + ) as flashinfer_quantize, + ): actual_x, packed_topk, actual_quant, actual_scale = ( _prepare_flashinfer_mxfp8_activations(x, 64) ) @@ -111,11 +123,14 @@ class TestMxfp4FlashinferActivationPrep(CustomTestCase): x_quant = torch.empty(3, 128, dtype=torch.float8_e4m3fn) x_scale = torch.arange(12, dtype=torch.uint8) - with patch( - "sglang.srt.layers.quantization.fp8_utils.flashinfer_mxfp8_quantize", - return_value=(x_quant, x_scale), - create=True, - ) as quantize, patch("sglang.srt.layers.moe.route_quant_handoff.take") as take: + with ( + patch( + "sglang.srt.layers.quantization.fp8_utils.flashinfer_mxfp8_quantize", + return_value=(x_quant, x_scale), + create=True, + ) as quantize, + patch("sglang.srt.layers.moe.route_quant_handoff.take") as take, + ): actual_x, packed_topk, actual_quant, actual_scale = ( _prepare_flashinfer_mxfp8_activations(x, 128) ) @@ -133,13 +148,16 @@ class TestMxfp4FlashinferActivationPrep(CustomTestCase): x_quant = torch.empty(2, 64, dtype=torch.float8_e4m3fn) x_scale = torch.arange(4, dtype=torch.uint8).reshape(2, 2) - with patch( - "sglang.srt.layers.moe.route_quant_handoff.take", - return_value=(packed_topk, x_quant, x_scale), - ), patch( - "sglang.srt.layers.quantization.fp8_utils.flashinfer_mxfp8_quantize", - create=True, - ) as quantize: + with ( + patch( + "sglang.srt.layers.moe.route_quant_handoff.take", + return_value=(packed_topk, x_quant, x_scale), + ), + patch( + "sglang.srt.layers.quantization.fp8_utils.flashinfer_mxfp8_quantize", + create=True, + ) as quantize, + ): actual_x, actual_packed, actual_quant, actual_scale = ( _prepare_flashinfer_mxfp8_activations(x, 64) ) diff --git a/test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py b/test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py index 8d440b2ce..f07503795 100644 --- a/test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py +++ b/test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py @@ -585,9 +585,9 @@ def test_humming_range_ignores_prerounded_hidden_tail(): "w13 Humming residual changed with the never-written hidden tail; " "the _UE8M0_ONE fill leaked into the per-expert E8M0 range" ) - assert torch.equal( - residuals[0][1], residuals[1][1] - ), "w2 Humming residual changed with the never-written hidden tail" + assert torch.equal(residuals[0][1], residuals[1][1]), ( + "w2 Humming residual changed with the never-written hidden tail" + ) @pytest.mark.skipif( diff --git a/test/registered/unit/layers/quantization/test_nvfp4_moe_backends.py b/test/registered/unit/layers/quantization/test_nvfp4_moe_backends.py index 6b2164b75..f59998758 100644 --- a/test/registered/unit/layers/quantization/test_nvfp4_moe_backends.py +++ b/test/registered/unit/layers/quantization/test_nvfp4_moe_backends.py @@ -43,17 +43,17 @@ class TestNvFp4MoeBackends(CustomTestCase): quant_config = ModelOptFp4Config( is_checkpoint_nvfp4_serialized=True, group_size=16 ) - with get_context().override_server_args( - model_path="dummy" - ), get_flags().moe.override( - runner_backend=MoeRunnerBackend(backend) - ), get_parallel().override( - moe_ep_size=1, - moe_ep_rank=0, - moe_tp_size=1, - moe_tp_rank=0, - tp_size=1, - tp_rank=0, + with ( + get_context().override_server_args(model_path="dummy"), + get_flags().moe.override(runner_backend=MoeRunnerBackend(backend)), + get_parallel().override( + moe_ep_size=1, + moe_ep_rank=0, + moe_tp_size=1, + moe_tp_rank=0, + tp_size=1, + tp_rank=0, + ), ): layer = FusedMoE( num_experts=E, diff --git a/test/registered/unit/layers/test_conv_layer.py b/test/registered/unit/layers/test_conv_layer.py index d296cdcd4..49d82953e 100644 --- a/test/registered/unit/layers/test_conv_layer.py +++ b/test/registered/unit/layers/test_conv_layer.py @@ -19,7 +19,6 @@ def _copy_weights(src, dst_nn): class TestConv2dLayer(unittest.TestCase): - def test_basic_patch_embedding(self): layer = Conv2dLayer(3, 768, kernel_size=14, stride=14, bias=False) ref = nn.Conv2d(3, 768, kernel_size=14, stride=14, bias=False) @@ -150,7 +149,6 @@ class TestConv2dLayer(unittest.TestCase): class TestConvValidation(unittest.TestCase): - def test_in_channels_not_divisible_by_groups(self): with self.assertRaises(ValueError): Conv2dLayer(3, 64, kernel_size=3, stride=1, groups=2) @@ -204,7 +202,6 @@ class TestConvValidation(unittest.TestCase): class TestConv3dLayer(unittest.TestCase): - def test_basic_temporal_patch_embedding(self): layer = Conv3dLayer( 3, 1152, kernel_size=[2, 14, 14], stride=[2, 14, 14], bias=False diff --git a/test/registered/unit/layers/test_dsv4_kv_splits_heuristic.py b/test/registered/unit/layers/test_dsv4_kv_splits_heuristic.py index 0983e6f91..79997eaf6 100644 --- a/test/registered/unit/layers/test_dsv4_kv_splits_heuristic.py +++ b/test/registered/unit/layers/test_dsv4_kv_splits_heuristic.py @@ -61,9 +61,9 @@ def test_splits_never_increase_with_token_count(): for tokens in [1, 2, 4, 8, 16, 32, 64, 128, 256, 512]: splits = _kv_splits_heuristic(tokens, HEADS, BLOCK_H, num_cu=NUM_CU) if prev is not None: - assert ( - splits <= prev - ), f"splits rose from {prev} to {splits} going to T={tokens}" + assert splits <= prev, ( + f"splits rose from {prev} to {splits} going to T={tokens}" + ) prev = splits diff --git a/test/registered/unit/layers/test_flashinfer_comm_fusion.py b/test/registered/unit/layers/test_flashinfer_comm_fusion.py index 63d607507..337068682 100644 --- a/test/registered/unit/layers/test_flashinfer_comm_fusion.py +++ b/test/registered/unit/layers/test_flashinfer_comm_fusion.py @@ -484,9 +484,12 @@ class TestTagGroupsForFlashInferAllReduceOnly(CustomTestCase): def _tag(self, *, attn_tp, moe_ep, moe_tp): from sglang.srt.distributed import parallel_state as ps - with patch.object(ps, "_ENABLE_FLASHINFER_ALLREDUCE_ONLY", True), patch.object( - ps, "_ATTN_TP", attn_tp - ), patch.object(ps, "_MOE_EP", moe_ep), patch.object(ps, "_MOE_TP", moe_tp): + with ( + patch.object(ps, "_ENABLE_FLASHINFER_ALLREDUCE_ONLY", True), + patch.object(ps, "_ATTN_TP", attn_tp), + patch.object(ps, "_MOE_EP", moe_ep), + patch.object(ps, "_MOE_TP", moe_tp), + ): ps._tag_groups_for_flashinfer_allreduce_only() def test_hybrid_ep_tp_tags_only_the_ep_group(self): diff --git a/test/registered/unit/managers/test_load_snapshot_backends.py b/test/registered/unit/managers/test_load_snapshot_backends.py index bb507bf7e..c3637889d 100644 --- a/test/registered/unit/managers/test_load_snapshot_backends.py +++ b/test/registered/unit/managers/test_load_snapshot_backends.py @@ -192,8 +192,10 @@ class TestZmqRoundTrip(CustomTestCase): loads = _read_until( lambda: reader.read_all(), - lambda snaps: len(snaps) == dp_size - and all(snap.timestamp == 3.0 for snap in snaps), + lambda snaps: ( + len(snaps) == dp_size + and all(snap.timestamp == 3.0 for snap in snaps) + ), ) self.assertEqual(len(loads), dp_size) for load in loads: @@ -373,8 +375,9 @@ class TestZmqReaderOwner(CustomTestCase): override = get_context().override_server_args(**fields) override.install() self.addCleanup(override.restore) - with mock.patch.object(rc, "_ROLE_NS_MODE", "enforce"), mock.patch.object( - rc._CONTEXT, "_publish_role", "dp_controller" + with ( + mock.patch.object(rc, "_ROLE_NS_MODE", "enforce"), + mock.patch.object(rc._CONTEXT, "_publish_role", "dp_controller"), ): self.assertTrue(zmq_reader_owner("DataParallelController")) @@ -442,8 +445,10 @@ class TestEndToEndZmqSimulation(CustomTestCase): loads = _read_until( lambda: reader.read_all(), - lambda snaps: len(snaps) == dp_size - and all(snap.timestamp == 1.0 for snap in snaps), + lambda snaps: ( + len(snaps) == dp_size + and all(snap.timestamp == 1.0 for snap in snaps) + ), ) self.assertEqual(len(loads), dp_size) self.assertEqual(loads[0].num_running_reqs, 10) @@ -464,8 +469,10 @@ class TestEndToEndZmqSimulation(CustomTestCase): loads = _read_until( lambda: reader.read_all(), - lambda snaps: len(snaps) == dp_size - and all(snap.timestamp == 2.0 for snap in snaps), + lambda snaps: ( + len(snaps) == dp_size + and all(snap.timestamp == 2.0 for snap in snaps) + ), ) self.assertEqual(loads[0].num_running_reqs, 20) self.assertEqual(loads[1].num_running_reqs, 21) diff --git a/test/registered/unit/managers/test_loadstat_wire.py b/test/registered/unit/managers/test_loadstat_wire.py index 3c392aeb9..56fc2692a 100644 --- a/test/registered/unit/managers/test_loadstat_wire.py +++ b/test/registered/unit/managers/test_loadstat_wire.py @@ -102,8 +102,7 @@ class TestLoadPublisherGating(CustomTestCase): by default (the feature is off without it). dp_size lives on the ps, which the publisher reads (no separate param to disagree with it).""" with patch( - "sglang.srt.managers.scheduler_components.load_publisher." - "_open_pub_socket" + "sglang.srt.managers.scheduler_components.load_publisher._open_pub_socket" ) as open_sock: pub = SchedulerLoadPublisher( kv_events_config=config, @@ -284,8 +283,7 @@ class TestLoadPublisherGating(CustomTestCase): import zmq with patch( - "sglang.srt.managers.scheduler_components.load_publisher." - "_open_pub_socket", + "sglang.srt.managers.scheduler_components.load_publisher._open_pub_socket", side_effect=zmq.ZMQError, ) as open_sock: pub = SchedulerLoadPublisher( diff --git a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py index 5d2209ca8..397e27666 100644 --- a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py +++ b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py @@ -451,18 +451,21 @@ class TestDecodePrebuilt(unittest.TestCase): scheduler.waiting_queue[0].priority = 1 scheduler.waiting_queue[1].priority = 10 scheduler.enable_priority_scheduling = True - scheduler.policy.calc_priority.side_effect = ( - lambda waiting_queue, _: waiting_queue.sort(key=lambda req: -req.priority) + scheduler.policy.calc_priority.side_effect = lambda waiting_queue, _: ( + waiting_queue.sort(key=lambda req: -req.priority) ) new_batch = MagicMock() # get_new_prebuilt_batch reads the published disagg config # (disaggregation_decode_enable_radix_cache). - with patch( - "sglang.srt.disaggregation.decode.ScheduleBatch.init_new", - return_value=new_batch, - ) as init_new, get_context().override_server_args( - disaggregation_decode_enable_radix_cache=False + with ( + patch( + "sglang.srt.disaggregation.decode.ScheduleBatch.init_new", + return_value=new_batch, + ) as init_new, + get_context().override_server_args( + disaggregation_decode_enable_radix_cache=False + ), ): ret = SchedulerDisaggregationDecodeMixin.get_new_prebuilt_batch( scheduler, scheduler.running_batch @@ -490,11 +493,14 @@ class TestDecodePrebuilt(unittest.TestCase): ) new_batch.process_prebuilt.side_effect = lambda *_: call_order.append("process") - with patch( - "sglang.srt.disaggregation.decode.ScheduleBatch.init_new", - return_value=new_batch, - ), get_context().override_server_args( - disaggregation_decode_enable_radix_cache=False + with ( + patch( + "sglang.srt.disaggregation.decode.ScheduleBatch.init_new", + return_value=new_batch, + ), + get_context().override_server_args( + disaggregation_decode_enable_radix_cache=False + ), ): ret = SchedulerDisaggregationDecodeMixin.get_new_prebuilt_batch( scheduler, scheduler.running_batch diff --git a/test/registered/unit/mem_cache/test_hicache_file_lru_unit.py b/test/registered/unit/mem_cache/test_hicache_file_lru_unit.py index 6a7e2243c..a083ae3ac 100644 --- a/test/registered/unit/mem_cache/test_hicache_file_lru_unit.py +++ b/test/registered/unit/mem_cache/test_hicache_file_lru_unit.py @@ -391,9 +391,10 @@ class TestMinFreeSpaceWatermark(HiCacheFileLRUTestBase): free[0] += os.path.getsize(p) return original_remove(p) - with mock.patch.object( - b._evictor, "_fs_stats", side_effect=fake_fs_stats - ), mock.patch("os.remove", side_effect=tracked_remove): + with ( + mock.patch.object(b._evictor, "_fs_stats", side_effect=fake_fs_stats), + mock.patch("os.remove", side_effect=tracked_remove), + ): self.assertTrue(b.set("newk", _t(60))) self.assertFalse(b.exists("victim")) self.assertTrue(b.exists("newk")) @@ -516,9 +517,10 @@ class TestHiCacheFileMetadataIntegration(HiCacheFileLRUTestBase): b.set("k2", _t(50)) # Now patch os.scandir and os.path.exists - with mock.patch("os.scandir") as mock_scandir, mock.patch( - "os.path.exists" - ) as mock_exists: + with ( + mock.patch("os.scandir") as mock_scandir, + mock.patch("os.path.exists") as mock_exists, + ): mock_exists.return_value = True # batch_exists_v2 for k1 and k2 should hit the metadata cache and NOT call os.scandir or os.path.exists diff --git a/test/registered/unit/mem_cache/test_hybrid_pool_assembler.py b/test/registered/unit/mem_cache/test_hybrid_pool_assembler.py index 1005400e4..66ed58723 100644 --- a/test/registered/unit/mem_cache/test_hybrid_pool_assembler.py +++ b/test/registered/unit/mem_cache/test_hybrid_pool_assembler.py @@ -49,9 +49,7 @@ class TestDeviceAllocEviction(CustomTestCase): def test_sufficient_capacity_skips_eviction(self): cache = MagicMock() cache.token_to_kv_pool_allocator.swa_available_size.return_value = 10 - cache.req_to_token_pool.mamba_allocator.schedulable_available_size.return_value = ( - 10 - ) + cache.req_to_token_pool.mamba_allocator.schedulable_available_size.return_value = 10 _evict_swa_for_device_alloc(cache, required_size=10) _evict_mamba_for_device_alloc(cache, required_size=10) diff --git a/test/registered/unit/mem_cache/test_kv_index_translator.py b/test/registered/unit/mem_cache/test_kv_index_translator.py index cf049b4e5..27fbaac78 100644 --- a/test/registered/unit/mem_cache/test_kv_index_translator.py +++ b/test/registered/unit/mem_cache/test_kv_index_translator.py @@ -182,7 +182,6 @@ def _alloc_and_fill(allocator, ps, lens): class TestReadTableBuild(unittest.TestCase): - def test_read_table_matches_reference_across_multipliers(self): """The load-bearing formula pin: full AND swa read tables equal the independent per-element derivation, across page sizes and both diff --git a/test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py b/test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py index 1fd938b2a..3bcdb9722 100644 --- a/test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py +++ b/test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py @@ -117,7 +117,9 @@ class TestMambaRatioEnvGate(unittest.TestCase): strategy = ( "extra_buffer_lazy" if lazy - else "extra_buffer" if extra_buffer else "no_buffer" + else "extra_buffer" + if extra_buffer + else "no_buffer" ) from sglang.srt import runtime_context as rc diff --git a/test/registered/unit/mem_cache/test_mamba_path_state_cap.py b/test/registered/unit/mem_cache/test_mamba_path_state_cap.py index e6e7b0476..ec4f425af 100644 --- a/test/registered/unit/mem_cache/test_mamba_path_state_cap.py +++ b/test/registered/unit/mem_cache/test_mamba_path_state_cap.py @@ -105,9 +105,12 @@ class TestMambaPathStateCap(unittest.TestCase): def test_server_arg_rejects_zero_and_values_below_negative_one(self): for value in (0, -2): args = ServerArgs(model_path="dummy", mamba_max_states_per_path=value) - with self.subTest(value=value), self.assertRaisesRegex( - ValueError, - "must be -1 \\(unlimited\\) or a positive integer", + with ( + self.subTest(value=value), + self.assertRaisesRegex( + ValueError, + "must be -1 \\(unlimited\\) or a positive integer", + ), ): handle_mamba_backend(args) diff --git a/test/registered/unit/mem_cache/test_mamba_unittest.py b/test/registered/unit/mem_cache/test_mamba_unittest.py index 313de93ae..78cdc559c 100755 --- a/test/registered/unit/mem_cache/test_mamba_unittest.py +++ b/test/registered/unit/mem_cache/test_mamba_unittest.py @@ -300,17 +300,17 @@ class TestMamba(unittest.TestCase): full_num_tokens = 1 print(f"evicting {full_num_tokens} full token") result = tree.evict(EvictParams(num_tokens=full_num_tokens)) - assert ( - result.num_tokens_evicted >= full_num_tokens - ), f"evicted {result.num_tokens_evicted} full tokens, expected {full_num_tokens}" + assert result.num_tokens_evicted >= full_num_tokens, ( + f"evicted {result.num_tokens_evicted} full tokens, expected {full_num_tokens}" + ) tree.pretty_print() mamba_num = 1 print(f"evicting {mamba_num} mamba") result = tree.evict(EvictParams(num_tokens=0, mamba_num=mamba_num)) - assert ( - result.mamba_num_evicted >= mamba_num - ), f"evicted {result.mamba_num_evicted} mamba states, expected {mamba_num}" + assert result.mamba_num_evicted >= mamba_num, ( + f"evicted {result.mamba_num_evicted} mamba states, expected {mamba_num}" + ) tree.pretty_print() req5_token_ids = [1, 2, 3, 4, 5] @@ -348,9 +348,9 @@ class TestMamba(unittest.TestCase): mamba_num = 1 print(f"evicting {mamba_num} mamba") result = tree.evict(EvictParams(num_tokens=0, mamba_num=mamba_num)) - assert ( - result.mamba_num_evicted >= mamba_num - ), f"evicted {result.mamba_num_evicted} mamba states, expected {mamba_num}" + assert result.mamba_num_evicted >= mamba_num, ( + f"evicted {result.mamba_num_evicted} mamba states, expected {mamba_num}" + ) tree.pretty_print() req8_token_ids = [1, 2, 3, 4, 5, 60, 70] diff --git a/test/registered/unit/mem_cache/test_mem_pool_host.py b/test/registered/unit/mem_cache/test_mem_pool_host.py index 61c4a0931..98e1b59b1 100644 --- a/test/registered/unit/mem_cache/test_mem_pool_host.py +++ b/test/registered/unit/mem_cache/test_mem_pool_host.py @@ -242,10 +242,11 @@ class TestHostMemoryBudget(CustomTestCase): # Deliberate single-accessor stub: isolates the budget math from the # topology derivation, which the ranks_per_host case below covers. fake_mem = unittest.mock.Mock(available=self._AVAILABLE) - with unittest.mock.patch.object( - base, "ranks_per_host", return_value=ranks - ), unittest.mock.patch.object( - base.psutil, "virtual_memory", return_value=fake_mem + with ( + unittest.mock.patch.object(base, "ranks_per_host", return_value=ranks), + unittest.mock.patch.object( + base.psutil, "virtual_memory", return_value=fake_mem + ), ): return base.host_memory_budget_bytes() @@ -264,9 +265,15 @@ class TestHostMemoryBudget(CustomTestCase): # The launcher slices ranks uniformly across nodes, so the co-located # rank count is world_size // nnodes — no hostname collective. fake_group = unittest.mock.Mock(world_size=16) - with get_context().override_server_args(nnodes=2), unittest.mock.patch.object( - torch.distributed, "is_initialized", return_value=True - ), unittest.mock.patch.object(base, "get_world_group", return_value=fake_group): + with ( + get_context().override_server_args(nnodes=2), + unittest.mock.patch.object( + torch.distributed, "is_initialized", return_value=True + ), + unittest.mock.patch.object( + base, "get_world_group", return_value=fake_group + ), + ): self.assertEqual(base.ranks_per_host(), 8) diff --git a/test/registered/unit/mem_cache/test_mmap_allocator.py b/test/registered/unit/mem_cache/test_mmap_allocator.py index d41d95eae..3090b1671 100644 --- a/test/registered/unit/mem_cache/test_mmap_allocator.py +++ b/test/registered/unit/mem_cache/test_mmap_allocator.py @@ -83,9 +83,12 @@ class TestMmapAllocator(unittest.TestCase): # MAP_POPULATE is unreachable on a 5.14+ kernel, so CI never runs it; # force the branch or it ships untested. - with self.subTest(path="map_populate"), unittest.mock.patch( - "sglang.srt.mem_cache.storage.mmap.mmap_allocator._has_madv_populate_write", - return_value=False, + with ( + self.subTest(path="map_populate"), + unittest.mock.patch( + "sglang.srt.mem_cache.storage.mmap.mmap_allocator._has_madv_populate_write", + return_value=False, + ), ): mm = _mmap_prefaulted(-1, alloc_bytes, flags) try: diff --git a/test/registered/unit/mem_cache/test_multi_ended_allocator.py b/test/registered/unit/mem_cache/test_multi_ended_allocator.py index e24050807..eee427d05 100644 --- a/test/registered/unit/mem_cache/test_multi_ended_allocator.py +++ b/test/registered/unit/mem_cache/test_multi_ended_allocator.py @@ -3141,12 +3141,16 @@ class TestFloatMultiEndedAllocator(unittest.TestCase): _, _, fla, _, kv = self._build_tri() v = fla.alloc(4) self._stamp(fla, kv, v) - with mock.patch.object( - torch.Tensor, "tolist", side_effect=AssertionError("tolist = D2H") - ), mock.patch.object( - torch.Tensor, "item", side_effect=AssertionError("item = D2H") - ), mock.patch.object( - torch, "unique", side_effect=AssertionError("unique = host sync") + with ( + mock.patch.object( + torch.Tensor, "tolist", side_effect=AssertionError("tolist = D2H") + ), + mock.patch.object( + torch.Tensor, "item", side_effect=AssertionError("item = D2H") + ), + mock.patch.object( + torch, "unique", side_effect=AssertionError("unique = host sync") + ), ): fla.free(v[:2], _pages=v[:2]) diff --git a/test/registered/unit/mem_cache/test_radix_cache_slru_accuracy.py b/test/registered/unit/mem_cache/test_radix_cache_slru_accuracy.py index 72d82f84f..78d47c62b 100644 --- a/test/registered/unit/mem_cache/test_radix_cache_slru_accuracy.py +++ b/test/registered/unit/mem_cache/test_radix_cache_slru_accuracy.py @@ -18,7 +18,6 @@ register_cpu_ci(est_time=8, suite="base-a-test-cpu") class TestSLRUAccuracy(unittest.TestCase): - def setUp(self): """Setup minimal memory pools for testing""" torch.set_default_device(None) diff --git a/test/registered/unit/mem_cache/test_rust_tree_core_integration.py b/test/registered/unit/mem_cache/test_rust_tree_core_integration.py index 3888412f3..945029445 100644 --- a/test/registered/unit/mem_cache/test_rust_tree_core_integration.py +++ b/test/registered/unit/mem_cache/test_rust_tree_core_integration.py @@ -882,8 +882,9 @@ def test_prefetch_node_accessors_round_trip(): assert not core.is_backuped(leaf) assert not core.is_root(leaf) - assert core.get_last_hash_value(leaf) == ( - mem_cache.get_hash_str(array("q", [1, 2]), None, 2)[-1] + assert ( + core.get_last_hash_value(leaf) + == (mem_cache.get_hash_str(array("q", [1, 2]), None, 2)[-1]) ) assert core.get_prefix_hash_values(leaf) == [] diff --git a/test/registered/unit/mem_cache/test_session_token_share_unit.py b/test/registered/unit/mem_cache/test_session_token_share_unit.py index b83c84489..9e08330e6 100644 --- a/test/registered/unit/mem_cache/test_session_token_share_unit.py +++ b/test/registered/unit/mem_cache/test_session_token_share_unit.py @@ -53,7 +53,6 @@ def _recv(rid, input_ids, max_new_tokens=8): class TestSessionTokenShare(CustomTestCase): - def setUp(self): self.session = Session(capacity_of_str_len=0, session_id="s", streaming=True) diff --git a/test/registered/unit/mem_cache/test_swa_eviction_boundary.py b/test/registered/unit/mem_cache/test_swa_eviction_boundary.py index 94a04d52a..773d5946c 100644 --- a/test/registered/unit/mem_cache/test_swa_eviction_boundary.py +++ b/test/registered/unit/mem_cache/test_swa_eviction_boundary.py @@ -134,7 +134,6 @@ def _make_batch(tree, allocator, pool): class TestSWAEvictionBoundary(unittest.TestCase): - # -- Eviction formula: page_size > window -- def test_formula_page_gt_window_sweep(self): diff --git a/test/registered/unit/mem_cache/test_swa_unittest.py b/test/registered/unit/mem_cache/test_swa_unittest.py index cf6267b1d..0df9e7f74 100644 --- a/test/registered/unit/mem_cache/test_swa_unittest.py +++ b/test/registered/unit/mem_cache/test_swa_unittest.py @@ -727,9 +727,9 @@ class TestSWA(unittest.TestCase): EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) ) assert isinstance(evict_result, EvictResult) - assert ( - evict_result.swa_num_tokens_evicted >= swa_num_tokens - ), f"evicted {evict_result.swa_num_tokens_evicted} swa tokens, expected {swa_num_tokens}" + assert evict_result.swa_num_tokens_evicted >= swa_num_tokens, ( + f"evicted {evict_result.swa_num_tokens_evicted} swa tokens, expected {swa_num_tokens}" + ) tree.pretty_print() full_num_tokens, swa_num_tokens = 1, 2 @@ -738,12 +738,12 @@ class TestSWA(unittest.TestCase): EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) ) assert isinstance(evict_result, EvictResult) - assert ( - evict_result.num_tokens_evicted >= full_num_tokens - ), f"evicted {evict_result.num_tokens_evicted} full tokens, expected {full_num_tokens}" - assert ( - evict_result.swa_num_tokens_evicted >= swa_num_tokens - ), f"evicted {evict_result.swa_num_tokens_evicted} swa tokens, expected {swa_num_tokens}" + assert evict_result.num_tokens_evicted >= full_num_tokens, ( + f"evicted {evict_result.num_tokens_evicted} full tokens, expected {full_num_tokens}" + ) + assert evict_result.swa_num_tokens_evicted >= swa_num_tokens, ( + f"evicted {evict_result.swa_num_tokens_evicted} swa tokens, expected {swa_num_tokens}" + ) tree.pretty_print() req5_token_ids = [1, 2, 3, 4, 5] diff --git a/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py b/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py index 0bee23ad4..b39989a47 100644 --- a/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py +++ b/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py @@ -457,10 +457,13 @@ class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase): alloc = self._swa_composite(lazy=True) v = alloc.alloc(8 * self.PS) self.assertIsNotNone(v) - with mock.patch.object( - torch, "unique", side_effect=AssertionError("unique = host sync") - ), mock.patch.object( - torch.Tensor, "item", side_effect=AssertionError("item = host sync") + with ( + mock.patch.object( + torch, "unique", side_effect=AssertionError("unique = host sync") + ), + mock.patch.object( + torch.Tensor, "item", side_effect=AssertionError("item = host sync") + ), ): alloc.free_swa(v[: 4 * self.PS], start_pos=0) alloc.free_swa(v[4 * self.PS :], start_pos=4 * self.PS) diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py index 9a07d97ad..1856aaf81 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py @@ -424,9 +424,9 @@ def bench_api( (excluded from latency measurement). """ items = setup_fn() - assert ( - len(items) >= num_ops + warmup - ), f"need {num_ops + warmup} items, got {len(items)}" + assert len(items) >= num_ops + warmup, ( + f"need {num_ops + warmup} items, got {len(items)}" + ) for i in range(warmup): op_fn(items[i]) diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index a65bfe49c..9637ba6e7 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -299,11 +299,11 @@ class TestUnifiedTreeCoreLoadBackPending(CustomTestCase): core.node_by_id.side_effect = nodes.__getitem__ core.components_by_type = {ComponentType.FULL: mock.Mock()} core.full_host_duplicates = {} - core._is_settled_full_host_duplicate.side_effect = ( - lambda node: UnifiedTreeCore._is_settled_full_host_duplicate(core, node) + core._is_settled_full_host_duplicate.side_effect = lambda node: ( + UnifiedTreeCore._is_settled_full_host_duplicate(core, node) ) - core._update_duplicate_tracking.side_effect = ( - lambda node: UnifiedTreeCore._update_duplicate_tracking(core, node) + core._update_duplicate_tracking.side_effect = lambda node: ( + UnifiedTreeCore._update_duplicate_tracking(core, node) ) return core, shared, anchor_a, anchor_b @@ -1265,7 +1265,6 @@ class TestUnifiedRadixCacheKVEvents(CustomTestCase): class UnifiedRadixCacheSuite: - cfg: CacheConfig _rid: int = 0 @@ -3572,8 +3571,9 @@ class UnifiedRadixCacheSuite: self.assertIn(n, pipeline.inflight_backup_node_ids) self._pump_hicache_until( cache, - lambda: not pipeline.inflight_backup_node_ids - and not pipeline.ongoing_backup, + lambda: ( + not pipeline.inflight_backup_node_ids and not pipeline.ongoing_backup + ), "buffer backup pipeline did not drain", ) @@ -3689,8 +3689,10 @@ class UnifiedRadixCacheSuite: self.assertEqual((stats["attempts"], stats["issued"]), (1, 1)) self._pump_hicache_until( cons, - lambda: cons.check_prefetch_progress(req_id) - and cons.buffer_pipeline.has_staged(req_id), + lambda: ( + cons.check_prefetch_progress(req_id) + and cons.buffer_pipeline.has_staged(req_id) + ), "prefetch did not stage", ) # Staged: bounce occupies host staging; nothing device-side; span @@ -3753,8 +3755,10 @@ class UnifiedRadixCacheSuite: # loaded KV bytes equal the producer's. self._pump_hicache_until( cons, - lambda: not cons.buffer_pipeline.ongoing_buffer_load_back - and self._host_avail_sizes(cons) == avail0, + lambda: ( + not cons.buffer_pipeline.ongoing_buffer_load_back + and self._host_avail_sizes(cons) == avail0 + ), "load-back ack did not free the bounce", ) mc = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) @@ -3818,8 +3822,10 @@ class UnifiedRadixCacheSuite: ) self._pump_hicache_until( cons, - lambda: cons.check_prefetch_progress(root_req) - and cons.buffer_pipeline.has_staged(root_req), + lambda: ( + cons.check_prefetch_progress(root_req) + and cons.buffer_pipeline.has_staged(root_req) + ), "salted root prefetch did not stage", ) held = cons.buffer_pipeline.staged_prefetches[root_req] @@ -3888,8 +3894,10 @@ class UnifiedRadixCacheSuite: ) self._pump_hicache_until( cons2, - lambda: cons2.check_prefetch_progress(anchored_req) - and cons2.buffer_pipeline.has_staged(anchored_req), + lambda: ( + cons2.check_prefetch_progress(anchored_req) + and cons2.buffer_pipeline.has_staged(anchored_req) + ), "salted mid-tree prefetch did not stage", ) self._consume_staged_prefetch( @@ -3950,8 +3958,10 @@ class UnifiedRadixCacheSuite: ) self._pump_hicache_until( cons, - lambda: cons.check_prefetch_progress(req_id) - and cons.buffer_pipeline.has_staged(req_id), + lambda: ( + cons.check_prefetch_progress(req_id) + and cons.buffer_pipeline.has_staged(req_id) + ), "retried prefetch did not stage", ) self.assertFalse(cons.pop_storage_prefetch_miss(req_id)) @@ -4047,8 +4057,10 @@ class UnifiedRadixCacheSuite: ) self._pump_hicache_until( cons, - lambda: cons.check_prefetch_progress(req_id) - and cons.buffer_pipeline.has_staged(req_id), + lambda: ( + cons.check_prefetch_progress(req_id) + and cons.buffer_pipeline.has_staged(req_id) + ), "prefetch did not stage", ) @@ -4151,8 +4163,10 @@ class UnifiedRadixCacheSuite: ) self._pump_hicache_until( cons, - lambda: cons.check_prefetch_progress(req_id) - and cons.buffer_pipeline.has_staged(req_id), + lambda: ( + cons.check_prefetch_progress(req_id) + and cons.buffer_pipeline.has_staged(req_id) + ), "prefetch did not stage", ) cons.pop_prefetch_loaded_tokens(req_id) @@ -4222,8 +4236,10 @@ class UnifiedRadixCacheSuite: ) self._pump_hicache_until( cons, - lambda: cons.check_prefetch_progress(req_id) - and cons.buffer_pipeline.has_staged(req_id), + lambda: ( + cons.check_prefetch_progress(req_id) + and cons.buffer_pipeline.has_staged(req_id) + ), "prefetch did not stage", ) cons.pop_prefetch_loaded_tokens(req_id) @@ -4331,8 +4347,10 @@ class UnifiedRadixCacheSuite: ) self._pump_hicache_until( cons, - lambda: cons.check_prefetch_progress(req_id) - and cons.buffer_pipeline.has_staged(req_id), + lambda: ( + cons.check_prefetch_progress(req_id) + and cons.buffer_pipeline.has_staged(req_id) + ), "prefetch did not stage", ) cons.pop_prefetch_loaded_tokens(req_id) @@ -4398,8 +4416,10 @@ class UnifiedRadixCacheSuite: ) self._pump_hicache_until( cons, - lambda: cons.check_prefetch_progress(req_id) - and cons.buffer_pipeline.has_staged(req_id), + lambda: ( + cons.check_prefetch_progress(req_id) + and cons.buffer_pipeline.has_staged(req_id) + ), "prefetch did not stage", ) cons.pop_prefetch_loaded_tokens(req_id) @@ -4519,8 +4539,10 @@ class UnifiedRadixCacheSuite: ) self._pump_hicache_until( cons2, - lambda: cons2.check_prefetch_progress("subwin-req") - and cons2.buffer_pipeline.has_staged("subwin-req"), + lambda: ( + cons2.check_prefetch_progress("subwin-req") + and cons2.buffer_pipeline.has_staged("subwin-req") + ), "sub-window prefetch did not stage", ) self.assertTrue( @@ -7081,7 +7103,6 @@ class UnifiedRadixCacheSuite: class UnifiedLRUListBoundedRefreshTest(CustomTestCase): - components = (ComponentType.FULL, ComponentType.SWA) def _make_node(self, key_len: int) -> UnifiedTreeNode: diff --git a/test/registered/unit/mem_cache/test_unified_tri_pool.py b/test/registered/unit/mem_cache/test_unified_tri_pool.py index 5aae3bdb0..8ca36f336 100644 --- a/test/registered/unit/mem_cache/test_unified_tri_pool.py +++ b/test/registered/unit/mem_cache/test_unified_tri_pool.py @@ -465,10 +465,13 @@ class TestTriFreeSwaNoHostSync(unittest.TestCase): self.assertIsNotNone(v) from unittest import mock - with mock.patch.object( - torch, "unique", side_effect=AssertionError("unique = host sync") - ), mock.patch.object( - torch.Tensor, "item", side_effect=AssertionError("item = host sync") + with ( + mock.patch.object( + torch, "unique", side_effect=AssertionError("unique = host sync") + ), + mock.patch.object( + torch.Tensor, "item", side_effect=AssertionError("item = host sync") + ), ): alloc.free_swa(v[: 4 * self.PS], start_pos=0) self.assertEqual(alloc.verify_byte_accounting(), []) diff --git a/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py b/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py index a1f6bb2cd..90a91ac40 100644 --- a/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py +++ b/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py @@ -60,13 +60,17 @@ class TestInitProfileBatchMode(CustomTestCase): env = {_BATCH_CAPTURE: "1"} if profiler_dir is not None: env["SGLANG_TORCH_PROFILER_DIR"] = profiler_dir - with mock.patch.dict(os.environ, env, clear=False), mock.patch.object( - mod, "get_parallel", return_value=SimpleNamespace(tp_rank=rank) - ), mock.patch.object(mod, "profile") as mock_profile, mock.patch( - "torch.profiler.schedule" - ) as mock_schedule, mock.patch( - "torch.cuda.memory._record_memory_history" - ) as mock_record_history: + with ( + mock.patch.dict(os.environ, env, clear=False), + mock.patch.object( + mod, "get_parallel", return_value=SimpleNamespace(tp_rank=rank) + ), + mock.patch.object(mod, "profile") as mock_profile, + mock.patch("torch.profiler.schedule") as mock_schedule, + mock.patch( + "torch.cuda.memory._record_memory_history" + ) as mock_record_history, + ): os.environ.pop(_CAPTURE_TRACE, None) # original flag off if profiler_dir is None: os.environ.pop("SGLANG_TORCH_PROFILER_DIR", None) @@ -110,19 +114,16 @@ class TestInitProfileBatchMode(CustomTestCase): # No SGLANG_TORCH_PROFILER_DIR -> falls back to the envs default base dir. # Patch makedirs so the test never writes to the cwd. fake_self = _make_fake_self([1]) - with mock.patch.dict( - os.environ, {_BATCH_CAPTURE: "1"}, clear=False - ), mock.patch.object( - mod, "get_parallel", return_value=SimpleNamespace(tp_rank=0) - ), mock.patch.object( - mod, "profile" - ), mock.patch( - "torch.profiler.schedule" - ), mock.patch( - "torch.cuda.memory._record_memory_history" - ), mock.patch.object( - mod.os, "makedirs" - ) as mock_makedirs: + with ( + mock.patch.dict(os.environ, {_BATCH_CAPTURE: "1"}, clear=False), + mock.patch.object( + mod, "get_parallel", return_value=SimpleNamespace(tp_rank=0) + ), + mock.patch.object(mod, "profile"), + mock.patch("torch.profiler.schedule"), + mock.patch("torch.cuda.memory._record_memory_history"), + mock.patch.object(mod.os, "makedirs") as mock_makedirs, + ): os.environ.pop("SGLANG_TORCH_PROFILER_DIR", None) os.environ.pop(_CAPTURE_TRACE, None) DecodeCudaGraphRunner._init_profile_context_and_memory_record(fake_self) @@ -143,12 +144,14 @@ class TestInitProfileOriginalMode(CustomTestCase): with tempfile.TemporaryDirectory() as tmp: environ = dict(env) environ["SGLANG_TORCH_PROFILER_DIR"] = tmp - with mock.patch.dict(os.environ, environ, clear=False), mock.patch.object( - mod, "get_parallel", return_value=SimpleNamespace(tp_rank=0) - ), mock.patch.object(mod, "profile") as mock_profile, mock.patch( - "torch.profiler.schedule" - ) as mock_schedule, mock.patch( - "torch.cuda.memory._record_memory_history" + with ( + mock.patch.dict(os.environ, environ, clear=False), + mock.patch.object( + mod, "get_parallel", return_value=SimpleNamespace(tp_rank=0) + ), + mock.patch.object(mod, "profile") as mock_profile, + mock.patch("torch.profiler.schedule") as mock_schedule, + mock.patch("torch.cuda.memory._record_memory_history"), ): for k in (_CAPTURE_TRACE, _BATCH_CAPTURE): if k not in environ: @@ -176,18 +179,18 @@ class TestInitProfileOriginalMode(CustomTestCase): class TestOnTraceReadyNaming(CustomTestCase): def _build_on_trace_ready(self, *, capture_bs, rank, tmp): fake_self = _make_fake_self(capture_bs) - with mock.patch.dict( - os.environ, - {"SGLANG_TORCH_PROFILER_DIR": tmp, _BATCH_CAPTURE: "1"}, - clear=False, - ), mock.patch.object( - mod, "get_parallel", return_value=SimpleNamespace(tp_rank=rank) - ), mock.patch.object( - mod, "profile" - ) as mock_profile, mock.patch( - "torch.profiler.schedule" - ), mock.patch( - "torch.cuda.memory._record_memory_history" + with ( + mock.patch.dict( + os.environ, + {"SGLANG_TORCH_PROFILER_DIR": tmp, _BATCH_CAPTURE: "1"}, + clear=False, + ), + mock.patch.object( + mod, "get_parallel", return_value=SimpleNamespace(tp_rank=rank) + ), + mock.patch.object(mod, "profile") as mock_profile, + mock.patch("torch.profiler.schedule"), + mock.patch("torch.cuda.memory._record_memory_history"), ): os.environ.pop(_CAPTURE_TRACE, None) DecodeCudaGraphRunner._init_profile_context_and_memory_record(fake_self) diff --git a/test/registered/unit/model_executor/runner/test_hidden_state_graph_recapture.py b/test/registered/unit/model_executor/runner/test_hidden_state_graph_recapture.py index eb65396bc..b3e52808c 100644 --- a/test/registered/unit/model_executor/runner/test_hidden_state_graph_recapture.py +++ b/test/registered/unit/model_executor/runner/test_hidden_state_graph_recapture.py @@ -118,9 +118,12 @@ class TestHiddenStateGraphRecapture(CustomTestCase): ): runner = self._make_runner(runner_cls, CaptureHiddenMode.NULL) - with self.subTest(runner_cls=runner_cls), self.assertRaisesRegex( - RuntimeError, - "exceeds the fixed (CUDA|CPU) graph capture mode", + with ( + self.subTest(runner_cls=runner_cls), + self.assertRaisesRegex( + RuntimeError, + "exceeds the fixed (CUDA|CPU) graph capture mode", + ), ): runner._validate_capture_hidden_mode( self._make_forward_batch(CaptureHiddenMode.LAST) diff --git a/test/registered/unit/model_executor/runner_backend/test_full_cuda_graph_backend.py b/test/registered/unit/model_executor/runner_backend/test_full_cuda_graph_backend.py index 4a25e37d1..32c767ff4 100644 --- a/test/registered/unit/model_executor/runner_backend/test_full_cuda_graph_backend.py +++ b/test/registered/unit/model_executor/runner_backend/test_full_cuda_graph_backend.py @@ -139,8 +139,11 @@ class TestCaptureOneWithProfiling(CustomTestCase): rf_names.append(name) return contextlib.nullcontext() - with mock.patch("torch.cuda.CUDAGraph", return_value="GRAPH"), mock.patch( - "torch.profiler.record_function", side_effect=_fake_record_function + with ( + mock.patch("torch.cuda.CUDAGraph", return_value="GRAPH"), + mock.patch( + "torch.profiler.record_function", side_effect=_fake_record_function + ), ): backend.capture_one(ShapeKey(size=size), forward_fn) diff --git a/test/registered/unit/model_executor/test_model_runner_decode_rows.py b/test/registered/unit/model_executor/test_model_runner_decode_rows.py index b96cf1e47..d30352109 100644 --- a/test/registered/unit/model_executor/test_model_runner_decode_rows.py +++ b/test/registered/unit/model_executor/test_model_runner_decode_rows.py @@ -35,14 +35,18 @@ class TestModelRunnerDecodeRows(unittest.TestCase): spec = SimpleNamespace( speculative_adaptive=True, speculative_adaptive_config=f.name ) - with patch( - "sglang.srt.model_executor.model_runner.get_spec", return_value=spec - ), patch( - "sglang.srt.model_executor.model_runner.max_speculative_num_draft_tokens", - return_value=6, - ), patch( - "sglang.srt.model_executor.model_runner.get_batch_sizes_to_capture", - side_effect=_alignment_8_capture_bs, + with ( + patch( + "sglang.srt.model_executor.model_runner.get_spec", return_value=spec + ), + patch( + "sglang.srt.model_executor.model_runner.max_speculative_num_draft_tokens", + return_value=6, + ), + patch( + "sglang.srt.model_executor.model_runner.get_batch_sizes_to_capture", + side_effect=_alignment_8_capture_bs, + ), ): self.assertEqual(runner.max_decode_logits_rows(), 72) diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index 3412a18d2..95d73e6af 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -813,8 +813,9 @@ class TestDSAIndexerAllocationPolicy(CustomTestCase): mr.model_config.hf_config.index_topk_freq = 4 mr.model_config.hf_config.index_skip_topk_offset = 3 - with get_memory().override(enable_hierarchical_cache=True), mock_cpu_env( - kv_size=1 + with ( + get_memory().override(enable_hierarchical_cache=True), + mock_cpu_env(kv_size=1), ): from sglang.srt.model_executor.pool_configurator import ( DefaultPoolConfigurator, diff --git a/test/registered/unit/model_loader/test_presharded_loader.py b/test/registered/unit/model_loader/test_presharded_loader.py index bc1aea8cc..f1b5fddef 100644 --- a/test/registered/unit/model_loader/test_presharded_loader.py +++ b/test/registered/unit/model_loader/test_presharded_loader.py @@ -374,7 +374,10 @@ class TestBuildDumpPlan(unittest.TestCase): # and depend on (name, content-SHA) pairs of every tensor a rank # owns. Permuting the manifest's insertion order must not change # the rank checksum. - with tempfile.TemporaryDirectory() as tmp_a, tempfile.TemporaryDirectory() as tmp_b: + with ( + tempfile.TemporaryDirectory() as tmp_a, + tempfile.TemporaryDirectory() as tmp_b, + ): base_entries = { "alpha.weight": { "checksum": "h_alpha", @@ -403,7 +406,10 @@ class TestBuildDumpPlan(unittest.TestCase): def test_rank_checksum_distinguishes_content(self): # Changing one tensor's content-SHA must change the rank checksum. - with tempfile.TemporaryDirectory() as tmp_a, tempfile.TemporaryDirectory() as tmp_b: + with ( + tempfile.TemporaryDirectory() as tmp_a, + tempfile.TemporaryDirectory() as tmp_b, + ): entries_a = { "x.weight": { "checksum": "ha", @@ -537,10 +543,10 @@ class TestBuildDumpPlan(unittest.TestCase): # read-only dump root (HF cache mounts). Mock OSError because root # can often still write to mode-0555 dirs under CAP_DAC_OVERRIDE. loader = object.__new__(PreshardedModelLoader) - with mock.patch.object( - loader, "_world_rank_and_size", return_value=(0, 1) - ), mock.patch.object(loader, "_world_barrier"), mock.patch( - "os.makedirs", side_effect=OSError("Read-only file system") + with ( + mock.patch.object(loader, "_world_rank_and_size", return_value=(0, 1)), + mock.patch.object(loader, "_world_barrier"), + mock.patch("os.makedirs", side_effect=OSError("Read-only file system")), ): with self.assertRaises(RuntimeError) as ctx: loader._ensure_presharded_dir_writable("/ro/presharded") @@ -551,9 +557,10 @@ class TestBuildDumpPlan(unittest.TestCase): loader = object.__new__(PreshardedModelLoader) with tempfile.TemporaryDirectory() as tmp: leaf = os.path.join(tmp, "TP-8-sig-test") - with mock.patch.object( - loader, "_world_rank_and_size", return_value=(0, 1) - ), mock.patch.object(loader, "_world_barrier") as barrier: + with ( + mock.patch.object(loader, "_world_rank_and_size", return_value=(0, 1)), + mock.patch.object(loader, "_world_barrier") as barrier, + ): loader._ensure_presharded_dir_writable(leaf) self.assertTrue(os.path.isdir(leaf)) barrier.assert_called_once() @@ -605,15 +612,21 @@ class TestStructuralSignature(unittest.TestCase): loader = object.__new__(PreshardedModelLoader) loader.load_config = SimpleNamespace() - with mock.patch( - "sglang.srt.model_loader.loader._get_quantization_config", - return_value=None, - ), mock.patch( - "sglang.srt.model_loader.loader._initialize_model", - side_effect=_init_stub, - ), mock.patch( - "sglang.srt.model_loader.loader.set_default_torch_dtype", - return_value=mock.MagicMock(__enter__=mock.Mock(), __exit__=mock.Mock()), + with ( + mock.patch( + "sglang.srt.model_loader.loader._get_quantization_config", + return_value=None, + ), + mock.patch( + "sglang.srt.model_loader.loader._initialize_model", + side_effect=_init_stub, + ), + mock.patch( + "sglang.srt.model_loader.loader.set_default_torch_dtype", + return_value=mock.MagicMock( + __enter__=mock.Mock(), __exit__=mock.Mock() + ), + ), ): sig_ab = loader._compute_local_structural_signature( SimpleNamespace( @@ -694,15 +707,21 @@ class TestStructuralSignature(unittest.TestCase): def _initialize_model_stub(model_config, load_config, quant_config): return nn.Linear(4, model_config.width, bias=False) - with mock.patch( - "sglang.srt.model_loader.loader._get_quantization_config", - return_value=None, - ), mock.patch( - "sglang.srt.model_loader.loader._initialize_model", - side_effect=_initialize_model_stub, - ), mock.patch( - "sglang.srt.model_loader.loader.set_default_torch_dtype", - return_value=mock.MagicMock(__enter__=mock.Mock(), __exit__=mock.Mock()), + with ( + mock.patch( + "sglang.srt.model_loader.loader._get_quantization_config", + return_value=None, + ), + mock.patch( + "sglang.srt.model_loader.loader._initialize_model", + side_effect=_initialize_model_stub, + ), + mock.patch( + "sglang.srt.model_loader.loader.set_default_torch_dtype", + return_value=mock.MagicMock( + __enter__=mock.Mock(), __exit__=mock.Mock() + ), + ), ): loader_narrow = FakeModelLoader(width=2) loader_wide = FakeModelLoader(width=8) @@ -738,15 +757,21 @@ class TestStructuralSignature(unittest.TestCase): loader = object.__new__(PreshardedModelLoader) loader.load_config = SimpleNamespace() - with mock.patch( - "sglang.srt.model_loader.loader._get_quantization_config", - return_value=None, - ), mock.patch( - "sglang.srt.model_loader.loader._initialize_model", - side_effect=RuntimeError("simulated init failure"), - ), mock.patch( - "sglang.srt.model_loader.loader.set_default_torch_dtype", - return_value=mock.MagicMock(__enter__=mock.Mock(), __exit__=mock.Mock()), + with ( + mock.patch( + "sglang.srt.model_loader.loader._get_quantization_config", + return_value=None, + ), + mock.patch( + "sglang.srt.model_loader.loader._initialize_model", + side_effect=RuntimeError("simulated init failure"), + ), + mock.patch( + "sglang.srt.model_loader.loader.set_default_torch_dtype", + return_value=mock.MagicMock( + __enter__=mock.Mock(), __exit__=mock.Mock() + ), + ), ): result = loader._compute_structural_signature( SimpleNamespace(quantization=None, dtype=torch.float32) @@ -816,23 +841,26 @@ class TestShardConfig(unittest.TestCase): ) override.install() self.addCleanup(override.restore) - with get_parallel().override( - tp_size=8, pp_size=1, moe_dp_size=2, moe_ep_size=4 - ), mock.patch( - "sglang.srt.layers.dp_attention.get_moe_cp_size", - return_value=2, - ), mock.patch( - "sglang.srt.model_loader.loader.get_exec", - return_value=SimpleNamespace( - features=SimpleNamespace(enable_fp32_lm_head=True), - moe=SimpleNamespace( - ep_num_redundant_experts=4, - enable_eplb=True, - init_expert_location="trivial", + with ( + get_parallel().override(tp_size=8, pp_size=1, moe_dp_size=2, moe_ep_size=4), + mock.patch( + "sglang.srt.layers.dp_attention.get_moe_cp_size", + return_value=2, + ), + mock.patch( + "sglang.srt.model_loader.loader.get_exec", + return_value=SimpleNamespace( + features=SimpleNamespace(enable_fp32_lm_head=True), + moe=SimpleNamespace( + ep_num_redundant_experts=4, + enable_eplb=True, + init_expert_location="trivial", + ), ), ), - ), mock.patch.object( - loader, "_compute_structural_signature", return_value="sig16" + mock.patch.object( + loader, "_compute_structural_signature", return_value="sig16" + ), ): cfg = loader._collect_shard_config(model_config) self.assertEqual(required, set(cfg.keys())) @@ -931,23 +959,28 @@ class TestShardConfig(unittest.TestCase): # Force rank 0 / world 1 so the method runs the rank-0 prologue # and then fails early on empty state (no need for full dump). - with mock.patch.object( - PreshardedModelLoader, - "_world_rank_and_size", - return_value=(0, 1), - ), mock.patch.object( - PreshardedModelLoader, "_world_barrier", return_value=None - ), mock.patch.object( - PreshardedModelLoader, - "_build_dump_plan", - return_value={ - "version": PreshardedModelLoader.PLAN_VERSION, - "files": [], - "rank_to_reads": {"0": []}, - "rank_checksums": {"0": "0"}, - }, - ), mock.patch.object( - PreshardedModelLoader, "_dump_files_for_rank", return_value=None + with ( + mock.patch.object( + PreshardedModelLoader, + "_world_rank_and_size", + return_value=(0, 1), + ), + mock.patch.object( + PreshardedModelLoader, "_world_barrier", return_value=None + ), + mock.patch.object( + PreshardedModelLoader, + "_build_dump_plan", + return_value={ + "version": PreshardedModelLoader.PLAN_VERSION, + "files": [], + "rank_to_reads": {"0": []}, + "rank_checksums": {"0": "0"}, + }, + ), + mock.patch.object( + PreshardedModelLoader, "_dump_files_for_rank", return_value=None + ), ): loader._dump_state_to_disk( state_dict={}, diff --git a/test/registered/unit/models/test_kimi_k3_vision.py b/test/registered/unit/models/test_kimi_k3_vision.py index 22ebdfe9e..91f4338ab 100644 --- a/test/registered/unit/models/test_kimi_k3_vision.py +++ b/test/registered/unit/models/test_kimi_k3_vision.py @@ -467,11 +467,13 @@ def test_kimi_k3_encoder_dp_defers_feature_materialization(monkeypatch): # The IPC consumer count asks for the *configured* TP size (matching # MmItemMemoryPool.try_to_recycle), so publish it; the live topology the # sharding helper reads is forced through the context's own override. - with mock_patch( - "sglang.srt.multimodal.mm_utils.run_dp_sharded_mrope_vision_model", - return_value=sharded_embeddings, - ) as run_dp, get_context().override_server_args(tp_size=1), get_parallel().override( - tp_size=1, attn_tp_size=1 + with ( + mock_patch( + "sglang.srt.multimodal.mm_utils.run_dp_sharded_mrope_vision_model", + return_value=sharded_embeddings, + ) as run_dp, + get_context().override_server_args(tp_size=1), + get_parallel().override(tp_size=1, attn_tp_size=1), ): output = model.get_image_feature(items) # Exercise the loader while the runtime topology is forced. @@ -558,14 +560,17 @@ def test_kimi_k3_preprocesses_only_dp_owner_images(monkeypatch): # Configured TP size (the IPC consumer count) comes from the published # bags; the live topology is forced through the context's own override. - with mock_patch( - "sglang.srt.multimodal.mm_utils.run_dp_sharded_mrope_vision_model", - return_value=torch.zeros(1, 2), - ) as run_dp, get_context().override_server_args(tp_size=1), get_parallel().override( - tp_size=1, attn_tp_size=1 - ), mock_patch( - "sglang.srt.multimodal.processors.kimi_k25._gpu_preprocess_images", - side_effect=fake_preprocess, + with ( + mock_patch( + "sglang.srt.multimodal.mm_utils.run_dp_sharded_mrope_vision_model", + return_value=torch.zeros(1, 2), + ) as run_dp, + get_context().override_server_args(tp_size=1), + get_parallel().override(tp_size=1, attn_tp_size=1), + mock_patch( + "sglang.srt.multimodal.processors.kimi_k25._gpu_preprocess_images", + side_effect=fake_preprocess, + ), ): model.get_image_feature(items) loader = run_dp.call_args.kwargs["load_local_pixel_values"] diff --git a/test/registered/unit/models/test_kimi_vl.py b/test/registered/unit/models/test_kimi_vl.py index 1cfa38a1c..10ae42d7a 100644 --- a/test/registered/unit/models/test_kimi_vl.py +++ b/test/registered/unit/models/test_kimi_vl.py @@ -78,9 +78,12 @@ def _image_item(feature, grid_hws): class TestKimiVLEncoderParallelism(CustomTestCase): def test_moonvit_uses_tensor_parallel_layers(self): - with get_parallel().override( - tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0 - ), get_context().override_server_args(): + with ( + get_parallel().override( + tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0 + ), + get_context().override_server_args(), + ): layer = MoonVitEncoderLayer( num_heads=2, hidden_dim=8, diff --git a/test/registered/unit/models/test_qwen3_vl_feature_materialization.py b/test/registered/unit/models/test_qwen3_vl_feature_materialization.py index 091cfda41..b0cbccf75 100644 --- a/test/registered/unit/models/test_qwen3_vl_feature_materialization.py +++ b/test/registered/unit/models/test_qwen3_vl_feature_materialization.py @@ -160,15 +160,19 @@ class TestQwen3VLFeatureMaterialization(CustomTestCase): ) return encoded - with patch( - "sglang.srt.models.qwen3_vl.run_dp_sharded_mrope_vision_model", - side_effect=run_dp, - ), patch( - "sglang.srt.models.qwen3_vl.materialize_multimodal_features", - return_value=local_features, - ) as materialize, patch( - "sglang.srt.models.qwen3_vl.get_parallel", - return_value=SimpleNamespace(tp_size=8), + with ( + patch( + "sglang.srt.models.qwen3_vl.run_dp_sharded_mrope_vision_model", + side_effect=run_dp, + ), + patch( + "sglang.srt.models.qwen3_vl.materialize_multimodal_features", + return_value=local_features, + ) as materialize, + patch( + "sglang.srt.models.qwen3_vl.get_parallel", + return_value=SimpleNamespace(tp_size=8), + ), ): output = feature_method(items) diff --git a/test/registered/unit/observability/test_request_metrics_exporter.py b/test/registered/unit/observability/test_request_metrics_exporter.py index fe7de76da..21a99b42b 100644 --- a/test/registered/unit/observability/test_request_metrics_exporter.py +++ b/test/registered/unit/observability/test_request_metrics_exporter.py @@ -78,12 +78,12 @@ def setUpModule(): if stub_modules: if "sglang.srt.managers.io_struct" in stub_modules: - stub_modules["sglang.srt.managers.io_struct"].GenerateReqInput = ( - _GenerateReqInput - ) - stub_modules["sglang.srt.managers.io_struct"].EmbeddingReqInput = ( - _EmbeddingReqInput - ) + stub_modules[ + "sglang.srt.managers.io_struct" + ].GenerateReqInput = _GenerateReqInput + stub_modules[ + "sglang.srt.managers.io_struct" + ].EmbeddingReqInput = _EmbeddingReqInput if "sglang.srt.server_args" in stub_modules: stub_modules["sglang.srt.server_args"].ServerArgs = _ServerArgs diff --git a/test/registered/unit/parser/test_template_manager.py b/test/registered/unit/parser/test_template_manager.py index 497077137..de3f26f26 100644 --- a/test/registered/unit/parser/test_template_manager.py +++ b/test/registered/unit/parser/test_template_manager.py @@ -37,7 +37,6 @@ def _patch_hf_transformers_utils(get_tokenizer, get_config=None): class TestTemplateManagerReasoningDetection(unittest.TestCase): - def _detect(self, template, vocab): force, config = detect_reasoning_pattern(template) parser = detect_reasoning_parser( @@ -145,15 +144,13 @@ class TestTemplateManagerReasoningDetection(unittest.TestCase): # An explicit boolean=false second argument is equivalent to the # one-argument form. ( - "{%- set enable_thinking = enable_thinking" - " | default(true, false) -%}", + "{%- set enable_thinking = enable_thinking | default(true, false) -%}", True, ), # Boolean mode with a false default still maps False -> False and # True -> True, so it is a working default-off toggle. ( - "{%- set enable_thinking = enable_thinking" - " | default(false, true) -%}", + "{%- set enable_thinking = enable_thinking | default(false, true) -%}", False, ), ( diff --git a/test/registered/unit/sampling/test_custom_logit_processor.py b/test/registered/unit/sampling/test_custom_logit_processor.py index 2b706f671..2e94311ec 100644 --- a/test/registered/unit/sampling/test_custom_logit_processor.py +++ b/test/registered/unit/sampling/test_custom_logit_processor.py @@ -39,7 +39,6 @@ def _make_req(origin_input_ids=None, output_ids=None): # Serialization round-trip class TestCustomLogitProcessorSerialization(CustomTestCase): - def test_to_str_produces_valid_json(self): """Test that to_str() produces valid JSON with a 'callable' field.""" s = DisallowedTokensLogitsProcessor.to_str() diff --git a/test/registered/unit/sampling/test_penaltylib.py b/test/registered/unit/sampling/test_penaltylib.py index 39576702e..510c60d30 100644 --- a/test/registered/unit/sampling/test_penaltylib.py +++ b/test/registered/unit/sampling/test_penaltylib.py @@ -53,7 +53,6 @@ def _make_batch(reqs): # BatchedPenalizerOrchestrator class TestBatchedPenalizerOrchestrator(CustomTestCase): - def test_init_detects_required_penalizers(self): """Test that orchestrator marks is_required=True when any request has nonzero penalty.""" reqs = [_make_req(freq=1.0)] @@ -143,7 +142,6 @@ class TestBatchedPenalizerOrchestrator(CustomTestCase): # BatchedFrequencyPenalizer class TestBatchedFrequencyPenalizer(CustomTestCase): - def _setup(self, freq_values): reqs = [_make_req(freq=f) for f in freq_values] batch = _make_batch(reqs) @@ -225,7 +223,6 @@ class TestBatchedFrequencyPenalizer(CustomTestCase): # BatchedPresencePenalizer class TestBatchedPresencePenalizer(CustomTestCase): - def _setup(self, presence_values): reqs = [_make_req(presence=p) for p in presence_values] batch = _make_batch(reqs) @@ -275,7 +272,6 @@ class TestBatchedPresencePenalizer(CustomTestCase): # BatchedMinNewTokensPenalizer class TestBatchedMinNewTokensPenalizer(CustomTestCase): - def _setup(self, configs): """configs: list of (min_tokens, stop_ids, eos_id).""" reqs = [_make_req(min_tokens=c[0], stop_ids=c[1], eos_id=c[2]) for c in configs] @@ -388,7 +384,6 @@ class TestBatchedMinNewTokensPenalizer(CustomTestCase): # _BatchedPenalizer base class edge cases class TestBatchedPenalizerBase(CustomTestCase): - def test_filter_when_not_prepared_is_noop(self): """Test that filter on an unprepared penalizer does not crash.""" reqs = [_make_req()] @@ -452,7 +447,6 @@ class TestBatchedPenalizerBase(CustomTestCase): # Orchestrator with multiple penalizer types class TestOrchestratorMultiplePenalizers(CustomTestCase): - def test_all_three_penalizers(self): """Test orchestrator managing frequency, presence, and min_new_tokens together.""" reqs = [_make_req(freq=1.0, presence=0.5, min_tokens=2, eos_id=2)] diff --git a/test/registered/unit/sampling/test_sampling_batch_info.py b/test/registered/unit/sampling/test_sampling_batch_info.py index 59f7f3fe3..a5c3cef1e 100644 --- a/test/registered/unit/sampling/test_sampling_batch_info.py +++ b/test/registered/unit/sampling/test_sampling_batch_info.py @@ -50,7 +50,6 @@ def _serial_batched_fill(entries, vocab_mask): class TestMergeBiasTensor(CustomTestCase): - def test_both_none_returns_none(self): """Test that merging two None tensors returns None.""" result = merge_bias_tensor(None, None, 2, 3, DEVICE, 0.0) @@ -95,7 +94,6 @@ class TestMergeBiasTensor(CustomTestCase): # SamplingBatchInfo.__len__ class TestSamplingBatchInfoLen(CustomTestCase): - def test_len_matches_batch_size(self): """Test that __len__ returns batch size (number of temperature rows).""" info = _make_info(batch_size=5) @@ -103,7 +101,6 @@ class TestSamplingBatchInfoLen(CustomTestCase): class TestMergeCustomLogitProcessor(CustomTestCase): - def test_both_none_returns_none(self): """Test that merging two None processor dicts returns None.""" result = SamplingBatchInfo.merge_custom_logit_processor( @@ -150,7 +147,6 @@ class TestMergeCustomLogitProcessor(CustomTestCase): # apply_logits_bias class TestApplyLogitsBias(CustomTestCase): - def test_applies_additive_penalties(self): """Test that pre-accumulated additive penalties are added to logits.""" info = _make_info(batch_size=1) @@ -239,8 +235,8 @@ class TestApplyLogitsBias(CustomTestCase): def make_info(): grammar = MagicMock() - grammar.apply_vocab_mask.side_effect = ( - lambda logits, vocab_mask: logits.add_(vocab_mask) + grammar.apply_vocab_mask.side_effect = lambda logits, vocab_mask: ( + logits.add_(vocab_mask) ) info = _make_info(batch_size=1) info.acc_additive_penalties = torch.linspace( @@ -272,7 +268,6 @@ class TestApplyLogitsBias(CustomTestCase): # update_penalties class TestUpdatePenalties(CustomTestCase): - def test_required_creates_penalties_tensor(self): """Test that update_penalties allocates a zero tensor and calls orchestrator methods.""" orch = MagicMock(is_required=True) @@ -296,7 +291,6 @@ class TestUpdatePenalties(CustomTestCase): # update_regex_vocab_mask class TestUpdateRegexVocabMask(CustomTestCase): - def test_no_grammars_clears_mask(self): """Test that None grammars clears the grammar_mask.""" info = _make_info(batch_size=1) @@ -379,7 +373,6 @@ class TestUpdateRegexVocabMask(CustomTestCase): # filter_batch class TestFilterBatch(CustomTestCase): - def test_filter_keeps_correct_indices(self): """Test that filter retains rows at indices 0 and 2, dropping index 1.""" info = _make_info(batch_size=3) @@ -434,7 +427,6 @@ class TestFilterBatch(CustomTestCase): # merge_batch class TestMergeBatch(CustomTestCase): - def test_merge_concatenates_tensors(self): """Test that merge concatenates temperature tensors from both batches.""" info1 = _make_info(batch_size=2) @@ -513,7 +505,6 @@ class TestMergeBatch(CustomTestCase): # copy_for_forward class TestCopyForForward(CustomTestCase): - def test_returns_copy_without_orchestrator(self): """Test that copy_for_forward returns a copy with orchestrator set to None.""" orch = MagicMock(is_required=False) @@ -526,7 +517,6 @@ class TestCopyForForward(CustomTestCase): # from_schedule_batch class TestFromScheduleBatch(CustomTestCase): - def setUp(self): super().setUp() # from_schedule_batch reads these two flags from the exec bag; give diff --git a/test/registered/unit/scripted_runtime/test_background_http_poster.py b/test/registered/unit/scripted_runtime/test_background_http_poster.py index 12da761df..3f592d6c5 100644 --- a/test/registered/unit/scripted_runtime/test_background_http_poster.py +++ b/test/registered/unit/scripted_runtime/test_background_http_poster.py @@ -15,7 +15,6 @@ register_cpu_ci(est_time=8, suite="base-a-test-cpu") class _FakeResponse: - def __init__(self) -> None: self.read_called = False @@ -25,7 +24,6 @@ class _FakeResponse: class _FakePostCM: - def __init__(self, response: _FakeResponse) -> None: self._response = response @@ -37,7 +35,6 @@ class _FakePostCM: class _FakeSession: - def __init__(self) -> None: self.closed = False self.calls: list[tuple[str, object]] = [] @@ -52,7 +49,6 @@ class _FakeSession: class TestBackgroundHttpPosterLifecycle(CustomTestCase): - def test_init_starts_running_loop_on_daemon_thread(self): poster = BackgroundHttpPoster() self.addCleanup(poster.close) @@ -80,7 +76,6 @@ class TestBackgroundHttpPosterLifecycle(CustomTestCase): class TestBackgroundHttpPosterSubmitCoro(CustomTestCase): - def test_submit_coro_runs_on_background_loop_thread(self): poster = BackgroundHttpPoster() self.addCleanup(poster.close) @@ -134,7 +129,6 @@ class TestBackgroundHttpPosterSubmitCoro(CustomTestCase): class TestBackgroundHttpPosterEnsureSession(CustomTestCase): - def test_ensure_session_creates_reuses_then_recreates_when_closed(self): poster = BackgroundHttpPoster() self.addCleanup(poster.close) @@ -160,7 +154,6 @@ class TestBackgroundHttpPosterEnsureSession(CustomTestCase): class TestBackgroundHttpPosterPost(CustomTestCase): - def _run_on_loop(self, poster: BackgroundHttpPoster, coro) -> None: asyncio.run_coroutine_threadsafe(coro, poster._loop).result(timeout=5.0) diff --git a/test/registered/unit/scripted_runtime/test_http_server.py b/test/registered/unit/scripted_runtime/test_http_server.py index 5ca579743..97c76056f 100644 --- a/test/registered/unit/scripted_runtime/test_http_server.py +++ b/test/registered/unit/scripted_runtime/test_http_server.py @@ -28,7 +28,6 @@ _NO_REPLY = object() class _PairSocketHarness: - def __init__(self, *, reply: object = _NO_REPLY) -> None: self._ctx = zmq.Context() self.server_socket = self._ctx.socket(zmq.PAIR) @@ -62,7 +61,6 @@ class _PairSocketHarness: class _FakeProcess: - def __init__(self, *, alive: bool) -> None: self._alive = alive @@ -79,7 +77,6 @@ def _make_server(socket: zmq.Socket, process: _FakeProcess) -> ScriptedHttpServe class TestExecuteScriptReplyMatching(CustomTestCase): - def test_returns_on_script_succeeded(self): with _PairSocketHarness(reply=ScriptSucceeded()) as pair: server = _make_server(pair.server_socket, _FakeProcess(alive=True)) @@ -116,7 +113,6 @@ class TestExecuteScriptReplyMatching(CustomTestCase): class TestExecuteScriptNoReply(CustomTestCase): - def test_timeout_when_process_still_alive(self): with _PairSocketHarness() as pair: server = _make_server(pair.server_socket, _FakeProcess(alive=True)) @@ -135,7 +131,6 @@ class TestExecuteScriptNoReply(CustomTestCase): class TestExecuteScriptDirtyGuard(CustomTestCase): - def test_refuses_to_run_when_already_dirty(self): with _PairSocketHarness() as pair: server = _make_server(pair.server_socket, _FakeProcess(alive=True)) diff --git a/test/registered/unit/scripted_runtime/test_scheduler_hook.py b/test/registered/unit/scripted_runtime/test_scheduler_hook.py index 2044fa4aa..d2d7856c1 100644 --- a/test/registered/unit/scripted_runtime/test_scheduler_hook.py +++ b/test/registered/unit/scripted_runtime/test_scheduler_hook.py @@ -26,7 +26,6 @@ def _raising_gen(): class TestAdvanceGenerator(CustomTestCase): - def test_not_done_when_generator_yields(self): done, exc_tb = scheduler_hook._advance_generator(_yielding_gen()) diff --git a/test/registered/unit/scripted_runtime/test_scripted_runtime_utils.py b/test/registered/unit/scripted_runtime/test_scripted_runtime_utils.py index 29464e4c1..68e3e42fa 100644 --- a/test/registered/unit/scripted_runtime/test_scripted_runtime_utils.py +++ b/test/registered/unit/scripted_runtime/test_scripted_runtime_utils.py @@ -13,7 +13,6 @@ register_cpu_ci(est_time=2, suite="base-a-test-cpu") class TestResolveFn(CustomTestCase): - def test_resolves_top_level_function(self): self.assertIs(resolve_fn("json:dumps"), json.dumps) @@ -46,7 +45,6 @@ class TestResolveFn(CustomTestCase): class TestEnsureScriptImportable(CustomTestCase): - _FAKE_ENTRY = "/tmp/__scripted_runtime_ut_fake_sys_path__" def setUp(self): diff --git a/test/registered/unit/scripted_runtime/test_tokenizer_recv_proxy.py b/test/registered/unit/scripted_runtime/test_tokenizer_recv_proxy.py index 8adbe40f6..4eb2ba124 100644 --- a/test/registered/unit/scripted_runtime/test_tokenizer_recv_proxy.py +++ b/test/registered/unit/scripted_runtime/test_tokenizer_recv_proxy.py @@ -18,18 +18,15 @@ import unittest @dataclass class _ControlMsg: - tag: str = "flush" @dataclass class _StartReq: - rid: str class _FakeUnderlyingSocket: - def __init__(self) -> None: self._ready: deque = deque() self._scheduled: list[list] = [] @@ -62,7 +59,6 @@ def _is_start_req(rid: str): class TestScriptedTokenizerRecvProxyRecv(CustomTestCase): - def test_recv_pyobj_drains_then_pops_fifo(self): underlying = _FakeUnderlyingSocket() proxy = ScriptedTokenizerRecvProxy(underlying=underlying) @@ -88,7 +84,6 @@ class TestScriptedTokenizerRecvProxyRecv(CustomTestCase): class TestScriptedTokenizerRecvProxyWaitUntilArrived(CustomTestCase): - def _proxy_with_stale_control(self): underlying = _FakeUnderlyingSocket() proxy = ScriptedTokenizerRecvProxy(underlying=underlying) diff --git a/test/registered/unit/server_args/test_resolution_declarations.py b/test/registered/unit/server_args/test_resolution_declarations.py index 622f23f87..6eaaf7170 100644 --- a/test/registered/unit/server_args/test_resolution_declarations.py +++ b/test/registered/unit/server_args/test_resolution_declarations.py @@ -714,8 +714,7 @@ class TestResolutionDeclarations(CustomTestCase): current = getattr(server_args, name, None) if current != raw_input[name]: moved.append( - f"{shape} -> {name}: raw={raw_input[name]!r} " - f"field={current!r}" + f"{shape} -> {name}: raw={raw_input[name]!r} field={current!r}" ) self.assertEqual( moved, diff --git a/test/registered/unit/server_args/test_resolution_reads_the_declarations.py b/test/registered/unit/server_args/test_resolution_reads_the_declarations.py index ee1283186..5ebb8eea6 100644 --- a/test/registered/unit/server_args/test_resolution_reads_the_declarations.py +++ b/test/registered/unit/server_args/test_resolution_reads_the_declarations.py @@ -61,7 +61,9 @@ def _holders(fn): else ( annotation.id if isinstance(annotation, ast.Name) - else annotation.attr if isinstance(annotation, ast.Attribute) else None + else annotation.attr + if isinstance(annotation, ast.Attribute) + else None ) ) if text == "ServerArgs": diff --git a/test/registered/unit/spec/test_dflash_extra_buffer_lazy.py b/test/registered/unit/spec/test_dflash_extra_buffer_lazy.py index 53238dd38..d5ea7e7fc 100644 --- a/test/registered/unit/spec/test_dflash_extra_buffer_lazy.py +++ b/test/registered/unit/spec/test_dflash_extra_buffer_lazy.py @@ -120,11 +120,14 @@ class TestDflashVerifyRunsMambaTrackHook(CustomTestCase): calls.append("init_new") return fake_forward_batch - with mock.patch( - "sglang.srt.speculative.spec_utils.prepare_mamba_track_for_verify", - side_effect=fake_hook, - ), mock.patch.object( - dflash_info.ForwardBatch, "init_new", side_effect=fake_init_new + with ( + mock.patch( + "sglang.srt.speculative.spec_utils.prepare_mamba_track_for_verify", + side_effect=fake_hook, + ), + mock.patch.object( + dflash_info.ForwardBatch, "init_new", side_effect=fake_init_new + ), ): out, can_run_cuda_graph = self._spec_input().prepare_for_verify( batch, target_worker diff --git a/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py b/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py index 9587dda35..2ce9e2d89 100644 --- a/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py +++ b/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py @@ -231,12 +231,15 @@ class TestEagleWorkerV2BackendFallback(CustomTestCase): seq_lens=torch.ones((1,), dtype=torch.int32, device=DEVICE), ) - with patch( - "sglang.srt.speculative.eagle_worker_common.build_tree_kernel_efficient", - return_value=tree_result, - ), patch( - "sglang.srt.speculative.eagle_worker_v2.prepare_for_draft", - return_value=(forward_batch, True), + with ( + patch( + "sglang.srt.speculative.eagle_worker_common.build_tree_kernel_efficient", + return_value=tree_result, + ), + patch( + "sglang.srt.speculative.eagle_worker_v2.prepare_for_draft", + return_value=(forward_batch, True), + ), ): worker.draft(batch) diff --git a/test/registered/unit/spec/test_ngram_mamba_verify_update.py b/test/registered/unit/spec/test_ngram_mamba_verify_update.py index a620b5f2f..78d408309 100644 --- a/test/registered/unit/spec/test_ngram_mamba_verify_update.py +++ b/test/registered/unit/spec/test_ngram_mamba_verify_update.py @@ -194,12 +194,15 @@ class TestNgramMambaVerifyUpdate(CustomTestCase): dtype=torch.int32, ) - with patch( - "sglang.srt.speculative.spec_utils.mambaish_config", - return_value={"some": "config"}, - ), patch( - "sglang.srt.speculative.spec_utils.mamba_track_grid", - return_value=256, + with ( + patch( + "sglang.srt.speculative.spec_utils.mambaish_config", + return_value={"some": "config"}, + ), + patch( + "sglang.srt.speculative.spec_utils.mamba_track_grid", + return_value=256, + ), ): commit_mamba_states_after_verify( target_worker, diff --git a/test/registered/unit/state_capturer/test_routed_experts_scattered_a2a.py b/test/registered/unit/state_capturer/test_routed_experts_scattered_a2a.py index 37b558789..898591ab2 100644 --- a/test/registered/unit/state_capturer/test_routed_experts_scattered_a2a.py +++ b/test/registered/unit/state_capturer/test_routed_experts_scattered_a2a.py @@ -49,10 +49,11 @@ class TestGetLocalSliceBackendBranch(CustomTestCase): def test_deepep_v2_reads_buffer_head(self): cap, buf = self._capturer() - with mock.patch.object( - re_mod, "is_dp_attention_enabled", return_value=True - ), mock.patch.object( - re_mod, "get_moe_a2a_backend", return_value=MoeA2ABackend("deepep_v2") + with ( + mock.patch.object(re_mod, "is_dp_attention_enabled", return_value=True), + mock.patch.object( + re_mod, "get_moe_a2a_backend", return_value=MoeA2ABackend("deepep_v2") + ), ): out = self._slice(cap, n_local=5) self.assertTrue(torch.equal(out, buf[0:5, :, : self.K])) @@ -61,22 +62,23 @@ class TestGetLocalSliceBackendBranch(CustomTestCase): cap, _ = self._capturer() outs = [] for backend in ("deepep", "deepep_v2"): - with mock.patch.object( - re_mod, "is_dp_attention_enabled", return_value=True - ), mock.patch.object( - re_mod, "get_moe_a2a_backend", return_value=MoeA2ABackend(backend) + with ( + mock.patch.object(re_mod, "is_dp_attention_enabled", return_value=True), + mock.patch.object( + re_mod, "get_moe_a2a_backend", return_value=MoeA2ABackend(backend) + ), ): outs.append(self._slice(cap, n_local=7)) self.assertTrue(torch.equal(outs[0], outs[1])) def test_tp_moe_reads_global_offset(self): cap, buf = self._capturer() - with mock.patch.object( - re_mod, "is_dp_attention_enabled", return_value=True - ), mock.patch.object( - re_mod, "get_moe_a2a_backend", return_value=MoeA2ABackend("none") - ), mock.patch.object( - re_mod, "get_dp_local_slice_cpu", return_value=(6, 4) + with ( + mock.patch.object(re_mod, "is_dp_attention_enabled", return_value=True), + mock.patch.object( + re_mod, "get_moe_a2a_backend", return_value=MoeA2ABackend("none") + ), + mock.patch.object(re_mod, "get_dp_local_slice_cpu", return_value=(6, 4)), ): out = self._slice(cap, n_local=999) self.assertTrue(torch.equal(out, buf[6:10, :, : self.K])) diff --git a/test/registered/unit/test_cuda_vmm_utils.py b/test/registered/unit/test_cuda_vmm_utils.py index 27f13bd77..fa4b16a2d 100644 --- a/test/registered/unit/test_cuda_vmm_utils.py +++ b/test/registered/unit/test_cuda_vmm_utils.py @@ -104,9 +104,7 @@ def _byte(rank: int, chunk: int) -> int: def _assert_region(va: int, expected: int, peer: int, chunk: int) -> None: host = np.empty(16, dtype=np.uint8) check_drv(drv.cuMemcpyDtoH(host.ctypes.data, va, host.nbytes), "cuMemcpyDtoH") - assert ( - host == expected - ).all(), ( + assert (host == expected).all(), ( f"read {host.tolist()} from peer {peer} chunk {chunk}, expected all {expected}" ) diff --git a/test/registered/unit/test_dsa_tilelang_fp8_validation.py b/test/registered/unit/test_dsa_tilelang_fp8_validation.py index 212cd5a4f..e31dd99c5 100644 --- a/test/registered/unit/test_dsa_tilelang_fp8_validation.py +++ b/test/registered/unit/test_dsa_tilelang_fp8_validation.py @@ -15,7 +15,6 @@ register_cpu_ci(est_time=10, suite="base-a-test-cpu") class TestDsaTilelangFp8Validation(CustomTestCase): - def test_cuda_fp8_tilelang_decode_rejected(self): with self.assertRaises(ValueError): _check_tilelang_dsa_fp8_kv("fp8_e4m3", "flashmla_kv", "tilelang", hip=False) diff --git a/test/registered/unit/test_eval_accuracy_kit_sgl_eval.py b/test/registered/unit/test_eval_accuracy_kit_sgl_eval.py index 48d075073..65f19ff86 100644 --- a/test/registered/unit/test_eval_accuracy_kit_sgl_eval.py +++ b/test/registered/unit/test_eval_accuracy_kit_sgl_eval.py @@ -62,8 +62,9 @@ class TestEvalKitBackendDispatch(CustomTestCase): host.model = "m" for k, v in attrs.items(): setattr(host, k, v) - with patch.object(kit, "run_eval", side_effect=fake_run_eval), patch.object( - kit.requests, "get", side_effect=_fake_get + with ( + patch.object(kit, "run_eval", side_effect=fake_run_eval), + patch.object(kit.requests, "get", side_effect=_fake_get), ): host.test_gsm8k() return captured["args"] @@ -108,8 +109,9 @@ class TestEvalKitBackendDispatch(CustomTestCase): host.model = "deployment-model" host.mmmu_pro_score_threshold = 0.75 host.mmmu_pro_load_preset_from_model_id = "moonshotai/Kimi-K3" - with patch.object(kit, "run_eval", side_effect=fake_run_eval), patch.object( - kit.requests, "get", side_effect=_fake_get + with ( + patch.object(kit, "run_eval", side_effect=fake_run_eval), + patch.object(kit.requests, "get", side_effect=_fake_get), ): host.test_mmmu_pro() return captured["args"] diff --git a/test/registered/unit/test_fork_test_worker.py b/test/registered/unit/test_fork_test_worker.py index 653aa6cd2..3e1f795d5 100644 --- a/test/registered/unit/test_fork_test_worker.py +++ b/test/registered/unit/test_fork_test_worker.py @@ -37,27 +37,32 @@ class TestForkTestWorker(CustomTestCase): os.fdopen(result_read_fd) as result_stream, ): first = Path(tmpdir) / "first.py" - first.write_text(textwrap.dedent(""" + first.write_text( + textwrap.dedent(""" import builtins import os builtins._sglang_fork_worker_marker = 41 os.environ["SGLANG_FORK_WORKER_TEST"] = "leaked" raise SystemExit(0) - """)) + """) + ) second = Path(tmpdir) / "second.py" - second.write_text(textwrap.dedent(""" + second.write_text( + textwrap.dedent(""" import builtins import os assert not hasattr(builtins, "_sglang_fork_worker_marker") assert "SGLANG_FORK_WORKER_TEST" not in os.environ raise SystemExit(3) - """)) + """) + ) helper = Path(tmpdir) / "sibling_helper.py" helper.write_text("VALUE = 42\n") sibling_import = Path(tmpdir) / "sibling_import.py" - sibling_import.write_text(textwrap.dedent(""" + sibling_import.write_text( + textwrap.dedent(""" import os import sys @@ -65,7 +70,8 @@ class TestForkTestWorker(CustomTestCase): assert sys.path[0] == os.path.dirname(__file__) assert VALUE == 42 - """)) + """) + ) results = [] for filename in (first, second, sibling_import): diff --git a/test/registered/unit/test_global_config_read_ratchet.py b/test/registered/unit/test_global_config_read_ratchet.py index fe4c1c07f..7e745f9f3 100644 --- a/test/registered/unit/test_global_config_read_ratchet.py +++ b/test/registered/unit/test_global_config_read_ratchet.py @@ -375,8 +375,7 @@ class TestNoRenamedAccessorImports(CustomTestCase): base = imported.name.rsplit(".", 1)[-1] if base == "get_server_args": offenders.append( - f"{rel}:{node.lineno}: {imported.name} as " - f"{imported.asname}" + f"{rel}:{node.lineno}: {imported.name} as {imported.asname}" ) self.assertFalse( offenders, diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index 27e18bc76..3285905f8 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -2350,10 +2350,13 @@ class TestGoldenModelOverrides(_IsolatedPublish): args._model_config = SimpleNamespace(attention_arch=AttentionArch.MHA) return args - with override_platform(is_sm100=True), patch.object( - qwen3_5_module, - "get_default_attn_backend", - lambda server_args, **_: server_args.default_backend_for_test, + with ( + override_platform(is_sm100=True), + patch.object( + qwen3_5_module, + "get_default_attn_backend", + lambda server_args, **_: server_args.default_backend_for_test, + ), ): # radix on + no extra buffer + no spec -> page_size=1 path self.assertEqual( diff --git a/test/registered/unit/test_runtime_context_config_bags.py b/test/registered/unit/test_runtime_context_config_bags.py index 73a0226de..db90fb3a1 100644 --- a/test/registered/unit/test_runtime_context_config_bags.py +++ b/test/registered/unit/test_runtime_context_config_bags.py @@ -281,8 +281,11 @@ class TestRoleNamespaceEnforcement(CustomTestCase): def test_enforce_blocks_reads_outside_the_declared_set(self): self._publish("test") - with mock.patch.object(rc, "_ROLE_NS_MODE", "enforce"), mock.patch.dict( - rc.ROLE_NAMESPACE_SETS, {"test": frozenset({"serving", "schedule"})} + with ( + mock.patch.object(rc, "_ROLE_NS_MODE", "enforce"), + mock.patch.dict( + rc.ROLE_NAMESPACE_SETS, {"test": frozenset({"serving", "schedule"})} + ), ): rc.get_serving() rc.get_schedule() @@ -316,8 +319,9 @@ class TestRoleNamespaceEnforcement(CustomTestCase): def test_record_mode_collects_the_audit(self): self._publish("test") - with mock.patch.object(rc, "_ROLE_NS_MODE", "record"), mock.patch.object( - rc, "_RECORDED_NS_READS", set() + with ( + mock.patch.object(rc, "_ROLE_NS_MODE", "record"), + mock.patch.object(rc, "_RECORDED_NS_READS", set()), ): rc.get_exec() rc.get_disagg() @@ -343,8 +347,9 @@ class TestRoleNamespaceEnforcement(CustomTestCase): def test_record_mode_registers_the_exit_summary_at_publish(self): # A role that reads no bags must still emit its audit line; the exit # hook therefore registers at publish, not at the first read. - with mock.patch.object(rc, "_ROLE_NS_MODE", "record"), mock.patch.object( - rc, "_RECORD_DUMP_REGISTERED", False + with ( + mock.patch.object(rc, "_ROLE_NS_MODE", "record"), + mock.patch.object(rc, "_RECORD_DUMP_REGISTERED", False), ): self._publish("test") self.assertTrue(rc._RECORD_DUMP_REGISTERED) @@ -356,8 +361,9 @@ class TestRoleNamespaceEnforcement(CustomTestCase): import torch self._publish("test") - with mock.patch.object(rc, "_ROLE_NS_MODE", "record"), mock.patch.object( - rc, "_RECORDED_NS_READS", set() + with ( + mock.patch.object(rc, "_ROLE_NS_MODE", "record"), + mock.patch.object(rc, "_RECORDED_NS_READS", set()), ): @torch.compile(fullgraph=True, backend="eager", dynamic=False) diff --git a/test/registered/unit/test_server_args_migration.py b/test/registered/unit/test_server_args_migration.py index 67b2d9eb7..b1a94065d 100644 --- a/test/registered/unit/test_server_args_migration.py +++ b/test/registered/unit/test_server_args_migration.py @@ -17,7 +17,6 @@ register_cpu_ci(est_time=10, suite="base-a-test-cpu") class TestServerArgsAnnotatedCli(CustomTestCase): - @classmethod def setUpClass(cls): cls.parser = argparse.ArgumentParser() diff --git a/test/registered/unit/test_supplied_instance_exposure_ratchet.py b/test/registered/unit/test_supplied_instance_exposure_ratchet.py index 36f5a5460..3805f54ad 100644 --- a/test/registered/unit/test_supplied_instance_exposure_ratchet.py +++ b/test/registered/unit/test_supplied_instance_exposure_ratchet.py @@ -206,17 +206,17 @@ def _expanded_override_keys(rel, tree, call, kw) -> set: return values def dict_keys(node) -> set: - assert isinstance( - node, ast.Dict - ), f"non-literal dict in override expansion at {rel}:{call.lineno}" + assert isinstance(node, ast.Dict), ( + f"non-literal dict in override expansion at {rel}:{call.lineno}" + ) keys = set() for key in node.keys: if isinstance(key, ast.Constant): keys.add(key.value) continue - assert isinstance( - key, ast.Name - ), f"non-literal dict key in override expansion at {rel}:{call.lineno}" + assert isinstance(key, ast.Name), ( + f"non-literal dict key in override expansion at {rel}:{call.lineno}" + ) bound = loop_variable_values(key.id) assert bound, ( f"dict key {key.id!r} at {rel}:{call.lineno} is not bound by a " @@ -239,9 +239,9 @@ def _expanded_override_keys(rel, tree, call, kw) -> set: f"unresolvable override expansion at {rel}:{call.lineno}" ) return keys - assert isinstance( - kw.value, ast.Name - ), f"unresolvable override expansion at {rel}:{call.lineno}" + assert isinstance(kw.value, ast.Name), ( + f"unresolvable override expansion at {rel}:{call.lineno}" + ) name = kw.value.id enclosing = None for fn in ast.walk(tree): @@ -253,9 +253,9 @@ def _expanded_override_keys(rel, tree, call, kw) -> set: ): if enclosing is None or fn.lineno > enclosing.lineno: enclosing = fn - assert ( - enclosing is not None - ), f"override expansion outside any function at {rel}:{call.lineno}" + assert enclosing is not None, ( + f"override expansion outside any function at {rel}:{call.lineno}" + ) keys = set() found = False for node in ast.walk(enclosing): diff --git a/test/registered/unit/utils/test_profile_merger.py b/test/registered/unit/utils/test_profile_merger.py index 081f0a147..7df24afe9 100644 --- a/test/registered/unit/utils/test_profile_merger.py +++ b/test/registered/unit/utils/test_profile_merger.py @@ -201,7 +201,6 @@ class TestProfileMerger(CustomTestCase): class TestProfileMergerIntegration(CustomTestCase): - def test_data_structures_merge_profiles(self): # Test ProfileReq req = ProfileReq() diff --git a/test/registered/unit/utils/test_weight_checker.py b/test/registered/unit/utils/test_weight_checker.py index e9d893a9d..8924e32e9 100644 --- a/test/registered/unit/utils/test_weight_checker.py +++ b/test/registered/unit/utils/test_weight_checker.py @@ -59,9 +59,9 @@ def _assert_entries_close( """Compare two streams of (name, should_compare, ComparableWeight).""" actual_list: List[CheckEntry] = list(actual) expected_list: List[CheckEntry] = list(expected) - assert len(actual_list) == len( - expected_list - ), f"length mismatch: actual={len(actual_list)} expected={len(expected_list)}" + assert len(actual_list) == len(expected_list), ( + f"length mismatch: actual={len(actual_list)} expected={len(expected_list)}" + ) for i, ((a_name, a_flag, a_ref), (e_name, e_flag, e_ref)) in enumerate( zip(actual_list, expected_list) ): @@ -150,7 +150,6 @@ class _FakeModelRunner: class TestRandomLike(CustomTestCase): - def test_floating_point_preserves_dtype_shape_device(self): for dtype in (torch.float32, torch.float16, torch.bfloat16): t = torch.zeros(8, 4, dtype=dtype) @@ -209,7 +208,6 @@ class TestRandomLike(CustomTestCase): class TestPostprocessTensors(CustomTestCase): - # --- non-quant / non-skip --- def test_no_quant_yields_raw_with_should_compare_true(self): @@ -312,7 +310,6 @@ class TestPostprocessTensors(CustomTestCase): class TestCheckTensors(CustomTestCase): - def test_passes_when_all_equal(self): t = torch.ones(2, 2) expect = [ @@ -388,7 +385,6 @@ def _quantize_block_fp8(weight: torch.Tensor, scale_margin: float): class TestCheckTensorsAllowQuantError(CustomTestCase): - def setUp(self): torch.manual_seed(0) weight = torch.randn(256, 256, device="cuda") * 0.02 @@ -442,7 +438,6 @@ class TestCheckTensorsAllowQuantError(CustomTestCase): class TestBuildQuantizedSet(CustomTestCase): - def test_fp8_block_module_pairs_weight_and_scale(self): from sglang.srt.layers.quantization.fp8 import Fp8LinearMethod @@ -493,7 +488,6 @@ class _WeightCheckerTestBase(CustomTestCase): class TestSnapshot(_WeightCheckerTestBase): - def test_captures_params_and_buffers(self): self.checker._snapshot() keys = set(self.checker._snapshot_tensors.keys()) @@ -519,7 +513,6 @@ class TestSnapshot(_WeightCheckerTestBase): class TestResetTensors(_WeightCheckerTestBase): - def test_changes_normal_params_in_place(self): before_w = self.model.w.clone() before_w_ptr = self.model.w.data_ptr() @@ -547,7 +540,6 @@ class TestResetTensors(_WeightCheckerTestBase): class TestCompare(_WeightCheckerTestBase): - def test_without_snapshot_raises(self): with self.assertRaises(AssertionError): self.checker._compare() @@ -586,7 +578,6 @@ class TestCompare(_WeightCheckerTestBase): class TestHandle(_WeightCheckerTestBase): - def test_routes_to_actions(self): with ( patch.object(self.checker, "_snapshot") as m_snap, @@ -628,7 +619,6 @@ class TestHandle(_WeightCheckerTestBase): class TestIsNonPersistentBufferName(CustomTestCase): - def test_matches_cos_sin_cache_substring(self): self.assertTrue( _is_non_persistent_buffer_name("model.rotary_emb.cos_sin_cache") @@ -651,7 +641,6 @@ class TestIsNonPersistentBufferName(CustomTestCase): class TestHashTensor(CustomTestCase): - def test_stable_for_same_input(self): t = torch.arange(64, dtype=torch.float32).cuda() self.assertEqual(_hash_tensor(t), _hash_tensor(t.clone())) @@ -680,7 +669,6 @@ class TestHashTensor(CustomTestCase): class _ChecksumTestBase(CustomTestCase): - def setUp(self): torch.manual_seed(0) self.model = _TinyModel().cuda() @@ -699,7 +687,6 @@ class _ChecksumTestBase(CustomTestCase): class TestComputeChecksum(_ChecksumTestBase): - def test_returns_dict_with_expected_top_level_keys(self): out = self.checker._compute_checksum() self.assertEqual( diff --git a/test/registered/unit/utils/test_weight_checker_comparator.py b/test/registered/unit/utils/test_weight_checker_comparator.py index 864d5e87d..5b59c75e7 100644 --- a/test/registered/unit/utils/test_weight_checker_comparator.py +++ b/test/registered/unit/utils/test_weight_checker_comparator.py @@ -62,7 +62,6 @@ def _build_fp8_quant_pair(device: str = "cuda"): class TestQuantUlp(CustomTestCase): - def test_matches_bruteforce_spacing_for_fp8(self): for dtype in (torch.float8_e4m3fn, torch.float8_e5m2): all_bits = torch.arange(256, dtype=torch.uint8).view(dtype) @@ -183,7 +182,6 @@ class TestCompareQuantPair(CustomTestCase): class TestSelectComparableWeight(CustomTestCase): - def test_returns_none_when_not_a_quant_method(self): self.assertIsNone(select_comparable_weight(None)) diff --git a/test/registered/utils/test_bench_typebaseddispatcher.py b/test/registered/utils/test_bench_typebaseddispatcher.py index dc0365a6f..d269b080e 100644 --- a/test/registered/utils/test_bench_typebaseddispatcher.py +++ b/test/registered/utils/test_bench_typebaseddispatcher.py @@ -118,10 +118,10 @@ def benchmark_with_inheritance(): cached_call_times.append(time_taken) print( - f"First call (with MRO lookup): {sum(first_call_times)/len(first_call_times):.6f}s avg" + f"First call (with MRO lookup): {sum(first_call_times) / len(first_call_times):.6f}s avg" ) - print(f"Cached call: {sum(cached_call_times)/len(cached_call_times):.6f}s avg") - print(f"Caching improvement: {sum(first_call_times)/sum(cached_call_times):.2f}x") + print(f"Cached call: {sum(cached_call_times) / len(cached_call_times):.6f}s avg") + print(f"Caching improvement: {sum(first_call_times) / sum(cached_call_times):.2f}x") def benchmark_dispatchers(): @@ -158,8 +158,8 @@ def benchmark_dispatchers(): print(f"for list: {list_time:.4f} s") print(f"for dict: {dict_time:.4f} s") - print(f"improvement: {list_time/dict_time:.2f} x") - print(f"time reduce: {(1-dict_time/list_time) * 100:.1f} %") + print(f"improvement: {list_time / dict_time:.2f} x") + print(f"time reduce: {(1 - dict_time / list_time) * 100:.1f} %") def test_memory_usage(): @@ -252,7 +252,7 @@ def simulate_real_workload(): print(f"list version: {list_time:.4f} s") print(f"dict version: {dict_time:.4f} s") - print(f"improvement: {list_time/dict_time:.2f} x") + print(f"improvement: {list_time / dict_time:.2f} x") if __name__ == "__main__": diff --git a/test/registered/utils/test_numa_utils.py b/test/registered/utils/test_numa_utils.py index d9f8261f8..25d060d64 100644 --- a/test/registered/utils/test_numa_utils.py +++ b/test/registered/utils/test_numa_utils.py @@ -196,12 +196,15 @@ class TestGetNumaNodeIfAvailable(unittest.TestCase): def test_auto_bind_disabled_skips_numa_detection(self, mock_avail, mock_query): args = self._make_server_args(numa_node=None) for bind_v2 in ("0", "1"): - with self.subTest(bind_v2=bind_v2), patch.dict( - os.environ, - { - "SGLANG_AUTO_NUMA_BIND": "0", - "SGLANG_NUMA_BIND_V2": bind_v2, - }, + with ( + self.subTest(bind_v2=bind_v2), + patch.dict( + os.environ, + { + "SGLANG_AUTO_NUMA_BIND": "0", + "SGLANG_NUMA_BIND_V2": bind_v2, + }, + ), ): self.assertIsNone(get_numa_node_if_available(args, 0)) mock_avail.assert_not_called() diff --git a/test/registered/vlm/test_vision_chunked_prefill.py b/test/registered/vlm/test_vision_chunked_prefill.py index 0f791decc..ba8279e42 100644 --- a/test/registered/vlm/test_vision_chunked_prefill.py +++ b/test/registered/vlm/test_vision_chunked_prefill.py @@ -38,7 +38,6 @@ logger = logging.getLogger(__name__) class TestVisionChunkedPrefill(CustomTestCase): - def prepare_video_messages(self, video_path, max_frames_num=8): from sglang.srt.utils.video_decoder import VideoDecoderWrapper @@ -189,12 +188,12 @@ class TestVisionChunkedPrefill(CustomTestCase): try: outputs_chunked = [] for i, (batch, num_frame) in enumerate(zip(batches, num_frames)): - logger.info(f"Chunked test iteration {i+1}/{len(batches)}") + logger.info(f"Chunked test iteration {i + 1}/{len(batches)}") output_chunked = self.generate_for_video( batch=batch, num_frame=num_frame ) outputs_chunked += [output_chunked] - logger.info(f"Chunked test iteration {i+1} completed") + logger.info(f"Chunked test iteration {i + 1} completed") finally: logger.info(f"Killing chunked server pid={chunked_server_pid}") kill_process_tree(chunked_server_pid) @@ -208,12 +207,12 @@ class TestVisionChunkedPrefill(CustomTestCase): logger.info(f"Non-chunked server started with pid={no_chunked_server_pid}") outputs_no_chunked = [] for i, (batch, num_frame) in enumerate(zip(batches, num_frames)): - logger.info(f"Non-chunked test iteration {i+1}/{len(batches)}") + logger.info(f"Non-chunked test iteration {i + 1}/{len(batches)}") output_no_chunked = self.generate_for_video( batch=batch, num_frame=num_frame ) outputs_no_chunked += [output_no_chunked] - logger.info(f"Non-chunked test iteration {i+1} completed") + logger.info(f"Non-chunked test iteration {i + 1} completed") finally: logger.info(f"Killing non-chunked server pid={no_chunked_server_pid}") diff --git a/test/registered/vlm/test_vision_openai_server_a.py b/test/registered/vlm/test_vision_openai_server_a.py index 8d71e351d..948733380 100644 --- a/test/registered/vlm/test_vision_openai_server_a.py +++ b/test/registered/vlm/test_vision_openai_server_a.py @@ -269,9 +269,9 @@ class TestDeepseekOCRServer(TestOpenAIMLLMServerBase): import re coord_pattern = r"\[\[[\d\s,]+\]\]" - assert re.search( - coord_pattern, text - ), f"OCR text: {text}, should contain coordinate format [[x1, y1, x2, y2]]" + assert re.search(coord_pattern, text), ( + f"OCR text: {text}, should contain coordinate format [[x1, y1, x2, y2]]" + ) # Verify basic response fields assert response.id diff --git a/test/registered/xpu/test_intel_xpu_backend.py b/test/registered/xpu/test_intel_xpu_backend.py index 0ef8e088f..dc884ae9e 100644 --- a/test/registered/xpu/test_intel_xpu_backend.py +++ b/test/registered/xpu/test_intel_xpu_backend.py @@ -57,7 +57,6 @@ def intel_xpu_benchmark( class TestIntelXPUBackend(CustomTestCase): - @intel_xpu_benchmark(min_throughput=10, mem_fraction_static="0.3") def test_latency_qwen_model(self): return DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN diff --git a/test/registered/xpu/test_moe_ld_padding.py b/test/registered/xpu/test_moe_ld_padding.py index 92cefa48a..46cadce79 100644 --- a/test/registered/xpu/test_moe_ld_padding.py +++ b/test/registered/xpu/test_moe_ld_padding.py @@ -88,11 +88,14 @@ class TestXpuMoeLdPadding(CustomTestCase): ) return layer.w13_weight, layer.w2_weight - with unittest.mock.patch( - "sglang.srt.layers.quantization.unquant.is_xpu", return_value=True - ), unittest.mock.patch( - "sglang.srt.layers.quantization.unquant.get_moe_runner_backend", - return_value=MoeRunnerBackend.AUTO, + with ( + unittest.mock.patch( + "sglang.srt.layers.quantization.unquant.is_xpu", return_value=True + ), + unittest.mock.patch( + "sglang.srt.layers.quantization.unquant.get_moe_runner_backend", + return_value=MoeRunnerBackend.AUTO, + ), ): # Backend on but building for CPU -> must stay contiguous. w13_cpu, w2_cpu = build("cpu") @@ -107,11 +110,14 @@ class TestXpuMoeLdPadding(CustomTestCase): self.assertTrue(w13_triton.is_contiguous()) # Backend forced to Triton -> never padded, even on XPU. - with unittest.mock.patch( - "sglang.srt.layers.quantization.unquant.is_xpu", return_value=True - ), unittest.mock.patch( - "sglang.srt.layers.quantization.unquant.get_moe_runner_backend", - return_value=MoeRunnerBackend.TRITON, + with ( + unittest.mock.patch( + "sglang.srt.layers.quantization.unquant.is_xpu", return_value=True + ), + unittest.mock.patch( + "sglang.srt.layers.quantization.unquant.get_moe_runner_backend", + return_value=MoeRunnerBackend.TRITON, + ), ): device = "xpu" if torch.xpu.is_available() else "cpu" w13, w2 = build(device) diff --git a/test/registered/xpu/test_topk.py b/test/registered/xpu/test_topk.py index bf3365e44..26b2f7ce1 100644 --- a/test/registered/xpu/test_topk.py +++ b/test/registered/xpu/test_topk.py @@ -182,7 +182,6 @@ class TestBiasedGroupedTopK(CustomTestCase): for M in num_tokens: for num_fused_shared_experts in num_fused_shared_experts_list: - topk_routed = topk_value - num_fused_shared_experts hidden_states = torch.randn(M, 100, dtype=torch.bfloat16, device=device) gating_output = torch.randn(M, E_num, dtype=gating_dtype, device=device) @@ -243,7 +242,6 @@ class TestBiasedGroupedTopK(CustomTestCase): for M in num_tokens: for num_fused_shared_experts in num_fused_shared_experts_list: - topk_routed = topk_value - num_fused_shared_experts hidden_states = torch.randn(M, 100, dtype=torch.bfloat16, device=device) gating_output = torch.randn(M, E_num, dtype=gating_dtype, device=device) @@ -301,7 +299,6 @@ class TestBiasedGroupedTopK(CustomTestCase): for M in num_tokens: for scoring_func in scoring_func_list: for num_fused_shared_experts in num_fused_shared_experts_list: - topk_routed = topk_value - num_fused_shared_experts hidden_states = torch.randn( M, 100, dtype=gating_dtype, device=device diff --git a/test/registered/xpu/test_triton_attention_backend.py b/test/registered/xpu/test_triton_attention_backend.py index ba0ea1688..2a26ad9a2 100644 --- a/test/registered/xpu/test_triton_attention_backend.py +++ b/test/registered/xpu/test_triton_attention_backend.py @@ -51,7 +51,6 @@ def triton_attention_benchmark(extra_args=None, mem_fraction_static="0.84"): class TestTritonAttentionBackend(CustomTestCase): - @triton_attention_benchmark( [ "--json-model-override-args", diff --git a/test/registered/xpu/test_xpu_basic.py b/test/registered/xpu/test_xpu_basic.py index c1e4c30ea..c79a39742 100644 --- a/test/registered/xpu/test_xpu_basic.py +++ b/test/registered/xpu/test_xpu_basic.py @@ -22,7 +22,6 @@ register_xpu_ci(est_time=300, suite="stage-a-test-1-gpu-xpu") class TestXPUBasic(CustomTestCase): - def test_basic_generation(self): """Server starts on XPU and completes at least one decode step.""" args = [