diff --git a/.buildkite/check-torch-abi.py b/.buildkite/check-torch-abi.py index 37e2b8dadb..5bfec8abd0 100644 --- a/.buildkite/check-torch-abi.py +++ b/.buildkite/check-torch-abi.py @@ -17,7 +17,6 @@ ALLOWED_UNSTABLE_LIBRARIES: tuple[str, ...] = ( # _C_stable_libtorch, but these two modules still bind libtorch directly. "_C.abi3.so", "_rocm_C.abi3.so", - "third_party/deep_gemm/_C*.so", ) diff --git a/.buildkite/intel_jobs/model_executor_intel.yaml b/.buildkite/intel_jobs/model_executor_intel.yaml index 2b31c222fd..3c9b80cbc2 100644 --- a/.buildkite/intel_jobs/model_executor_intel.yaml +++ b/.buildkite/intel_jobs/model_executor_intel.yaml @@ -28,7 +28,7 @@ steps: export VLLM_WORKER_MULTIPROC_METHOD=spawn && export PYTHONFAULTHANDLER=1 && cd tests && - pytest -v -s model_executor -m "not slow_test" --ignore="model_executor/layers/test_rocm_unquantized_gemm.py" --ignore="model_executor/model_loader/tensorizer_loader/test_tensorizer.py" --ignore="model_executor/model_loader/test_reload.py"' + pytest -v -s model_executor -m "not slow_test" --ignore="model_executor/layers/test_rocm_unquantized_gemm.py" --ignore="model_executor/model_loader/tensorizer_loader/test_tensorizer.py" --ignore="model_executor/model_loader/test_reload.py" --deselect="tests/model_executor/test_enabled_custom_ops.py::test_topk_gating_enabled_by_cdna"' - label: ":intel: (BMG) Tensorizer Model Loader Tests (XPU)" key: tensorizer-model-loader-tests-xpu diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 39ac57b8af..7a9876032d 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -3439,6 +3439,26 @@ steps: commands: - pytest -v -s tests/kernels/test_rocm_gemm_rs_ar.py +- label: ":amd: (MI355) Qwen4 Exp SP" + timeout_in_minutes: 30 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/config/parallel.py + - vllm/models/common/ops/sequence_parallel.py + - vllm/models/qwen4_exp/ + - vllm/model_executor/layers/vocab_parallel_embedding.py + - vllm/model_executor/models/qwen3_next.py + - vllm/platforms/rocm.py + - tests/models/qwen4_exp/test_ple.py + - tests/models/qwen4_exp/test_tpsp.py + commands: + - pytest -v -s models/qwen4_exp/test_tpsp.py + - label: ":amd: (MI355) Distributed DP Basic" # TBD timeout_in_minutes: 45 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] @@ -4684,16 +4704,32 @@ steps: kernels/test_fused_inv_rope_fp8_quant.py kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py kernels/test_fused_quant_activation.py - kernels/test_gate_linear_rocm_dispatch.py kernels/test_gemm_rs_ar.py kernels/test_kpool_decode_update_batched.py kernels/test_minimax_m3_amd_ops.py kernels/test_minimax_m3_sparse_attn_fp8_scale.py kernels/test_onednn.py kernels/test_relu2_fp8_quant.py - kernels/test_rocm_fp32_router_gemm.py --deselect=tests/kernels/test_bf16_skinny_gemm.py::test_every_dsv3_routed_shape_is_instantiated +- label: ":amd: (MI355) Gate Linear Router Kernels" + timeout_in_minutes: 20 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + optional: true + working_dir: "/vllm-workspace/" + source_file_dependencies: + - vllm/model_executor/layers/fused_moe/router/gate_linear.py + - vllm/model_executor/layers/fused_moe/router/rocm_fp32_router_gemm.py + - vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_rocm.py + - vllm/platforms/rocm.py + - tests/kernels/test_rocm_fp32_router_gemm.py + - tests/kernels/test_gate_linear_rocm_dispatch.py + - tests/kernels/test_bf16x3_router_gemm_rocm.py + commands: + - pytest -v -s tests/kernels/test_rocm_fp32_router_gemm.py tests/kernels/test_gate_linear_rocm_dispatch.py tests/kernels/test_bf16x3_router_gemm_rocm.py + - label: ":amd: (MI355 DPX) Attention DiffKV Kernels" timeout_in_minutes: 35 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index d96071dfd0..5a982efc3c 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -14,6 +14,7 @@ steps: - tests/distributed commands: - pytest -v -s distributed/test_comm_ops.py + - pytest -v -s distributed/test_cpu_all_reduce.py - pytest -v -s distributed/test_shm_broadcast.py - pytest -v -s distributed/test_shm_buffer.py - pytest -v -s distributed/test_shm_storage.py @@ -109,6 +110,7 @@ steps: amd: label: ":amd: (MI355) Distributed DP Basic" dind: false + automatic_shard: true device: mi355_2 timeout_in_minutes: 45 depends_on: @@ -331,6 +333,7 @@ steps: amd: label: ':amd: (MI355) Distributed DP Extended' dind: false + automatic_shard: true device: mi355_4 timeout_in_minutes: 60 working_dir: /vllm-workspace/tests @@ -637,7 +640,7 @@ steps: - tests/distributed/ - examples/features/data_parallel/data_parallel_offline.py commands: - - ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 $IMAGE_TAG "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/features/data_parallel/data_parallel_offline.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=0 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py" "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/features/data_parallel/data_parallel_offline.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=1 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code" + - ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 $IMAGE_TAG "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/features/data_parallel/data_parallel_offline.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=0 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py" "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/features/data_parallel/data_parallel_offline.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=1 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code" - label: ":nvidia: (L4) Pipeline + Context Parallelism" key: pipeline-context-parallelism-4-gpus @@ -706,6 +709,7 @@ steps: amd: label: ':amd: (MI355) RayExecutorV2' dind: false + automatic_shard: true device: mi355_4 timeout_in_minutes: 75 working_dir: /vllm-workspace/tests diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index 4d8bfdc55f..ae59d3cb53 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -41,6 +41,7 @@ steps: amd: label: ":amd: (MI355 DPX) Entrypoints Unit" dind: false + automatic_shard: true device: mi355_dpx timeout_in_minutes: 50 depends_on: @@ -207,6 +208,7 @@ steps: amd: label: ":amd: (MI355 DPX) Entrypoints Integration (Speech to Text)" dind: false + automatic_shard: true device: mi355_dpx timeout_in_minutes: 60 working_dir: "/vllm-workspace/tests" @@ -252,6 +254,7 @@ steps: amd: label: ":amd: (MI355 DPX) Entrypoints Integration (Pooling)" dind: false + automatic_shard: true device: mi355_dpx timeout_in_minutes: 65 working_dir: "/vllm-workspace/tests" diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 126fff9bdc..12b8516926 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -138,11 +138,15 @@ steps: - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu - vllm/models/deepseek_v4/common/ops/ - vllm/models/deepseek_v4/nvidia/ + - vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py + - vllm/model_executor/layers/fused_moe/prepare_megamoe.py - vllm/utils/deep_gemm.py - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py + - tests/kernels/moe/test_deep_gemm_mega_moe.py - tests/models/test_deepseek_v4_mega_moe.py commands: - pytest -v -s kernels/test_fused_deepseek_v4_*.py + - pytest -v -s kernels/moe/test_deep_gemm_mega_moe.py - pytest -v -s models/test_deepseek_v4_mega_moe.py mirror: amd: @@ -286,7 +290,6 @@ steps: kernels/test_apply_repetition_penalties.py kernels/test_awq_int4_to_int8.py kernels/test_bf16_skinny_gemm.py - kernels/test_bf16x3_router_gemm_rocm.py kernels/test_cache_kernels.py kernels/test_compressor_kv_cache.py kernels/test_cp_gather_fp8.py @@ -301,14 +304,12 @@ steps: kernels/test_fused_inv_rope_fp8_quant.py kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py kernels/test_fused_quant_activation.py - kernels/test_gate_linear_rocm_dispatch.py kernels/test_gemm_rs_ar.py kernels/test_kpool_decode_update_batched.py kernels/test_minimax_m3_amd_ops.py kernels/test_minimax_m3_sparse_attn_fp8_scale.py kernels/test_onednn.py kernels/test_relu2_fp8_quant.py - kernels/test_rocm_fp32_router_gemm.py --deselect=tests/kernels/test_bf16_skinny_gemm.py::test_every_dsv3_routed_shape_is_instantiated - label: ":nvidia: (L4) Attention Kernels Shard %N" @@ -528,6 +529,7 @@ steps: amd: label: ":amd: (MI250) Mamba Kernels" dind: false + automatic_shard: true device: mi250_1 timeout_in_minutes: 55 depends_on: @@ -720,6 +722,24 @@ steps: - nvidia-smi - pytest -v -s tests/kernels/test_ll_bf16_gemm.py - pytest -v -s tests/kernels/test_ll_fp32w_gemm.py + mirror: + amd: + label: ":amd: (MI355) Gate Linear Router Kernels" + dind: false + device: mi355_1 + timeout_in_minutes: 20 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/model_executor/layers/fused_moe/router/gate_linear.py + - vllm/model_executor/layers/fused_moe/router/rocm_fp32_router_gemm.py + - vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_rocm.py + - vllm/platforms/rocm.py + - tests/kernels/test_rocm_fp32_router_gemm.py + - tests/kernels/test_gate_linear_rocm_dispatch.py + - tests/kernels/test_bf16x3_router_gemm_rocm.py + commands: + - pytest -v -s tests/kernels/test_rocm_fp32_router_gemm.py tests/kernels/test_gate_linear_rocm_dispatch.py tests/kernels/test_bf16x3_router_gemm_rocm.py - label: ":nvidia: (DGX) Spark b12x Linear Kernels Nightly" key: b12x-linear-kernels-dgx-spark-nightly diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 64f5d8047f..29e0582ec5 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -32,6 +32,7 @@ steps: amd: label: ":amd: (MI355 DPX) V1 Sample" dind: false + automatic_shard: true device: mi355_dpx timeout_in_minutes: 70 depends_on: @@ -219,6 +220,7 @@ steps: - pytest -v -s -m 'cpu_test' v1/core - pytest -v -s v1/structured_output - pytest -v -s v1/test_serial_utils.py + - pytest -v -s v1/worker/test_cpu_worker_kv_budget.py - pytest -v -s v1/test_kv_cache_spec_registry.py - pytest -v -s v1/cudagraph/test_cudagraph_manager.py - pytest -v -s -m 'cpu_test' v1/kv_connector/unit @@ -707,6 +709,7 @@ steps: amd: label: ":amd: (MI300) Batch Invariance" dind: false + automatic_shard: true device: mi300_1 timeout_in_minutes: 60 depends_on: diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index ebc17c1cce..029ccbbffa 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -69,10 +69,11 @@ steps: - tests/models/test_deepseek_v4_vl_rocm.py - tests/models/test_deepseek_v41_decoder_replay_layers.py - tests/models/test_deepseek_v41_replay_batch.py + - tests/models/test_deepseek_v41_pcp.py - tests/models/transformers/test_layer_registry_e2e.py - tests/models/test_deepseek_v41_replay_start.py commands: - - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py + - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_pcp.py mirror: amd: label: ":amd: (MI355) Basic Models (Other)" diff --git a/.buildkite/test_areas/models_distributed.yaml b/.buildkite/test_areas/models_distributed.yaml index 8d2bc5dbaf..38a2e7c09f 100644 --- a/.buildkite/test_areas/models_distributed.yaml +++ b/.buildkite/test_areas/models_distributed.yaml @@ -18,6 +18,22 @@ steps: - tests/models/qwen4_exp/test_tpsp.py commands: - pytest -v -s models/qwen4_exp/test_tpsp.py + mirror: + amd: + label: ":amd: (MI355) Qwen4 Exp SP" + dind: false + device: mi355_4 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/config/parallel.py + - vllm/models/common/ops/sequence_parallel.py + - vllm/models/qwen4_exp/ + - vllm/model_executor/layers/vocab_parallel_embedding.py + - vllm/model_executor/models/qwen3_next.py + - vllm/platforms/rocm.py + - tests/models/qwen4_exp/test_ple.py + - tests/models/qwen4_exp/test_tpsp.py - label: ":nvidia: (L4) Distributed Models Shard %N" key: distributed-model-tests-2-gpus diff --git a/.buildkite/test_areas/pytorch.yaml b/.buildkite/test_areas/pytorch.yaml index 60061af414..1dc6659e8e 100644 --- a/.buildkite/test_areas/pytorch.yaml +++ b/.buildkite/test_areas/pytorch.yaml @@ -160,6 +160,7 @@ steps: amd: label: ':amd: (MI300) PyTorch Compilation Passes' dind: false + automatic_shard: true device: mi300_1 timeout_in_minutes: 80 working_dir: /vllm-workspace/tests diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml index cdf8e7bbe7..7c9ddf9438 100644 --- a/.buildkite/test_areas/quantization.yaml +++ b/.buildkite/test_areas/quantization.yaml @@ -11,6 +11,7 @@ steps: - csrc/ - vllm/model_executor/layers/quantization - vllm/models/deepseek_v41/quant_config.py + - vllm/utils/humming.py - tests/quantization commands: # Pin torchao to the version compatible with the CI PyTorch and CUDA stack. @@ -136,6 +137,7 @@ steps: amd: label: ":amd: (MI355 DPX) Quantized Models" dind: false + automatic_shard: true device: mi355_dpx timeout_in_minutes: 80 working_dir: "/vllm-workspace/tests" diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index eceb914c86..c2742b4c51 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -28,6 +28,7 @@ steps: amd: label: ":amd: (MI300) V1 Spec Decode" dind: false + automatic_shard: true device: mi300_1 timeout_in_minutes: 50 depends_on: diff --git a/benchmarks/kernels/benchmark_fused_topk.py b/benchmarks/kernels/benchmark_fused_topk.py index 72bf2d97cc..ab4d954c87 100644 --- a/benchmarks/kernels/benchmark_fused_topk.py +++ b/benchmarks/kernels/benchmark_fused_topk.py @@ -11,7 +11,7 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser num_tokens_range = [2**i for i in range(0, 8, 2)] num_experts_range = [16, 32, 64, 128, 256, 512] -topk_range = [3, 4] +topk_range = [3, 4, 8, 10] configs = list(itertools.product(num_tokens_range, num_experts_range, topk_range)) @@ -39,9 +39,9 @@ def get_benchmark(scoring_func): x_names=["num_tokens", "num_experts", "topk"], x_vals=[list(_) for _ in configs], line_arg="provider", - line_vals=["torch", "vllm"], - line_names=["Torch", "vLLM"], - styles=[("blue", "-"), ("red", "-")], + line_vals=["torch", "vllm", "aiter_gating"], + line_names=["Torch", "vLLM", "AITER topk_gating"], + styles=[("blue", "-"), ("red", "-"), ("green", "-")], ylabel="us", plot_name=f"fused-topk-perf-{scoring_func}", args={}, @@ -70,6 +70,28 @@ def get_benchmark(scoring_func): ), quantiles=quantiles, ) + elif provider == "aiter_gating": + from vllm._aiter_ops import rocm_aiter_ops + + if scoring_func != "softmax" or not rocm_aiter_ops.is_topk_gating_enabled(): + return float("nan"), float("nan"), float("nan") + topk_weights = torch.empty( + (num_tokens, topk), dtype=torch.float32, device="cuda" + ) + topk_ids = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda") + token_expert_indices = torch.empty( + (num_tokens, topk), dtype=torch.int32, device="cuda" + ) + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: rocm_aiter_ops.topk_gating( + topk_weights, + topk_ids, + token_expert_indices, + gating_output, + renormalize, + ), + quantiles=quantiles, + ) else: ms, min_ms, max_ms = triton.testing.do_bench( lambda: fused_topk( diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index b796623b7a..26eb5ffe6b 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -32,8 +32,9 @@ else() # Pinned to the fork's dev branch: upstream's 26/09/30 release (2.8.1) # plus the SM120 port, the SM90 paged-MQA kv_block=32/next_n=4 port, # configurable SwiGLU alpha/beta, SiTU for FP8/FP4 Mega MoE, and NVFP4 - # Mega MoE. TORCH_LIBRARY build with a CPython 3.10 stable-ABI floor. - set(_DEEPGEMM_UPSTREAM_TAG "6ce1c5ec6627a6bf11dbcc698b0b1e19c9248432") + # Mega MoE, plus swiglustep for BF16/FP8/NVFP4 Mega MoE. + # Uses the PyTorch 2.11 stable ABI and CPython 3.10 limited API. + set(_DEEPGEMM_UPSTREAM_TAG "a7a927731b109fdfec770df7c2a5be5d7e2216b5") set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}") if(NOT _deepgemm_fc_root) diff --git a/cmake/external_projects/flashmla.cmake b/cmake/external_projects/flashmla.cmake index cfa8243f72..daec85069c 100644 --- a/cmake/external_projects/flashmla.cmake +++ b/cmake/external_projects/flashmla.cmake @@ -19,7 +19,7 @@ else() FetchContent_Declare( flashmla GIT_REPOSITORY https://github.com/vllm-project/FlashMLA - GIT_TAG 0eee43b12f034b657133cf2afca6a72ebb6efccf + GIT_TAG 5dceadbd9eabc9acfc547ef18ec8f7b14b980682 GIT_PROGRESS TRUE CONFIGURE_COMMAND "" BUILD_COMMAND "" @@ -81,72 +81,83 @@ if(FLASH_MLA_ARCHS) ${flashmla_SOURCE_DIR}/csrc/api/dense_decode.cpp ${flashmla_SOURCE_DIR}/csrc/api/fused_norm_rope_attn_rope_cast_fwd.cpp - # Misc kernels for decoding - ${flashmla_SOURCE_DIR}/csrc/kernels/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/smxx/decode/combine/combine.cu - - # sm90 dense decode - ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/dense/instantiations/fp16.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/dense/instantiations/bf16.cu - - # sm90 sparse decode - ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h64.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h128.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h64.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h128.cu - - # sm90 sparse prefill - ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu + # Current kernels (DeepSeek V4.1, token-interleaved KV cache scales) + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/smxx/decode/combine/combine.cu # sm100 dense prefill (inference-only: dense backward stays out) - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu # sm100 sparse prefill - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k512.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k576.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_k512.cu - - # sm100 fused norm + rope + attn + rope + cast - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_norm.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_nonorm.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_norm.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_nonorm.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_norm.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_nonorm.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_norm.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_nonorm.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_norm.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_nonorm.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_norm.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_nonorm.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_norm.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_nonorm.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_norm.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_nonorm.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_q_b_proj/kernel.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_wv_proj/kernel.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k512.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_k512.cu # sm100 sparse decode - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64_no_split.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64_no_split.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64_no_split.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64_no_split.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/nvfp4_head64/instantiations/v32_nvfp4_fp8rope.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_splitkv.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41_splitkv.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4.cu - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4_splitkv.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/decode/sparse/head64/instantiations/v41_h64.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/decode/sparse/head64/instantiations/v41_h64_no_split.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64_no_split.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41_splitkv.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4_splitkv.cu + + # sm100 fused norm + rope + attn + rope + cast + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_prefill_norm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_prefill_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_prefill_norm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_prefill_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_q_b_proj/kernel.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_wv_proj/kernel.cu + + # Legacy kernels (upstream keeps them under cuda_kernels/legacy/): + # V4, V3.2 and V3.2-NVFP4, and everything on Hopper + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/smxx/decode/combine/combine.cu + + # sm90 dense decode + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/decode/dense/instantiations/fp16.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/decode/dense/instantiations/bf16.cu + + # sm90 sparse decode + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/decode/sparse/instantiations/v4_persistent_h64.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/decode/sparse/instantiations/v4_persistent_h128.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/decode/sparse/instantiations/v32_persistent_h64.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/decode/sparse/instantiations/v32_persistent_h128.cu + + # sm90 sparse prefill + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/prefill/sparse/instantiations/phase1_k512.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/prefill/sparse/instantiations/phase1_k576.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu + + # sm100 sparse prefill, d_qk = 576 + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k576.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu + + # sm100 sparse decode + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/decode/sparse/head64/instantiations/v32_h64.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/decode/sparse/head64/instantiations/v32_h64_no_split.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/decode/sparse/head64/instantiations/v4_h64.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/decode/sparse/head64/instantiations/v4_h64_no_split.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/decode/sparse/nvfp4_head64/instantiations/v32_nvfp4_fp8rope.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_splitkv.cu + + # sm100 fused norm + rope + attn + rope + cast, V4 decode + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_nonorm.cu ) set(FlashMLA_Extension_SOURCES @@ -158,17 +169,17 @@ if(FLASH_MLA_ARCHS) set(FlashMLA_INCLUDES ${flashmla_SOURCE_DIR}/csrc - ${flashmla_SOURCE_DIR}/csrc/kerutils/include - ${flashmla_SOURCE_DIR}/csrc/cutlass/include - ${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include + ${flashmla_SOURCE_DIR}/csrc/3rdparty/kerutils/include + ${flashmla_SOURCE_DIR}/csrc/3rdparty/cutlass/include + ${flashmla_SOURCE_DIR}/csrc/3rdparty/cutlass/tools/util/include ) set(FlashMLA_Extension_INCLUDES ${flashmla_SOURCE_DIR}/csrc - ${flashmla_SOURCE_DIR}/csrc/kerutils/include + ${flashmla_SOURCE_DIR}/csrc/3rdparty/kerutils/include ${flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/ - ${flashmla_SOURCE_DIR}/csrc/cutlass/include - ${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include + ${flashmla_SOURCE_DIR}/csrc/3rdparty/cutlass/include + ${flashmla_SOURCE_DIR}/csrc/3rdparty/cutlass/tools/util/include ) set_gencode_flags_for_srcs( @@ -202,6 +213,12 @@ if(FLASH_MLA_ARCHS) target_compile_definitions(_flashmla_C PRIVATE USE_CUDA) endif() + # Select the CUDA branch of FlashMLA's params.h and kerutils; kerutils + # detects CUDA by itself in .cu files, so only host files need the flag. + target_compile_definitions(_flashmla_C PRIVATE + FLASH_MLA_IS_BUILD_ON_CUDA + $<$:KERUTILS_IS_BUILD_ON_CUDA>) + define_extension_target( _flashmla_extension_C DESTINATION vllm diff --git a/csrc/cache_kernels.cu b/csrc/cache_kernels.cu index 24453ff1c1..351bcf24f9 100644 --- a/csrc/cache_kernels.cu +++ b/csrc/cache_kernels.cu @@ -273,7 +273,7 @@ __global__ void reshape_and_cache_kernel( const int64_t* __restrict__ slot_mapping, // [num_tokens] const int key_stride, const int value_stride, const int num_heads, const int head_size, const int block_size, const int x, - const float* k_scale, const float* v_scale) { + const float* k_scale, const float* v_scale, const bool interleaved_v) { const int64_t token_idx = blockIdx.x; const int64_t slot_idx = slot_mapping[token_idx]; if (slot_idx < 0) { @@ -301,10 +301,20 @@ __global__ void reshape_and_cache_kernel( key_cache + block_idx * num_heads * h_block_count * block_size * x + head_idx * h_block_count * block_size * x + h_block * block_size * x + block_offset * x; - const int64_t tgt_value_start = - block_idx * num_heads * h_block_count * x * block_size + - head_idx * h_block_count * x * block_size + h_block * x * block_size + - block_offset; + // V cache target: standard [.., head_size, block_size], or the interleaved + // [.., block_size/x, head_size, x] the fused AITER writer/reader use. + const int64_t v_head_base = block_idx * num_heads * head_size * block_size + + head_idx * head_size * block_size; + int64_t tgt_value_start; + int v_dim_stride; + if (interleaved_v) { + tgt_value_start = v_head_base + (block_offset / x) * (head_size * x) + + (h_block * x) * x + (block_offset % x); + v_dim_stride = x; + } else { + tgt_value_start = v_head_base + h_block * x * block_size + block_offset; + v_dim_stride = block_size; + } constexpr int VEC_SIZE = (sizeof(scalar_t) == 2) ? 8 : 4; float k_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) ? 0.f : *k_scale; @@ -318,7 +328,7 @@ __global__ void reshape_and_cache_kernel( cache_t* __restrict__ value_dst = value_cache + tgt_value_start; #pragma unroll for (int i = 0; i < x; i++) { - v_op(value_dst[i * block_size], value_src[i]); + v_op(value_dst[i * v_dim_stride], value_src[i]); } } @@ -775,7 +785,7 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel( slot_mapping.const_data_ptr(), key_stride, value_stride, \ num_heads, head_size, block_size, x, \ reinterpret_cast(k_scale.data_ptr()), \ - reinterpret_cast(v_scale.data_ptr())); + reinterpret_cast(v_scale.data_ptr()), interleaved_v); void reshape_and_cache( torch::stable::Tensor& key, // [num_tokens, num_heads, head_size] @@ -783,10 +793,12 @@ void reshape_and_cache( torch::stable::Tensor& key_cache, // [num_blocks, num_heads, head_size/x, block_size, x] torch::stable::Tensor& - value_cache, // [num_blocks, num_heads, head_size, block_size] + value_cache, // [num_blocks, num_heads, head_size, block_size], or the + // interleaved [.., block_size/x, head_size, x] when + // interleaved_v is set torch::stable::Tensor& slot_mapping, // [num_tokens] const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, - torch::stable::Tensor& v_scale) { + torch::stable::Tensor& v_scale, bool interleaved_v) { int num_tokens = slot_mapping.size(0); int num_heads = key.size(1); int head_size = key.size(2); diff --git a/csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu b/csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu index 4eb8d7756b..c90c48d1ba 100644 --- a/csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu +++ b/csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu @@ -32,9 +32,10 @@ * [0, bs*576): token data, 448 fp8 + 128 bf16 each * [bs*576, bs*576 + bs*8): UE8M0 scales, 7 real + 1 pad per token * V4.1 (528 B/token, kv_mxfp8=true) -- RoPE is quantized as well, so the - * whole row is MXFP8 (fp8 e4m3 with one UE8M0 scale per 32 dims) - * [0, bs*512): token data, 512 fp8 each - * [bs*512, bs*512 + bs*16): UE8M0 scales, 16 per token + * whole row is MXFP8 (fp8 e4m3 with one UE8M0 scale per 32 dims), and each + * token's scales follow its own data + * [i*528, i*528 + 512): token i data, 512 fp8 + * [i*528 + 512, (i+1)*528): token i UE8M0 scales, 16 */ #include "torch_utils.h" @@ -104,6 +105,8 @@ constexpr int kTokenDataBytes = kNopeDim + kRopeDim * 2; // 448 + 128 = 576 constexpr int kMxfp8QuantBlock = 32; constexpr int kMxfp8NumQuantBlocks = kHeadDim / kMxfp8QuantBlock; // 16 constexpr int kMxfp8TokenDataBytes = kHeadDim; // 512 +constexpr int kMxfp8BytesPerToken = + kMxfp8TokenDataBytes + kMxfp8NumQuantBlocks; // 528 // FNUZ on gfx942 / OCP elsewhere. FNUZ uses 224.0 (not the dtype's raw // 240.0) to match the rest of vLLM's FNUZ pipeline. #if defined(USE_ROCM) && defined(__gfx942__) @@ -356,20 +359,20 @@ __device__ __forceinline__ void processDeepseekV4Slot( // V4.1 (kv_mxfp8) quantizes them too. int64_t const slot_id = slot_mapping[tokenIdx]; if (slot_id >= 0) { - int const token_data_bytes = - kv_mxfp8 ? kMxfp8TokenDataBytes : kTokenDataBytes; - int const scale_bytes = - kv_mxfp8 ? kMxfp8NumQuantBlocks : kScaleBytesPerToken; int64_t const block_idx = slot_id / cache_block_size; int64_t const pos_in_block = slot_id % cache_block_size; uint8_t* block_base = k_cache + block_idx * static_cast(kv_block_stride); - uint8_t* token_fp8_ptr = block_base + pos_in_block * token_data_bytes; + uint8_t* token_fp8_ptr = + block_base + + pos_in_block * (kv_mxfp8 ? kMxfp8BytesPerToken : kTokenDataBytes); uint8_t* token_bf16_ptr = token_fp8_ptr + kNopeDim; uint8_t* token_scale_ptr = - block_base + - static_cast(cache_block_size) * token_data_bytes + - pos_in_block * scale_bytes; + kv_mxfp8 + ? token_fp8_ptr + kMxfp8TokenDataBytes + : block_base + + static_cast(cache_block_size) * kTokenDataBytes + + pos_in_block * kScaleBytesPerToken; #pragma unroll for (int i = 0; i < kElemsPerLane; i++) { @@ -790,7 +793,7 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert( // backend. Differences from the legacy path: // * No Q head padding — output Q layout matches the input num_heads_q. // * KV is written as a *contiguous* 512-wide row per token (token-strided), -// not the legacy UE8M0 paged layout with a separate scale tail. +// not the V4 paged record with a separate scale tail. // * Q/KV are stored either as bf16 or as per-tensor E4M3 FP8 (one global // scale), selected by the STORE_Q_FP8 / STORE_KV_FP8 template flags. // diff --git a/csrc/ops.h b/csrc/ops.h index 3266d28b1e..549dc84a1c 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -650,7 +650,7 @@ void reshape_and_cache(torch::stable::Tensor& key, torch::stable::Tensor& value, torch::stable::Tensor& slot_mapping, const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, - torch::stable::Tensor& v_scale); + torch::stable::Tensor& v_scale, bool interleaved_v); void reshape_and_cache_flash( torch::stable::Tensor& key, torch::stable::Tensor& value, diff --git a/csrc/quantization/activation_kernels.cu b/csrc/quantization/activation_kernels.cu index 9d7da85a22..0e01c0d912 100644 --- a/csrc/quantization/activation_kernels.cu +++ b/csrc/quantization/activation_kernels.cu @@ -4,6 +4,7 @@ #include "core/utils.hpp" #include "core/cuda_compat.h" +#include "cuda_vec_utils.cuh" #include "dispatch_utils.h" #include "quantization/w8a8/fp8/common.cuh" @@ -43,7 +44,7 @@ template 16 ? num_tokens > 32 ? 1 : 2 : 4); \ - dim3 block(std::min(d, 512)); \ - const torch::stable::accelerator::DeviceGuard device_guard( \ - input.get_device_index()); \ - const cudaStream_t stream = \ - get_current_cuda_stream(input.get_device_index()); \ - VLLM_STABLE_DISPATCH_FLOATING_TYPES( \ - input.scalar_type(), "act_and_mul_kernel", [&] { \ - VLLM_STABLE_DISPATCH_FP8_TYPES( \ - out.scalar_type(), "act_and_mul_quant_kernel_fp8_type", [&] { \ - vllm::act_and_mul_quant_kernel, \ - fp8_t> \ - <<>>( \ - out.mutable_data_ptr(), \ - input.const_data_ptr(), \ - scale.const_data_ptr(), d); \ - }); \ +#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL) \ + int d = input.size(-1) / 2; \ + int64_t num_tokens = input.numel() / input.size(-1); \ + dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4); \ + dim3 block(std::min(d, 512)); \ + const torch::stable::accelerator::DeviceGuard device_guard( \ + input.get_device_index()); \ + const cudaStream_t stream = \ + get_current_cuda_stream(input.get_device_index()); \ + VLLM_STABLE_DISPATCH_FLOATING_TYPES( \ + input.scalar_type(), "act_and_mul_kernel", [&] { \ + VLLM_STABLE_DISPATCH_FP8_TYPES( \ + out.scalar_type(), "act_and_mul_quant_kernel_fp8_type", [&] { \ + const scalar_t* in_ptr = input.const_data_ptr(); \ + fp8_t* out_ptr = out.mutable_data_ptr(); \ + const bool vectorize = d % (16 / sizeof(scalar_t)) == 0 && \ + vllm::is_16byte_aligned(in_ptr) && \ + vllm::is_16byte_aligned(out_ptr); \ + vllm::act_and_mul_quant_kernel, \ + fp8_t> \ + <<>>(out_ptr, in_ptr, \ + scale.const_data_ptr(), \ + d, vectorize); \ + }); \ }); void silu_and_mul_quant(torch::stable::Tensor& out, // [..., d] diff --git a/csrc/rocm/attention.cu b/csrc/rocm/attention.cu index caf6ee9a14..7106840964 100644 --- a/csrc/rocm/attention.cu +++ b/csrc/rocm/attention.cu @@ -355,7 +355,8 @@ __device__ float warpReduceMax(float val) { // clang-format off template + int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO, MFMAType MFMA_TYPE, + bool USE_INTERLEAVED_V_CACHE> __global__ __launch_bounds__(NUM_THREADS, 4) void paged_attention_ll4mi_QKV_mfma16_kernel( const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] @@ -593,26 +594,59 @@ __launch_bounds__(NUM_THREADS, 4) void paged_attention_ll4mi_QKV_mfma16_kernel( _B16x8 Vlocal[VTLOOP][VHELOOP][VTLANELOOP]; // this could be B8x16 too - const cache_t* v_ptr = v_cache + wg_start_kv_head_idx * kv_head_stride + - ((rowid * VTOKENS_PER_LANE) % BLOCK_SIZE); + // KX_V = number of contiguous token elements per interleaved chunk + // (e.g. 8 for bf16, 16 for fp8). This is also the 'x' dimension size. + constexpr int KX_V = 16 / sizeof(cache_t); - // v fetches are 16head elems across lanes x 16 tokens per lane - for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { - const int vhead_elem = vhe_depth * NWARPS * 16 + warpid * 16 + lane16id; - const cache_t* v_ptr2 = v_ptr + vhead_elem * BLOCK_SIZE; - - for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { - for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) { + if constexpr (USE_INTERLEAVED_V_CACHE) { + // Interleaved V-cache layout: [num_blocks, num_heads, block_size/x, + // head_dim, x] Address for v[block][head][token t][dim d]: + // chunk = t / x, sub = t % x + // addr = base + chunk * head_dim * x + d * x + sub + // Each 16B fetch reads KX_V contiguous elements within a single chunk, so + // a lane's token offset must land on a chunk boundary. This holds only + // when BLOCK_SIZE is a multiple of KX_V. + static_assert(BLOCK_SIZE % KX_V == 0, + "Interleaved V-cache read requires BLOCK_SIZE % KX_V == 0"); + const cache_t* v_ptr_head = v_cache + wg_start_kv_head_idx * kv_head_stride; + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { + const int vhead_elem = vhe_depth * NWARPS * 16 + warpid * 16 + lane16id; + for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { const int vblock_depth = 0; const int64_t vblock_number = static_cast( vphysical_block_number[vtoken_depth][vblock_depth]); - const cache_t* v_ptr3 = v_ptr2 + (vblock_number * kv_block_stride); - - const cache_t* v_fetch_ptr = - v_ptr3 + vfetch_depth * CONTIGUOUS_KV_ELEMS_16B_LOAD; - const _B16x8* v_fetch_ptr_16B = - reinterpret_cast(v_fetch_ptr); - Vlocal[vtoken_depth][vhe_depth][vfetch_depth] = *v_fetch_ptr_16B; + const cache_t* v_block_ptr = + v_ptr_head + vblock_number * kv_block_stride; + for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) { + const int bo = ((rowid * VTOKENS_PER_LANE) % BLOCK_SIZE) + + vfetch_depth * CONTIGUOUS_KV_ELEMS_16B_LOAD; + const int bo_chunk = bo / KX_V; // which chunk of x tokens + const cache_t* v_fetch_ptr = + v_block_ptr + (bo_chunk * HEAD_SIZE + vhead_elem) * KX_V + + (bo % KX_V); + Vlocal[vtoken_depth][vhe_depth][vfetch_depth] = + *reinterpret_cast(v_fetch_ptr); + } + } + } + } else { + const cache_t* v_ptr = v_cache + wg_start_kv_head_idx * kv_head_stride + + ((rowid * VTOKENS_PER_LANE) % BLOCK_SIZE); + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { + const int vhead_elem = vhe_depth * NWARPS * 16 + warpid * 16 + lane16id; + const cache_t* v_ptr2 = v_ptr + vhead_elem * BLOCK_SIZE; + for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { + for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) { + const int vblock_depth = 0; + const int64_t vblock_number = static_cast( + vphysical_block_number[vtoken_depth][vblock_depth]); + const cache_t* v_ptr3 = v_ptr2 + (vblock_number * kv_block_stride); + const cache_t* v_fetch_ptr = + v_ptr3 + vfetch_depth * CONTIGUOUS_KV_ELEMS_16B_LOAD; + const _B16x8* v_fetch_ptr_16B = + reinterpret_cast(v_fetch_ptr); + Vlocal[vtoken_depth][vhe_depth][vfetch_depth] = *v_fetch_ptr_16B; + } } } } @@ -996,7 +1030,8 @@ __launch_bounds__(NUM_THREADS, 4) void paged_attention_ll4mi_QKV_mfma16_kernel( template + int GQA_RATIO, + bool USE_INTERLEAVED_V_CACHE> __global__ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel( const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] @@ -1177,25 +1212,41 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel( } } + // KX_V = number of contiguous token elements per interleaved chunk + constexpr int KX_V = 16 / sizeof(cache_t); const cache_t* v_ptr = v_cache + wg_start_kv_head_idx * kv_head_stride; - // fetch vcache in kv cache auto case + // fetch vcache in kv cache auto case (bf16/fp16) if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { const _B16x8* v_ptrh8 = reinterpret_cast(v_ptr); - // iterate over each v block for (int b = 0; b < VBLOCKS; b++) { - // int32 physical_block_number leads to overflow when multiplied with - // kv_block_stride const int64_t vphysical_block_number = static_cast(vphysical_blocks[b]); const _B16x8* v_ptrh8b = v_ptrh8 + (vphysical_block_number * kv_block_stride) / 8; - // iterate over each head elem (within head_size) for (int h = 0; h < VHELOOP; h++) { const int head_size_elem = h * WARP_SIZE + laneid; - const _B16x8* v_ptrh8be = v_ptrh8b + head_size_elem * BLOCK_SIZE / 8; - // iterate over all velems within block - for (int d = 0; d < BLOCK_SIZE / 8; d++) { - Vlocal[h][b * BLOCK_SIZE / 8 + d] = v_ptrh8be[d]; + if constexpr (USE_INTERLEAVED_V_CACHE) { + // Interleaved layout: [num_blocks, num_heads, block_size/x, + // head_dim, x] Iterate over chunks of x tokens; within each chunk, + // x values for the same head_size_elem are contiguous (16-byte + // aligned). + for (int bo_chunk = 0; bo_chunk < BLOCK_SIZE / KX_V; bo_chunk++) { + const _B16x8* v_chunk = + v_ptrh8b + (bo_chunk * HEAD_SIZE + head_size_elem) * KX_V / 8; + for (int xd = 0; xd < KX_V / 8; xd++) { + Vlocal[h][b * BLOCK_SIZE / 8 + bo_chunk * (KX_V / 8) + xd] = + v_chunk[xd]; + } + } + } else { + // Standard PagedAttention layout: [num_blocks, num_heads, head_dim, + // block_size] All block_size tokens for one head_size_elem are + // contiguous. + const _B16x8* v_ptrh8be = + v_ptrh8b + head_size_elem * BLOCK_SIZE / 8; + for (int d = 0; d < BLOCK_SIZE / 8; d++) { + Vlocal[h][b * BLOCK_SIZE / 8 + d] = v_ptrh8be[d]; + } } } } @@ -1203,21 +1254,29 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel( // fetch vcache in fp8 case else { // if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) const _B8x8* v_ptrh8 = reinterpret_cast(v_ptr); - // iterate over each v block for (int b = 0; b < VBLOCKS; b++) { - // int32 physical_block_number leads to overflow when multiplied with - // kv_block_stride const int64_t vphysical_block_number = static_cast(vphysical_blocks[b]); const _B8x8* v_ptrh8b = v_ptrh8 + (vphysical_block_number * kv_block_stride) / 8; - // iterate over each head elem (within head_size) for (int h = 0; h < VHELOOP; h++) { const int head_size_elem = h * WARP_SIZE + laneid; - const _B8x8* v_ptrh8be = v_ptrh8b + head_size_elem * BLOCK_SIZE / 8; - // iterate over all velems within block - for (int d = 0; d < BLOCK_SIZE / 8; d++) { - Vlocalb8[h][b * BLOCK_SIZE / 8 + d] = v_ptrh8be[d]; + if constexpr (USE_INTERLEAVED_V_CACHE) { + // Same interleaved addressing as bf16 above, but with fp8 + // elements (KX_V=16, using _B8x8 loads instead of _B16x8). + for (int bo_chunk = 0; bo_chunk < BLOCK_SIZE / KX_V; bo_chunk++) { + const _B8x8* v_chunk = + v_ptrh8b + (bo_chunk * HEAD_SIZE + head_size_elem) * KX_V / 8; + for (int xd = 0; xd < KX_V / 8; xd++) { + Vlocalb8[h][b * BLOCK_SIZE / 8 + bo_chunk * (KX_V / 8) + xd] = + v_chunk[xd]; + } + } + } else { + const _B8x8* v_ptrh8be = v_ptrh8b + head_size_elem * BLOCK_SIZE / 8; + for (int d = 0; d < BLOCK_SIZE / 8; d++) { + Vlocalb8[h][b * BLOCK_SIZE / 8 + d] = v_ptrh8be[d]; + } } } } @@ -1831,7 +1890,8 @@ __device__ __forceinline__ _B16x8 from_floatx8(const floatx8& inp) { template + MFMAType MFMA_TYPE, + bool USE_INTERLEAVED_V_CACHE> __global__ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel( const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] @@ -2285,8 +2345,8 @@ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel( template + int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO, + bool USE_INTERLEAVED_V_CACHE> __global__ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel( const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] @@ -2611,7 +2671,8 @@ __device__ __forceinline__ _B16x8 from_floatx8(const floatx8& inp) { template + MFMAType MFMA_TYPE, + bool USE_INTERLEAVED_V_CACHE> __global__ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel( const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] @@ -3030,8 +3091,8 @@ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel( template + int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO, + bool USE_INTERLEAVED_V_CACHE> __global__ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel( const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] @@ -3261,7 +3322,8 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel( template + int GQA_RATIO, MFMAType MFMA_TYPE, + bool USE_INTERLEAVED_V_CACHE> __global__ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma16_kernel( const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] @@ -3288,7 +3350,8 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma16_kernel( template + int GQA_RATIO, + bool USE_INTERLEAVED_V_CACHE> __global__ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel( const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] @@ -3330,10 +3393,11 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel( #endif +// clang-format off #define LAUNCH_CUSTOM_ATTENTION_MFMA16(GQA_RATIO) \ - paged_attention_ll4mi_QKV_mfma16_kernel \ + paged_attention_ll4mi_QKV_mfma16_kernel< \ + T, KVT, KV_DTYPE, OUTT, BLOCK_SIZE, HEAD_SIZE, NTHR, ALIBI_ENABLED, \ + GQA_RATIO, MFMA_TYPE, USE_INTERLEAVED_V_CACHE> \ <<>>( \ query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale, \ block_tables_ptr, seq_lens_ptr, query_start_loc_ptr, \ @@ -3342,15 +3406,16 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel( max_ctx_blocks, k_scale_ptr, v_scale_ptr); #define LAUNCH_CUSTOM_ATTENTION_MFMA4(GQA_RATIO) \ - paged_attention_ll4mi_QKV_mfma4_kernel \ + paged_attention_ll4mi_QKV_mfma4_kernel< \ + T, KVT, KV_DTYPE, OUTT, BLOCK_SIZE, HEAD_SIZE, NTHR, ALIBI_ENABLED, \ + GQA_RATIO, USE_INTERLEAVED_V_CACHE> \ <<>>( \ query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale, \ block_tables_ptr, seq_lens_ptr, query_start_loc_ptr, \ max_num_blocks_per_seq, alibi_slopes_ptr, q_stride, kv_block_stride, \ kv_head_stride, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, out_ptr, \ max_ctx_blocks, k_scale_ptr, v_scale_ptr); +// clang-format on #define LAUNCH_CUSTOM_REDUCTION(NPAR_LOOPS) \ paged_attention_ll4mi_reduce_kernel + bool ALIBI_ENABLED, MFMAType MFMA_TYPE, bool USE_INTERLEAVED_V_CACHE> void paged_attention_custom_launcher( torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits, torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache, @@ -3517,7 +3582,7 @@ void paged_attention_custom_launcher( template + bool ALIBI_ENABLED, MFMAType MFMA_TYPE, bool USE_INTERLEAVED_V_CACHE> void paged_attention_custom_launcher_navi( torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits, torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache, @@ -3688,21 +3753,42 @@ void paged_attention_custom_launcher_navi( } } -#define CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, \ - PSIZE, ALIBI_ENABLED, MFMA_TYPE) \ - if (!is_navi) { \ - paged_attention_custom_launcher( \ - out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \ - num_kv_heads, scale, block_tables, seq_lens, query_start_loc, \ - max_seq_len, alibi_slopes, k_scale, v_scale, fp8_out_scale); \ - } else { \ - paged_attention_custom_launcher_navi( \ - out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \ - num_kv_heads, scale, block_tables, seq_lens, query_start_loc, \ - max_seq_len, alibi_slopes, k_scale, v_scale); \ +// USE_INTERLEAVED_V_CACHE gates the V-cache read path at compile time. +// Currently only the GFX9 kernels (MI300x/MI325x/MI350x/MI355x) have the +// actual interleaved V-fetch logic implemented. The GFX11 (RDNA 3) and +// GFX12 (RDNA 4) mfma16 kernels do NOT support interleaved V-cache addressing. +// The template parameter is carried on all variants for compilation +// compatibility; a TORCH_CHECK prevents dispatch to NAVI with interleaved V. +// TODO: implement interleaved V-cache support in the NAVI (GFX11/GFX12) +// kernels if shuffle KV cache layout is needed on RDNA GPUs. +#define CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, \ + PSIZE, ALIBI_ENABLED, MFMA_TYPE) \ + if (use_interleaved_v_cache) { \ + TORCH_CHECK(!is_navi, \ + "Reading interleaved V-cache layout from paged decode " \ + "attention not supported on NAVI GPUs"); \ + paged_attention_custom_launcher( \ + out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \ + num_kv_heads, scale, block_tables, seq_lens, query_start_loc, \ + max_seq_len, alibi_slopes, k_scale, v_scale, fp8_out_scale); \ + } else { \ + if (!is_navi) { \ + paged_attention_custom_launcher( \ + out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \ + num_kv_heads, scale, block_tables, seq_lens, query_start_loc, \ + max_seq_len, alibi_slopes, k_scale, v_scale, fp8_out_scale); \ + } else { \ + paged_attention_custom_launcher_navi( \ + out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \ + num_kv_heads, scale, block_tables, seq_lens, query_start_loc, \ + max_seq_len, alibi_slopes, k_scale, v_scale); \ + } \ } #define CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \ @@ -3799,7 +3885,8 @@ void paged_attention( const std::string& kv_cache_dtype, torch::Tensor& k_scale, torch::Tensor& v_scale, const std::optional& fp8_out_scale, - const std::string& mfma_type) { + const std::string& mfma_type, + bool use_interleaved_v_cache) { // clang-format on bool is_navi = is_navi_gpu(); const int head_size = query.size(2); diff --git a/csrc/rocm/ops.h b/csrc/rocm/ops.h index 78751841ca..de963913ea 100644 --- a/csrc/rocm/ops.h +++ b/csrc/rocm/ops.h @@ -50,4 +50,4 @@ void paged_attention( int64_t max_seq_len, const std::optional& alibi_slopes, const std::string& kv_cache_dtype, torch::Tensor& k_scale, torch::Tensor& v_scale, const std::optional& fp8_out_scale, - const std::string& mfma_type); + const std::string& mfma_type, bool use_interleaved_v_cache); diff --git a/csrc/rocm/torch_bindings.cpp b/csrc/rocm/torch_bindings.cpp index 14d1822707..077c3f35ce 100644 --- a/csrc/rocm/torch_bindings.cpp +++ b/csrc/rocm/torch_bindings.cpp @@ -92,7 +92,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) { " str kv_cache_dtype," " Tensor k_scale, Tensor v_scale," " Tensor? fp8_out_scale," - " str mfma_type) -> ()"); + " str mfma_type," + " bool use_interleaved_v_cache=False) -> ()"); rocm_ops.impl("paged_attention", torch::kCUDA, &paged_attention); } diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 3d4f9ab273..a92d9cc114 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -955,7 +955,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) { " Tensor! key_cache, Tensor! value_cache," " Tensor slot_mapping," " str kv_cache_dtype," - " Tensor k_scale, Tensor v_scale) -> ()"); + " Tensor k_scale, Tensor v_scale," + " bool interleaved_v=False) -> ()"); // Reshape the key and value tensors and cache them. ops.def( diff --git a/docker/Dockerfile.rocm_base_gfx1250 b/docker/Dockerfile.rocm_base_gfx1250 index 3fad6800ab..ae41096a6e 100644 --- a/docker/Dockerfile.rocm_base_gfx1250 +++ b/docker/Dockerfile.rocm_base_gfx1250 @@ -1,16 +1,16 @@ ARG BASE_IMAGE=ubuntu:24.04 -ARG ROCM_WHEEL_INDEX=https://rocm.devreleases.amd.com/whl-multi-arch/ -ARG ROCM_SDK_VERSION=7.14.0a20260623 -ARG TORCH_VERSION=2.11.0+rocm7.14.0a20260623 -ARG TORCHVISION_VERSION=0.26.0+rocm7.14.0a20260623 -ARG TORCHAUDIO_VERSION=2.11.0+rocm7.14.0a20260623 -ARG TRITON_VERSION=3.7.1+git110cd8e2.rocm7.14.0a20260623 -ARG APEX_VERSION=1.11.0+rocm7.14.0a20260623 +ARG ROCM_WHEEL_INDEX=https://rocm.nightlies.amd.com/whl-multi-arch/ +ARG ROCM_SDK_VERSION=10.1.0a20260822 +ARG TORCH_VERSION=2.11.0+rocm10.1.0a20260822 +ARG TORCHVISION_VERSION=0.26.0+rocm10.1.0a20260822 +ARG TORCHAUDIO_VERSION=2.11.0+rocm10.1.0a20260822 +ARG TRITON_VERSION=3.8.0+git4cff872c.rocm10.1.0a20260822 +ARG APEX_VERSION=1.11.0+rocm10.1.0a20260822 ARG FA_BRANCH="jpvillam/gfx1250_wip" ARG FA_REPO="https://github.com/jpvillam-amd/flash-attention.git" -ARG AITER_BRANCH="main" +ARG AITER_BRANCH="f85b24b" ARG AITER_REPO="https://github.com/ROCm/aiter.git" ARG MORI_BRANCH="v1.1.0" ARG MORI_REPO="https://github.com/ROCm/mori.git" @@ -119,10 +119,10 @@ ENV PYTHONPATH=${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi # which keeps the set consistent and avoids pip backtracking. RUN pip install --pre --index-url "${ROCM_WHEEL_INDEX}" \ --extra-index-url https://pypi.org/simple \ - "torch[device-all]==${TORCH_VERSION}" \ + "torch[device-gfx1250]==${TORCH_VERSION}" \ "torchvision==${TORCHVISION_VERSION}" \ "torchaudio==${TORCHAUDIO_VERSION}" \ - "rocm[libraries,devel,device-all]==${ROCM_SDK_VERSION}" && \ + "rocm[libraries,devel,device-gfx1250]==${ROCM_SDK_VERSION}" && \ rocm-sdk init # Torch runtime deps that may not be published on the ROCm wheel index; @@ -200,15 +200,6 @@ ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}} ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0} -### -### AMD SMI Build -### -FROM base AS build_amdsmi -RUN cd ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi \ - && pip wheel . --wheel-dir=dist -RUN mkdir -p /app/install && cp ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi/dist/*.whl /app/install - - ### ### MORI Build TODO(Build needs fixing) ### @@ -253,7 +244,10 @@ FROM base AS build_aiter ARG AITER_BRANCH ARG AITER_REPO ARG USE_SCCACHE -RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO} + +RUN git clone ${AITER_REPO} \ + && cd aiter \ + && git checkout --detach ${AITER_BRANCH} RUN cd aiter \ && git submodule update --init --recursive \ && pip install -r requirements.txt @@ -278,8 +272,6 @@ FROM base AS debs_wheel_release RUN mkdir /app/debs RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi -RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ cp /install/*.whl /app/debs @@ -288,8 +280,6 @@ FROM base AS debs RUN mkdir /app/debs RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi -RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \ diff --git a/docker/Dockerfile.rocm_gfx1250 b/docker/Dockerfile.rocm_gfx1250 index 0609c73c04..634e522e0d 100644 --- a/docker/Dockerfile.rocm_gfx1250 +++ b/docker/Dockerfile.rocm_gfx1250 @@ -15,6 +15,7 @@ ARG NIC_BACKEND=all # Users can specify a custom version compatible with their host drivers. # The default version has been tested with ionic-dkms=26.03.3.001 # (libionic1=54.0-187-1), required for MoRI-EP/NIXL on Pollara/AINIC. +# TODO (jmathias) : Check if we can revert this Pinned AINIC version ARG AINIC_VERSION=1.117.5-a-77 ARG UBUNTU_CODENAME=jammy @@ -115,8 +116,8 @@ WORKDIR ${COMMON_WORKDIR} FROM base AS fetch_vllm_0 ONBUILD COPY ./ vllm/ FROM base AS fetch_vllm_1 -ARG VLLM_REPO="https://github.com/ROCm/vllm.git" -ARG VLLM_BRANCH="455_wip" +ARG VLLM_REPO="https://github.com/vllm-project/vllm.git" +ARG VLLM_BRANCH="main" ENV VLLM_REPO=${VLLM_REPO} ENV VLLM_BRANCH=${VLLM_BRANCH} ONBUILD RUN git clone ${VLLM_REPO} \ @@ -726,13 +727,14 @@ RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \ # Download bench scripts for gfx1250 RUN curl -k -O https://raw.githubusercontent.com/ROCm/vllm/refs/heads/gfx1250_bench/vllm_smoketest.sh +RUN curl -k -O https://raw.githubusercontent.com/ROCm/vllm/refs/heads/gfx1250_bench/vllm_benchserve.sh ### Install triton from upstream for AITER Deps TODO: (JPVILLAM) If possible to get this on whls it would be better RUN pip3 uninstall -y triton && \ git clone https://github.com/triton-lang/triton.git && \ cd triton && \ - git checkout c517f38c && \ + git checkout aedd1f6 && \ TRITON_APPEND_CMAKE_ARGS="-DCMAKE_FIND_USE_CMAKE_ENVIRONMENT_PATH=FALSE" pip3 install . CMD ["/bin/bash"] diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index ada261dd8f..5e1e2984f4 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -170,8 +170,8 @@ FROM vllm-base AS ucx-nixl-build # Keep UCX_VERSION aligned with the UCX that NIXL itself builds against for this # NIXL_VERSION (nixl contrib/Dockerfile.manylinux: ARG UCX_REF), so XPU runs the # same UCX commit as the released nixl-cu* wheels rather than an untested pairing. -ARG UCX_VERSION=v1.21.x -ARG NIXL_VERSION=v1.3.2 +ARG UCX_VERSION=v1.23.x +ARG NIXL_VERSION=v1.5.0 # Build-time only: compiler, autotools, and verbs dev headers RUN apt-get update -y && apt-get install -y --no-install-recommends \ @@ -209,7 +209,7 @@ ENV UV_COMPILE_BYTECODE=1 # Must match the ucx-nixl-build stage above: ARG does not cross stages, so a stale # value here installs a meta package that disagrees with the wheel built there. -ARG NIXL_VERSION=v1.3.2 +ARG NIXL_VERSION=v1.5.0 # Copy compiled UCX runtime libraries and the pre-built NIXL wheel. # No compiler or autotools are installed in this stage. diff --git a/docs/design/cuda_graphs_multimodal.md b/docs/design/cuda_graphs_multimodal.md index 0a8a831f01..b1a85ce2da 100644 --- a/docs/design/cuda_graphs_multimodal.md +++ b/docs/design/cuda_graphs_multimodal.md @@ -12,6 +12,7 @@ For two-tower vision encoders (e.g., DeepSeek-OCR's SAM + CLIP with dynamic tili | Architecture | Models | CG for Image | CG for Video | Multi-Path Graph | | ------------ | ------ | ------------ | ------------ | --------------- | | `DeepseekOCRForCausalLM` | `DeepSeek-OCR` | ✅︎ | ❌︎ | ✅︎ | +| `DeepseekVLV2ForCausalLM` | `DeepSeek-VL2` | ✅︎ | ❌︎ | ❌︎ | | `DeepseekV41ForCausalLM` | `DeepSeek-V4.1-Flash` | ✅︎ | ❌︎ | ❌︎ | | `Ernie4_5_VLMoeForConditionalGeneration` | `ERNIE-4.5-VL` | ✅︎ | ❌︎ | ❌︎ | | `Gemma3ForConditionalGeneration` | `Gemma3` | ✅︎ | ❌︎ | ❌︎ | diff --git a/docs/features/speculative_decoding/mtp.md b/docs/features/speculative_decoding/mtp.md index 3b637c9de8..34d8f9539b 100644 --- a/docs/features/speculative_decoding/mtp.md +++ b/docs/features/speculative_decoding/mtp.md @@ -74,3 +74,12 @@ vllm serve XiaomiMiMo/MiMo-7B-Base \ is a good default to start with. - If your model does not support MTP, use another method such as EAGLE or draft model speculation. +- When the MTP head shares the target's unquantized lm_head (most MTP models), + `"draft_lm_head_quantization": "fp8"` or `"nvfp4"` in `--speculative-config` + gives the drafter a quantized copy of that head (Model Runner V2). The target + still verifies with its own head, so outputs are unchanged; only the + acceptance length can move. It helps most on bandwidth-bound GPUs with a + large vocabulary. The copy costs about one byte (FP8) or 0.56 bytes (NVFP4) + per lm_head weight, e.g. 0.6 or 0.35 GiB for a 248k x 2560 head; it is + counted in the model weights, so with `--kv-cache-memory-bytes` leave room + for it. diff --git a/docs/features/watermarking.md b/docs/features/watermarking.md index f51ec72eb1..adac09311c 100644 --- a/docs/features/watermarking.md +++ b/docs/features/watermarking.md @@ -125,8 +125,10 @@ generation. ## Speculative decoding Watermarking requires speculative decoding to use probabilistic draft sampling, -standard rejection sampling, and an autoregressive model-based method (`dspark`, -`eagle`, `eagle3`, or `mtp`). Parallel drafting is supported only by `dspark`. +standard rejection sampling, and a model-based method (`dflash`, `dspark`, +`eagle`, `eagle3`, or `mtp`). Parallel drafting is supported only by `dflash` +and `dspark`. DFlash drafters with a candidate head (`DFlash2DraftModel`, +`LiLiCorrDraftModel`) are not supported. A watermarking algorithm without native speculative-decoding support is rejected before model loading. Set diff --git a/docs/models/structured_decisions_models.md b/docs/models/structured_decisions_models.md index b1aea36812..3b53150a5f 100644 --- a/docs/models/structured_decisions_models.md +++ b/docs/models/structured_decisions_models.md @@ -304,12 +304,16 @@ vllm serve Cloudflare/clef-omni ### Decision 3.0 -The Decision 3.0 models ([`vllm-sr/d3`](https://huggingface.co/vllm-sr/d3) and -its smaller variants such as `vllm-sr/d3-mini` and `vllm-sr/d3-edge`) are a -Qwen3.5 backbone with an answer-code readout: one forward pass per question, -and the readout scores the question's options at the last token. vLLM detects -the checkpoint's `decision_config.json` and serves it as a classification -model. +The Decision 3.0 models are a Qwen3.5 backbone with an answer-code readout: +one forward pass per question, and the readout scores the question's options +at the last token. vLLM detects the checkpoint's `decision_config.json` and +serves it as a classification model. All six sizes take the same command: +[`vllm-sr/d3`](https://huggingface.co/vllm-sr/d3) (27B), +[`vllm-sr/d3-flash`](https://huggingface.co/vllm-sr/d3-flash) (9B), +[`vllm-sr/d3-mini`](https://huggingface.co/vllm-sr/d3-mini) (4B), +[`vllm-sr/d3-nano`](https://huggingface.co/vllm-sr/d3-nano) (2B), +[`vllm-sr/d3-lite`](https://huggingface.co/vllm-sr/d3-lite) (0.8B) and +[`vllm-sr/d3-edge`](https://huggingface.co/vllm-sr/d3-edge) (0.6B). ```bash vllm serve vllm-sr/d3 diff --git a/docs/serving/online_serving/derenderer.md b/docs/serving/online_serving/derenderer.md index f85433a963..b7b5e5b8f0 100644 --- a/docs/serving/online_serving/derenderer.md +++ b/docs/serving/online_serving/derenderer.md @@ -89,7 +89,7 @@ The derenderer builds its tool and reasoning parsers from its own server flags p A mismatch doesn't fail. The parser just splits `reasoning`, `content` and `tool_calls` differently from what `vllm serve` would return. -`/v1/completions/derender` only detokenizes. It never runs tool or reasoning parsers, the same as `/v1/completions` on `vllm serve`. It only reads `skip_special_tokens` from `completion_request`. +`/v1/completions/derender` only detokenizes. It never runs tool or reasoning parsers, the same as `/v1/completions` on `vllm serve`. From `completion_request` it reads `skip_special_tokens`, and `logprobs` and `logprob_token_ids` to cut `top_logprobs` the way `/v1/completions` does. ## Streaming @@ -105,7 +105,7 @@ The server keeps no state between calls. Everything the next call needs is in `s - **Send `prompt_token_ids` on the first chunk, even without a parser.** It seeds detokenization so the first token's leading space matches the coupled endpoint (see [Request format](#request-format)). Later chunks already carry that context in `stream_state`. - **Don't forward `[DONE]`.** It marks the end of the generate stream and isn't a chunk. -`/inference/v1/generate` returns `GenerateLogProbs` with integer `token_id`s (the generate server has no tokenizer), and derender turns those into `ChatCompletionLogProbs` / `CompletionLogProbs`, filling `token` and `bytes` from the tokenizer with the usual U+FFFD byte-fallback correction. Streaming derender does the same per chunk on the plain detokenization path (see [Streaming state and logprobs](#streaming-state-and-logprobs)). The parser path doesn't resolve them yet and drops them. +`/inference/v1/generate` returns `GenerateLogProbs` with integer `token_id`s (the generate server has no tokenizer), and derender turns those into `ChatCompletionLogProbs` / `CompletionLogProbs`, filling `token` and `bytes` from the tokenizer with the usual U+FFFD byte-fallback correction. Generate returns every top-k candidate the engine produced; derender keeps the ones the coupled endpoint would, with the same helpers: the first `top_logprobs` on chat (all of them for `top_logprobs=-1`) and the first `logprobs + 1` on completions (all of them for `logprobs=-1`). It reads those values from `chat_request` / `completion_request`, so without the original request, or with one that does not ask for logprobs, it returns every candidate. Requests that use `logprob_token_ids` get no output logprobs from generate (their `sampling_params.logprobs` is unset), so derender has none to return for them either. Streaming derender does the same per chunk on the plain detokenization path (see [Streaming state and logprobs](#streaming-state-and-logprobs)). The parser path doesn't resolve them yet and drops them. ## Streaming cost diff --git a/docs/serving/online_serving/renderer.md b/docs/serving/online_serving/renderer.md index e44a9de319..adf9c4e7c1 100644 --- a/docs/serving/online_serving/renderer.md +++ b/docs/serving/online_serving/renderer.md @@ -108,10 +108,13 @@ choice carries a `GenerateLogProbs`; `output_mode: "text"` returns decoded - `content` has one entry per generated token, in generation order. - `top_logprobs` is a list, not a dict: JSON turns dict keys into strings and the ordering would be implicit. It follows the engine's order: the sampled - token first, then the remaining candidates in rank order. With non-greedy - sampling the sampled token can sit outside the top k (for example ranks - `[5, 1, 2]` at `logprobs=2`); it then takes one of the `logprobs` slots and - the rank-k candidate is left out, as on the OpenAI endpoints. + token first, then the remaining candidates in rank order. It holds every + candidate the engine returned and is not cut to `logprobs`: k entries when + the sampled token is in the top k, k + 1 when non-greedy sampling picked a + token outside it (for example ranks `[5, 1, 2]` at `logprobs=2`). The OpenAI + endpoints cut this differently (`/v1/chat/completions` keeps the first k, + `/v1/completions` the first k + 1), so [derender](derenderer.md) applies the + cut of the endpoint it renders for. - `rank` is the token's rank in the vocabulary distribution (1 = most likely) on every entry, the sampled one included; a top-k candidate's rank is its top-k position. The list is not sorted by it, so sort by `rank` if you need @@ -132,6 +135,16 @@ choice carries a `GenerateLogProbs`; `output_mode: "text"` returns decoded `return_tokens_as_token_ids` on `/v1/chat/completions` and `/v1/completions` is unchanged: it is a user-facing OpenAI option and still uses the `token_id:N` format. + `top_logprobs` is no longer cut to `max(logprobs, 1)` entries, in either + `output_mode`: it now holds every candidate the engine returned, which is + `logprobs + 1` entries when the sampled token is outside the top + `logprobs`. Clients that assume `len(top_logprobs) == logprobs` should + take the first `logprobs` entries after sorting by `rank`. + Derender applies each endpoint's cut from the original request, so send + `chat_request` / `completion_request` to derender whenever logprobs are + requested: without it derender returns every candidate, which on chat is + `top_logprobs + 1` entries where `/v1/chat/completions` returns + `top_logprobs`. ## Multimodal Render Features diff --git a/examples/generate/multimodal/vision_language_offline.py b/examples/generate/multimodal/vision_language_offline.py index 95e21090a9..750bce8928 100644 --- a/examples/generate/multimodal/vision_language_offline.py +++ b/examples/generate/multimodal/vision_language_offline.py @@ -2274,6 +2274,7 @@ MODELS_SUPPORT_VIT_CUDA_GRAPH = [ "stepvl", "glm4_1v", "deepseek_ocr", + "deepseek_vl_v2", "ernie45_vl", "minicpmv2_5_vl", "minicpmv2_6_vl", diff --git a/requirements/cuda.txt b/requirements/cuda.txt index ad2bc63df1..e3d659545c 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -10,7 +10,7 @@ torchaudio==2.11.0 torchvision==0.29.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version torchcodec >= 0.14 PyNvVideoCodec==2.0.4 -flashinfer-python==0.7.1 +flashinfer-python==0.7.2 apache-tvm-ffi==0.1.11 tilelang==0.1.12 nvidia-cudnn-frontend>=1.19.1 @@ -25,7 +25,7 @@ nvidia-cutlass-dsl[cu13]==4.8.0 quack-kernels==0.6.5 # Required by MSA # Tokenspeed_MLA for faster mla with spec decode -tokenspeed-mla==2.0.2; platform_system == "Linux" +tokenspeed-mla==2.0.3; platform_system == "Linux" # Humming kernels for quantization gemm humming-kernels[cu13]==0.1.17 diff --git a/requirements/test/cpu.txt b/requirements/test/cpu.txt index d7d4692733..bfa4aede0b 100644 --- a/requirements/test/cpu.txt +++ b/requirements/test/cpu.txt @@ -157,8 +157,6 @@ cryptography==50.0.1 # azure-storage-blob # msal # pyjwt -cupy-cuda12x==13.6.0 - # via ray cycler==0.12.1 # via matplotlib datamodel-code-generator==0.79.0 @@ -239,8 +237,6 @@ fastavro==1.12.2 # via cohere fastparquet==2024.11.0 # via genai-perf -fastrlock==0.8.2 - # via cupy-cuda12x fastsafetensors==0.3.3 # via -r requirements/test/cuda.in filelock==4.0.7 @@ -569,7 +565,6 @@ numpy==2.2.6 # albumentations # bm25s # contourpy - # cupy-cuda12x # datasets # decord # einx diff --git a/requirements/test/cuda.in b/requirements/test/cuda.in index 8c9b8570a6..7a69c92969 100644 --- a/requirements/test/cuda.in +++ b/requirements/test/cuda.in @@ -22,7 +22,7 @@ vector_quantize_pytorch # required for minicpmo_26 test vocos # required for minicpmo_26 test peft>=0.20.0 # required for phi-4-mm test pqdm -ray[cgraph,default]>=2.56.0,<2.57 # Ray Compiled Graph, required by pipeline parallelism tests; sharded-RDT weight transfer needs Ray Direct Transport (>=2.56) +ray[default]>=2.56.0,<2.57 # sharded-RDT weight transfer needs Ray Direct Transport (>=2.56) sentence-transformers>=5.2.0 # required for embedding tests soundfile # required for audio tests jiwer # required for audio tests diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index d464efaf67..08a45361bc 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -170,8 +170,6 @@ cuda-pathfinder==1.3.3 # via cuda-bindings cuda-toolkit==13.0.3.0 # via torch -cupy-cuda12x==13.6.0 - # via ray cycler==0.12.1 # via matplotlib datamodel-code-generator==0.79.0 @@ -255,8 +253,6 @@ fastavro==1.12.2 # via cohere fastparquet==2024.11.0 # via genai-perf -fastrlock==0.8.2 - # via cupy-cuda12x fastsafetensors==0.3.3 # via # -c requirements/cuda.txt @@ -604,7 +600,6 @@ numpy==2.2.6 # albumentations # bm25s # contourpy - # cupy-cuda12x # datasets # decord # einx diff --git a/requirements/test/nightly-torch.txt b/requirements/test/nightly-torch.txt index ec8ca4b9c9..6e13a5118f 100644 --- a/requirements/test/nightly-torch.txt +++ b/requirements/test/nightly-torch.txt @@ -16,7 +16,7 @@ librosa # required for audio tests vocos # required for minicpmo_26 test peft pqdm -ray[cgraph,default]>=2.56.0,<2.57 # Ray Compiled Graph, required by pipeline parallelism tests; sharded-RDT weight transfer needs Ray Direct Transport (>=2.56) +ray[default]>=2.56.0,<2.57 # sharded-RDT weight transfer needs Ray Direct Transport (>=2.56) sentence-transformers>=5.2.0 # required for embedding tests soundfile # required for audio tests jiwer # required for audio tests diff --git a/requirements/test/rocm.in b/requirements/test/rocm.in index 0ad0b25511..e26765ec9c 100644 --- a/requirements/test/rocm.in +++ b/requirements/test/rocm.in @@ -21,7 +21,7 @@ vector_quantize_pytorch # required for minicpmo_26 test vocos # required for minicpmo_26 test peft>=0.20.0 # required for phi-4-mm test pqdm -ray[cgraph,default]==2.55.1 # Includes worker startup getenv/setenv race fix +ray[default]==2.55.1 # Includes worker startup getenv/setenv race fix sentence-transformers>=5.2.0 # required for embedding tests soundfile # required for audio tests jiwer # required for audio tests diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index a0ca0fd833..b08a77f0c9 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -606,7 +606,6 @@ numpy==2.2.6 # bitsandbytes # bm25s # contourpy - # cupy-cuda12x # datasets # decord # einx @@ -1439,7 +1438,6 @@ zipp==3.23.0 # cuda-bindings # cuda-pathfinder # cuda-toolkit -# cupy-cuda12x # nvidia-cublas # nvidia-cuda-cupti # nvidia-cuda-nvrtc diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index 42716de8de..5bf434f522 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -5,7 +5,7 @@ mod convert; mod types; mod validate; -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::convert::Infallible; use std::result::Result; use std::sync::Arc; @@ -158,7 +158,7 @@ async fn generate_chunk_stream( ResponseOptions { include_usage, include_continuous_usage, - logprobs: output_logprobs, + include_logprobs, // Ignored: raw generate streaming has no prompt-logprobs wire shape. include_prompt_logprobs: _, return_token_ids, @@ -213,15 +213,13 @@ async fn generate_chunk_stream( continue; } - let logprobs = if let Some(requested) = output_logprobs - && !token_ids.is_empty() - { + let logprobs = if include_logprobs && !token_ids.is_empty() { let logprobs = output.logprobs.as_ref().ok_or_else(|| { server_error!( "raw generate stream requested logprobs but generation returned none" ) })?; - Some(raw_logprobs_to_generate(logprobs, requested)?) + Some(raw_logprobs_to_generate(logprobs)?) } else { None }; @@ -289,19 +287,19 @@ fn collect_generate( include_usage: _, // Ignored: continuous usage is a streaming-only option. include_continuous_usage: _, - logprobs: output_logprobs, + include_logprobs, include_prompt_logprobs, return_token_ids, }: ResponseOptions, mm_placeholders: Option, ) -> Result { - let logprobs = if let Some(requested) = output_logprobs { + let logprobs = if include_logprobs { let logprobs = collected.logprobs.as_ref().ok_or_else(|| { ApiError::server_error( "raw generate response requested logprobs but generation returned none".to_string(), ) })?; - Some(raw_logprobs_to_generate(logprobs, requested)?) + Some(raw_logprobs_to_generate(logprobs)?) } else { None }; @@ -381,14 +379,11 @@ fn extract_mm_placeholders(features: Option<&[MmFeatureSpec]>) -> Option Result { +fn raw_logprobs_to_generate(logprobs: &Logprobs) -> Result { let content = logprobs .positions .iter() - .map(|position| position_to_generate_logprobs_content(position, requested)) + .map(position_to_generate_logprobs_content) .collect::, _>>()?; Ok(GenerateLogProbs { @@ -409,12 +404,12 @@ fn raw_prompt_logprobs_to_maps( .collect() } -/// `top_logprobs` matches the Python generate server: the engine's entries (the -/// sampled token first, then ranks 1..k) with repeated token ids dropped, cut to -/// `max(requested, 1)` entries, or all of them for `-1`. +/// `top_logprobs` matches the Python generate server: every candidate the +/// engine returned (the sampled token first, then ranks 1..k) with the repeated +/// sampled entry dropped, and no cut. Derender applies each OpenAI endpoint's +/// cut. fn position_to_generate_logprobs_content( position: &PositionLogprobs, - requested: i32, ) -> Result { let chosen = position.entries.first().ok_or_else(|| { ApiError::server_error( @@ -422,9 +417,9 @@ fn position_to_generate_logprobs_content( ) })?; - // One pass: the first entry of each token id wins, as in Python's dict. The - // sampled token is the only one the engine can list twice. - let mut seen = HashSet::with_capacity(position.entries.len()); + // The engine row is the sampled token followed by the top k, which are + // distinct; only the sampled token can repeat, at its own rank (anywhere in + // 1..=k). Skip that repeat, as Python's dict does: linear, no allocation. Ok(GenerateLogProbsContent { token_id: chosen.token_id, logprob: clamp_logprob(chosen.logprob), @@ -432,9 +427,9 @@ fn position_to_generate_logprobs_content( top_logprobs: position .entries .iter() - .filter(|entry| seen.insert(entry.token_id)) - .take(usize::try_from(requested).map_or(usize::MAX, |n| n.max(1))) - .map(|entry| GenerateLogProb { + .enumerate() + .filter(|(i, entry)| *i == 0 || entry.token_id != chosen.token_id) + .map(|(_, entry)| GenerateLogProb { token_id: entry.token_id, logprob: clamp_logprob(entry.logprob), rank: wire_rank(entry.rank), @@ -563,8 +558,8 @@ mod tests { } } - fn top(position: &PositionLogprobs, requested: i32) -> Vec<(u32, Option)> { - position_to_generate_logprobs_content(position, requested) + fn top(position: &PositionLogprobs) -> Vec<(u32, Option)> { + position_to_generate_logprobs_content(position) .unwrap() .top_logprobs .iter() @@ -576,24 +571,31 @@ mod tests { fn generate_top_logprobs_drop_the_repeated_sampled_token() { // Sampled token 7 is also top-1: the engine row is [7, 7, 8]. let pos = position(&[(7, -0.1, 1), (7, -0.1, 1), (8, -2.0, 2)]); - assert_eq!(top(&pos, 2), vec![(7, Some(1)), (8, Some(2))]); + assert_eq!(top(&pos), vec![(7, Some(1)), (8, Some(2))]); } #[test] - fn generate_top_logprobs_keep_sampled_first_when_outside_top_k() { - // Sampled token 50 has vocab rank 5; it takes a slot and rank 2 drops - // out, as on the Python generate server and the OpenAI endpoints. + fn generate_top_logprobs_keep_every_candidate_when_sampled_is_outside_top_k() { + // Sampled token 50 has vocab rank 5 at logprobs=2: sampled first, then + // ranks 1 and 2, all kept (derender cuts per endpoint). let pos = position(&[(50, -3.0, 5), (10, -0.2, 1), (20, -1.0, 2)]); - assert_eq!(top(&pos, 2), vec![(50, Some(5)), (10, Some(1))]); + assert_eq!(top(&pos), vec![(50, Some(5)), (10, Some(1)), (20, Some(2))]); } #[test] - fn generate_top_logprobs_respect_zero_and_all() { - let pos = position(&[(7, -0.1, 1), (7, -0.1, 1), (8, -2.0, 2), (9, -3.0, 3)]); - assert_eq!(top(&pos, 0), vec![(7, Some(1))]); + fn generate_top_logprobs_drop_the_sampled_repeat_at_its_own_rank() { + // Sampled token 30 is rank 3: the engine row is [30, 10, 20, 30, 40], + // so the repeat is at entries[3], not entries[1]. + let pos = position(&[ + (30, -1.0, 3), + (10, -0.2, 1), + (20, -0.5, 2), + (30, -1.0, 3), + (40, -2.0, 4), + ]); assert_eq!( - top(&pos, -1), - vec![(7, Some(1)), (8, Some(2)), (9, Some(3))] + top(&pos), + vec![(30, Some(3)), (10, Some(1)), (20, Some(2)), (40, Some(4))] ); } @@ -606,7 +608,7 @@ mod tests { entries.extend((0..VOCAB).map(|t| (t, -1.0 - t as f32, t + 1))); let pos = position(&entries); let start = std::time::Instant::now(); - let got = top(&pos, -1); + let got = top(&pos); assert!(start.elapsed() < std::time::Duration::from_secs(5)); assert_eq!(got.len(), VOCAB as usize); assert_eq!(got[0], (5, Some(1))); @@ -616,10 +618,10 @@ mod tests { fn generate_rank_zero_from_a_nan_logprob_is_none() { // The engine reports rank 0 for a sampled token whose logprob is NaN. let pos = position(&[(7, f32::NAN, 0), (8, -0.2, 1)]); - let content = position_to_generate_logprobs_content(&pos, 1).unwrap(); + let content = position_to_generate_logprobs_content(&pos).unwrap(); assert_eq!(content.rank, None); assert_eq!(content.logprob, -9999.0); - assert_eq!(top(&pos, 1), vec![(7, None)]); + assert_eq!(top(&pos), vec![(7, None), (8, Some(1))]); } #[tokio::test] diff --git a/rust/src/server/src/routes/inference/generate/convert.rs b/rust/src/server/src/routes/inference/generate/convert.rs index 313e9fa1f9..cd65682f6c 100644 --- a/rust/src/server/src/routes/inference/generate/convert.rs +++ b/rust/src/server/src/routes/inference/generate/convert.rs @@ -26,10 +26,8 @@ pub(super) struct ResponseOptions { pub include_usage: bool, /// Whether the caller asked for usage on every streamed chunk. pub include_continuous_usage: bool, - /// Requested output logprobs (`sampling_params.logprobs`): `None` omits - /// them, `Some(n)` returns `max(n, 1)` candidates per position, `Some(-1)` - /// all of them. - pub logprobs: Option, + /// Whether the caller requested output logprobs on generate choices. + pub include_logprobs: bool, /// Whether the caller requested top-level prompt logprobs. pub include_prompt_logprobs: bool, /// Whether the caller requested final prompt token metadata. @@ -58,7 +56,7 @@ pub(super) fn prepare_generate_request( .as_ref() .and_then(|options| options.continuous_usage_stats) .unwrap_or(false); - let logprobs = request.sampling_params.inner.logprobs; + let include_logprobs = request.sampling_params.inner.logprobs.is_some(); let include_prompt_logprobs = request.sampling_params.inner.prompt_logprobs.is_some(); let return_token_ids = request.return_token_ids.unwrap_or(false); let mut sampling_params = request.sampling_params.inner; @@ -98,7 +96,7 @@ pub(super) fn prepare_generate_request( options: ResponseOptions { include_usage, include_continuous_usage, - logprobs, + include_logprobs, include_prompt_logprobs, return_token_ids, }, diff --git a/tests/basic_correctness/models/test_basic_correctness.py b/tests/basic_correctness/models/test_basic_correctness.py index 47777d6740..5c1ed00fbc 100644 --- a/tests/basic_correctness/models/test_basic_correctness.py +++ b/tests/basic_correctness/models/test_basic_correctness.py @@ -216,15 +216,6 @@ def test_models_distributed( pytest.skip(f"Skip test for {TARGET_TEST_SUITE}") with monkeypatch.context() as monkeypatch_context: - if ( - model == "meta-llama/Llama-3.2-1B-Instruct" - and distributed_executor_backend == "ray" - and attention_backend == "" - and TARGET_TEST_SUITE == "L4" - and enable_prompt_embeds - ): # noqa - pytest.skip("enable_prompt_embeds does not work with ray compiled dag.") - for k, v in extra_env.items(): monkeypatch_context.setenv(k, v) diff --git a/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py b/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py index c32469bfe2..cb421687a4 100644 --- a/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py +++ b/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py @@ -157,10 +157,17 @@ class QKNormRoPEKVCacheTestModel(torch.nn.Module): ) if self.kv_cache_dtype != self.dtype: - self.attn._k_scale = torch.tensor(1.0, dtype=torch.float32, device=device) - self.attn._v_scale = torch.tensor(1.0, dtype=torch.float32, device=device) - self.attn._k_scale_float = 1.0 - self.attn._v_scale_float = 1.0 + # Non-unit scales so the fused and unfused writers must agree on them. + self.attn._k_scale = torch.tensor(0.5, dtype=torch.float32, device=device) + self.attn._v_scale = torch.tensor(0.25, dtype=torch.float32, device=device) + self.attn._k_scale_cpu = torch.tensor( + 0.5, dtype=torch.float32, device="cpu" + ) + self.attn._v_scale_cpu = torch.tensor( + 0.25, dtype=torch.float32, device="cpu" + ) + self.attn._k_scale_float = 0.5 + self.attn._v_scale_float = 0.25 else: self.attn._k_scale = self.attn._k_scale.to(device) self.attn._v_scale = self.attn._v_scale.to(device) @@ -527,16 +534,6 @@ def _run_qk_norm_rope_kvcache_fusion_test( else: torch.testing.assert_close(q_unfused, q_fused, atol=ATOL, rtol=RTOL) - if not expect_fusion or (not is_fp8_cache and mrope_section is None): - # The AITER PTS kernel populates k_out only for non-FP8 caches. - # With FP8, the kernel writes quantized K directly to the cache - # and may leave k_out uninitialised. In production this is fine - # because downstream attention reads K from the cache. - torch.testing.assert_close(k_unfused, k_fused, atol=ATOL, rtol=RTOL) - - # Should be bit exact since no processing had been done on v for both paths - torch.testing.assert_close(v_unfused, v_fused, atol=0.0, rtol=0.0) - # Fused and unfused arithmetic can straddle an FP8 rounding boundary. # Allow one E4M3 quantization step while still comparing every block. if is_fp8_cache: @@ -544,6 +541,18 @@ def _run_qk_norm_rope_kvcache_fusion_test( else: cache_atol, cache_rtol = ATOL, RTOL + # The fused k_out is the dequantized cache write, so it carries the same + # rounding; FA and ROCM_ATTN prefill consume it directly. Unified reads K + # from the cache and asks the kernel not to write k_out at all. + if attn_backend != AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN: + torch.testing.assert_close( + k_unfused, k_fused, atol=cache_atol, rtol=cache_rtol + ) + + # Should be bit exact since no processing had been done on v for both paths + torch.testing.assert_close(v_unfused, v_fused, atol=0.0, rtol=0.0) + + # Whole cache, so the interleaved V write (ROCM_ATTN kv_cache[1]) is checked. torch.testing.assert_close( kv_cache_unfused.float(), kv_cache_fused.float(), @@ -570,20 +579,52 @@ _FUSION_CONFIGS = [ _FUSION_CONFIGS, ) @pytest.mark.parametrize( - "attn_backend", + "attn_backend, use_shuffle_kv_layout, kv_layout", [ - AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN, - AttentionBackendEnum.ROCM_AITER_FA, + # ROCM_ATTN always writes the shuffled, interleaved-V layout and needs each + # K/V side contiguous per block, so head-major only. + pytest.param( + AttentionBackendEnum.ROCM_ATTN, + "0", + KVCacheLayout.LBHNC, + id="rocm_attn", + ), + # Unified never reads the shuffle env; its strided views take any layout. + pytest.param( + AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN, + "0", + KVCacheLayout.LBHNC, + id="unified-head_major", + ), + pytest.param( + AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN, + "0", + KVCacheLayout.LBNHC, + id="unified-token_major", + ), + # FA takes any layout unshuffled; its shuffle writer needs each K/V side + # contiguous per block, so head-major only. + pytest.param( + AttentionBackendEnum.ROCM_AITER_FA, + "0", + KVCacheLayout.LBHNC, + id="fa-head_major", + ), + pytest.param( + AttentionBackendEnum.ROCM_AITER_FA, + "0", + KVCacheLayout.LBNHC, + id="fa-token_major", + ), + pytest.param( + AttentionBackendEnum.ROCM_AITER_FA, + "1", + KVCacheLayout.LBHNC, + id="fa-shuffle", + ), ], ) @pytest.mark.parametrize("num_tokens", [5, 2048]) -@pytest.mark.parametrize( - "kv_layout", - [ - pytest.param(KVCacheLayout.LBHNC, id="head_major"), - pytest.param(KVCacheLayout.LBNHC, id="token_major"), - ], -) @pytest.mark.parametrize("enable_aiter_triton_rope", [True, False]) @pytest.mark.parametrize("block_size", [16, 32, 64]) @pytest.mark.parametrize("dtype", [torch.bfloat16]) @@ -615,6 +656,7 @@ def test_qk_norm_rope_kvcache_fusion( kv_cache_dtype: CacheDType, rms_norm_eps: float, custom_op: str, + use_shuffle_kv_layout: str, monkeypatch: pytest.MonkeyPatch, ): _run_qk_norm_rope_kvcache_fusion_test( @@ -627,7 +669,7 @@ def test_qk_norm_rope_kvcache_fusion( rotary_dim=rotary_dim, block_size=block_size, is_neox=is_neox, - use_shuffle_kv_layout="0", + use_shuffle_kv_layout=use_shuffle_kv_layout, kv_layout=kv_layout, dtype=dtype, kv_cache_dtype=kv_cache_dtype, diff --git a/tests/distributed/test_cpu_all_reduce.py b/tests/distributed/test_cpu_all_reduce.py new file mode 100644 index 0000000000..13e2e818e6 --- /dev/null +++ b/tests/distributed/test_cpu_all_reduce.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for cpu_all_reduce_ (gloo, CPU only).""" + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from vllm.distributed.parallel_state import GroupCoordinator +from vllm.distributed.utils import ( + cpu_all_reduce_, + stateless_init_torch_distributed_process_group, +) +from vllm.utils.network_utils import get_open_port + +pytestmark = pytest.mark.cpu_test + + +def _make_input(rank: int, world_size: int, step: int) -> torch.Tensor: + # Same shape as the V2 DP metadata tensor; values differ per rank/step. + t = torch.zeros(6, world_size, dtype=torch.int32) + t[:, rank] = torch.arange(6, dtype=torch.int32) * (rank + 1) + step + return t + + +def _worker(rank: int, world_size: int, port: int, stateless_port: int): + dist.init_process_group( + "gloo", + init_method=f"tcp://127.0.0.1:{port}", + rank=rank, + world_size=world_size, + ) + coord = GroupCoordinator( + group_ranks=[list(range(world_size))], + local_rank=0, + torch_distributed_backend="gloo", + use_device_communicator=False, + group_name="cpu_all_reduce_test", + ) + + # Several steps to exercise tag reuse across consecutive calls. + for step in range(5): + expected = _make_input(rank, world_size, step) + dist.all_reduce(expected, group=coord.cpu_group) + t = _make_input(rank, world_size, step) + out = coord.cpu_all_reduce_(t) + assert out is t + torch.testing.assert_close(t, expected, rtol=0, atol=0) + + # Float sums must be bitwise identical on every rank. + t = torch.randn(6, world_size, generator=torch.Generator().manual_seed(rank)) + expected = t.clone() + dist.all_reduce(expected, group=coord.cpu_group) + coord.cpu_all_reduce_(t) + torch.testing.assert_close(t, expected) + gathered = [torch.empty_like(t) for _ in range(world_size)] + dist.all_gather(gathered, t, group=coord.cpu_group) + for other in gathered: + assert torch.equal(other, t) + + # Unregistered (stateless) process group, as used by the engine-core DP + # group and elastic EP. + stateless_pg = stateless_init_torch_distributed_process_group( + "127.0.0.1", stateless_port, rank, world_size, backend="gloo" + ) + expected = _make_input(rank, world_size, 0) + dist.all_reduce(expected) + t = _make_input(rank, world_size, 0) + cpu_all_reduce_(t, stateless_pg) + torch.testing.assert_close(t, expected, rtol=0, atol=0) + + dist.destroy_process_group() + + +# Non-power-of-two sizes exercise the cross-block exchange. +@pytest.mark.parametrize("world_size", [2, 3, 4, 6, 8, 9, 12]) +def test_cpu_all_reduce(world_size: int): + mp.spawn( + _worker, + args=(world_size, get_open_port(), get_open_port()), + nprocs=world_size, + join=True, + ) diff --git a/tests/distributed/test_distributed_oot.py b/tests/distributed/test_distributed_oot.py index 5f7f3ffa8a..51ae8b35f8 100644 --- a/tests/distributed/test_distributed_oot.py +++ b/tests/distributed/test_distributed_oot.py @@ -1,6 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json + +import pytest + from tests.plugins_tests.test_oot_registration_online import ( run_and_test_dummy_opt_api_server, ) @@ -8,3 +12,17 @@ from tests.plugins_tests.test_oot_registration_online import ( def test_distributed_oot(dummy_opt_path: str): run_and_test_dummy_opt_api_server(dummy_opt_path, tp=2) + + +@pytest.mark.parametrize("enforce_eager", [True, False]) +def test_distributed_oot_pp_raw_input_tokens(dummy_opt_path: str, enforce_eager: bool): + # The model asserts that input_ids reaches every pipeline stage. + overrides = {"architectures": ["MyOPTRawInputTokensForCausalLM"]} + extra_args = ["-pp", "2", "--hf-overrides", json.dumps(overrides)] + if enforce_eager: + extra_args.append("--enforce-eager") + run_and_test_dummy_opt_api_server( + dummy_opt_path, + extra_args=extra_args, + env_dict={"VLLM_USE_V2_MODEL_RUNNER": "1"}, + ) diff --git a/tests/distributed/test_elastic_ep.py b/tests/distributed/test_elastic_ep.py index 1dc6b5d3bb..d99aaeff7d 100644 --- a/tests/distributed/test_elastic_ep.py +++ b/tests/distributed/test_elastic_ep.py @@ -51,7 +51,7 @@ def _traffic_loop( stop: threading.Event, finished: threading.Event, is_probe: bool = False, -) -> list[tuple[float, float, int | None]]: +) -> list[tuple[float, float, int | None, str | None]]: url = server.url_for("is_scaling_elastic_ep" if is_probe else "v1/completions") payload = {"model": MODEL_NAME, "prompt": "Hello", "max_tokens": 4} headers = None if dp_rank is None else {"X-data-parallel-rank": str(dp_rank)} @@ -60,15 +60,18 @@ def _traffic_loop( is_ready = False while not stop.is_set(): request_start = time.perf_counter() + phase = None try: response = requests.post( url, json=request_payload, headers=headers, timeout=120 ) status_code = response.status_code + if is_probe and status_code == 200: + phase = response.json()["phase"] except requests.exceptions.RequestException: status_code = None - responses.append((request_start, time.perf_counter(), status_code)) - if status_code == 200: + responses.append((request_start, time.perf_counter(), status_code, phase)) + if status_code == 200 and (not is_probe or phase in ("idle", "completed")): if not is_ready: ready.wait(timeout=120) is_ready = True @@ -78,14 +81,16 @@ def _traffic_loop( return responses -def _downtime(responses: list[tuple[float, float, int | None]]) -> float: - rejected = [end for _, end, status in responses if status == 503] - if not rejected: +def _downtime(responses: list[tuple[float, float, int | None, str | None]]) -> float: + committing = [end for _, end, _, phase in responses if phase == "committing"] + if not committing: return 0 recovered = next( - end for _, end, status in responses if status == 200 and end > rejected[-1] + end + for _, end, status, phase in responses + if status == 200 and phase in ("idle", "completed") and end > committing[-1] ) - return recovered - rejected[0] + return recovered - committing[0] def _scale_with_traffic( @@ -124,16 +129,20 @@ def _scale_with_traffic( bad_statuses = { status for responses in [probe_result, *results] - for _, _, status in responses + for _, _, status, _ in responses if status not in (200, 503) } assert not bad_statuses, f"traffic got unexpected statuses {bad_statuses}" - probe_503 = [start for start, _, status in probe_result if status == 503] - assert probe_503, "Scaling probe did not observe commit" + probe_committing = [ + start for start, _, _, phase in probe_result if phase == "committing" + ] + assert probe_committing, "Scaling probe did not observe commit" assert not results or any( - status == 200 and start_time <= request_start and request_end < probe_503[0] + status == 200 + and start_time <= request_start + and request_end < probe_committing[0] for responses in results - for request_start, request_end, status in responses + for request_start, request_end, status, _ in responses ), "No request completed successfully during preparation" print( diff --git a/tests/distributed/test_engram_dp_shard.py b/tests/distributed/test_engram_dp_shard.py index 8b7d3e8cec..6f14859dec 100644 --- a/tests/distributed/test_engram_dp_shard.py +++ b/tests/distributed/test_engram_dp_shard.py @@ -585,6 +585,7 @@ def _check_dummy_hash_model_forward( model = model_ops.DeepseekV4Model.__new__(model_ops.DeepseekV4Model) torch.nn.Module.__init__(model) + model.parallel_config = vllm_config.parallel_config model.use_native_mega_moe = False model.use_sequence_parallel = False model.fuse_mhc_all_reduce = False diff --git a/tests/distributed/test_multi_node_assignment.py b/tests/distributed/test_multi_node_assignment.py deleted file mode 100644 index 6346272429..0000000000 --- a/tests/distributed/test_multi_node_assignment.py +++ /dev/null @@ -1,79 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Make sure ray assigns GPU workers to the correct node. - -Run: -```sh -cd $VLLM_PATH/tests - -pytest distributed/test_multi_node_assignment.py -``` -""" - -import os -import time - -import pytest -import ray -from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy - -from vllm import initialize_ray_cluster -from vllm.config import ParallelConfig -from vllm.utils.network_utils import get_ip -from vllm.v1.executor.ray_utils import available_resources_per_node - -VLLM_MULTI_NODE = os.getenv("VLLM_MULTI_NODE", "0") == "1" - - -@pytest.mark.skipif( - not VLLM_MULTI_NODE, reason="Need at least 2 nodes to run the test." -) -def test_multi_node_assignment() -> None: - # NOTE: important to keep this class definition here - # to let ray use cloudpickle to serialize it. - class Actor: - def get_ip(self): - return get_ip() - - for _ in range(10): - config = ParallelConfig(1, 2) - initialize_ray_cluster(config) - - current_ip = get_ip() - workers = [] - for bundle_id, bundle in enumerate(config.placement_group.bundle_specs): - if not bundle.get("GPU", 0): - continue - scheduling_strategy = PlacementGroupSchedulingStrategy( - placement_group=config.placement_group, - placement_group_capture_child_tasks=True, - placement_group_bundle_index=bundle_id, - ) - - worker = ray.remote( - num_cpus=0, - num_gpus=1, - scheduling_strategy=scheduling_strategy, - )(Actor).remote() - worker_ip = ray.get(worker.get_ip.remote()) - assert worker_ip == current_ip - workers.append(worker) - - for worker in workers: - ray.kill(worker) - - ray.util.remove_placement_group(config.placement_group) - - gpus_needed = len(workers) - deadline = time.monotonic() + 60 - node_id = ray.get_runtime_context().get_node_id() - while True: - free_gpus = available_resources_per_node().get(node_id, {}).get("GPU", 0) - if free_gpus >= gpus_needed: - break - - remaining = deadline - time.monotonic() - if remaining <= 0: - raise TimeoutError("placement group resources were not released") - - time.sleep(min(1, remaining)) diff --git a/tests/entrypoints/cohere/test_api_router.py b/tests/entrypoints/cohere/test_api_router.py index 209ceb4e3e..2c2fce1902 100644 --- a/tests/entrypoints/cohere/test_api_router.py +++ b/tests/entrypoints/cohere/test_api_router.py @@ -598,6 +598,21 @@ class TestCohereErrorEnvelope: assert "message" in body assert isinstance(body["message"], str) and body["message"] + @pytest.mark.parametrize("path", ["/cohere/v2/chat", "/cohere/v2/chat/render"]) + def test_required_tool_choice_without_tools_returns_400(self, path): + """Reject an impossible tool choice before reaching either handler.""" + app = _build_app_with_vllm_handlers(handler=None) + with TestClient(app) as client: + response = client.post( + path, + json={**_minimal_request_body(), "tool_choice": "REQUIRED"}, + ) + assert response.status_code == HTTPStatus.BAD_REQUEST + body = response.json() + assert "error" not in body + assert "REQUIRED" in body["message"] + assert "tools" in body["message"] + def test_http_error_body_is_cohere_shaped(self): # A raised ``HTTPException`` from anywhere in the request cycle # is routed through ``http_exception_handler`` (producing the diff --git a/tests/entrypoints/cohere/test_protocol.py b/tests/entrypoints/cohere/test_protocol.py index dc37ff7cdd..30e625cac2 100644 --- a/tests/entrypoints/cohere/test_protocol.py +++ b/tests/entrypoints/cohere/test_protocol.py @@ -95,6 +95,27 @@ class TestCohereChatV2Request: tool_choice="ANY", # not REQUIRED/NONE ) + @pytest.mark.parametrize("tool_fields", [{}, {"tools": None}, {"tools": []}]) + def test_required_tool_choice_requires_tools(self, tool_fields): + with pytest.raises(VLLMValidationError, match="REQUIRED.*tools") as exc_info: + CohereChatV2Request( + model="m", + messages=[{"role": "user", "content": "hi"}], + tool_choice="REQUIRED", + **tool_fields, + ) + assert exc_info.value.parameter == "tool_choice" + + @pytest.mark.parametrize("tool_choice", [None, "NONE"]) + def test_nonrequired_tool_choice_allows_empty_tools(self, tool_choice): + request = CohereChatV2Request( + model="m", + messages=[{"role": "user", "content": "hi"}], + tool_choice=tool_choice, + tools=[], + ) + assert request.tool_choice == tool_choice + def test_stop_sequences_bounded_by_env_limit(self): messages = [{"role": "user", "content": "hi"}] limit = envs.VLLM_MAX_STOP_STRINGS diff --git a/tests/entrypoints/openai/completion/test_prompt_validation.py b/tests/entrypoints/openai/completion/test_prompt_validation.py index 6c40037e07..7eecf6afbf 100644 --- a/tests/entrypoints/openai/completion/test_prompt_validation.py +++ b/tests/entrypoints/openai/completion/test_prompt_validation.py @@ -64,6 +64,7 @@ def test_load_prompt_embeds( model_config = Mock(spec=ModelConfig) model_config.enable_prompt_embeds = True model_config.get_hidden_size.return_value = hidden_size + model_config.hf_text_config = object() model_config.dtype = dtype # construct arbitrary tensors of various dtypes, layouts, and sizes. diff --git a/tests/entrypoints/scale_out/derender/test_derender_stream.py b/tests/entrypoints/scale_out/derender/test_derender_stream.py index a1c4e2a330..372ef31ea3 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_stream.py +++ b/tests/entrypoints/scale_out/derender/test_derender_stream.py @@ -348,6 +348,85 @@ def test_non_streaming_derender_initializes_parser_prefix( assert choices[0].finish_reason == "length" +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("stop", "include_stop_str_in_output", "dropped"), + [ + ("eos", False, True), + ("stop_token_id", False, True), + ("stop_string", False, False), + ("eos", True, False), + ], +) +async def test_parsed_derender_leaves_out_stop_token_like_engine( + tokenizer, stop, include_stop_str_in_output, dropped +): + """A parser keeps special tokens in the decoded text, so the token that + stopped generation must be dropped the way the engine drops it.""" + eos = tokenizer.eos_token_id + answer_ids = tokenizer.encode("answer", add_special_tokens=False) + bang = tokenizer.encode("!", add_special_tokens=False)[-1] + last_id, stop_reason = { + "eos": (eos, None), + "stop_token_id": (bang, bang), + "stop_string": (bang, "!"), + }[stop] + + parser = MagicMock() + parser.parse.side_effect = lambda text, *args, **kwargs: (None, text, []) + parser.parse_delta.side_effect = lambda text, *args, **kwargs: DeltaMessage( + content=text + ) + request = _chat_request(include_stop_str_in_output=include_stop_str_in_output) + + with ThreadPoolExecutor(max_workers=1) as executor: + renderer = MagicMock() + renderer.get_tokenizer.return_value = tokenizer + renderer.get_eos_token_id.return_value = eos + renderer._executor = executor + model_config = MagicMock() + model_config.hf_config.model_type = "llama" + model_config.model = MODEL_NAME + model_config.try_get_generation_config.return_value = {} + derenderer = OnlineDerenderer( + model_config=model_config, + renderer=renderer, + request_logger=None, + chat_template=None, + chat_template_content_format="string", + ) + derenderer.parser = MagicMock(return_value=parser) + + def batch_content(token_ids: list[int], finish_reason: str) -> str: + choice = GenerateTokensChoice( + index=0, + token_ids=token_ids, + finish_reason=finish_reason, + stop_reason=stop_reason if finish_reason == "stop" else None, + ) + response = GenerateTokensResponse( + request_id="test", prompt_token_ids=[11, 12], choices=[choice] + ) + return derenderer._derender_chat(response, request)[0].message.content + + # Same decode window and seeding; "length" never drops a token. + expected = batch_content( + answer_ids if dropped else answer_ids + [last_id], "length" + ) + batch = batch_content(answer_ids + [last_id], "stop") + stream, _ = await derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk( + answer_ids + [last_id], finish_reason="stop", stop_reason=stop_reason + ), + chat_request=request, + prompt_token_ids=[11, 12], + ) + + assert batch == expected + assert stream.choices[0].delta.content == expected + + class TestDetokenizeDelta: """_detokenize_delta: chunked decode must equal one shot decode.""" diff --git a/tests/entrypoints/scale_out/derender/test_generate_logprobs_conversion.py b/tests/entrypoints/scale_out/derender/test_generate_logprobs_conversion.py index 96d937ad5a..221b2c3994 100644 --- a/tests/entrypoints/scale_out/derender/test_generate_logprobs_conversion.py +++ b/tests/entrypoints/scale_out/derender/test_generate_logprobs_conversion.py @@ -139,7 +139,7 @@ def test_sentencepiece_leading_space_matches_engine(): assert [t.token for t in entry.top_logprobs] == convert_ids_list_to_tokens( tokenizer, [spaced, bare] ) - completion = _convert_chat_logprobs_to_completion_logprobs(resolved) + completion = _convert_chat_logprobs_to_completion_logprobs(resolved, logprobs) assert completion.top_logprobs == [{" true": -0.5, "true": -1.5}] diff --git a/tests/entrypoints/scale_out/derender/test_logprobs_endpoint_parity.py b/tests/entrypoints/scale_out/derender/test_logprobs_endpoint_parity.py new file mode 100644 index 0000000000..fd3910851d --- /dev/null +++ b/tests/entrypoints/scale_out/derender/test_logprobs_endpoint_parity.py @@ -0,0 +1,212 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""generate + derender must return the same `top_logprobs` as the coupled +`/v1/chat/completions` and `/v1/completions` endpoints (#59513). + +Both sides start from the same per-position engine dicts, built with the +engine's own `append_logprobs_for_next_position` and detokenized the way the +engine does, and go through the real coupled serving objects on one side and +`/inference/v1/generate` + derender on the other. The cases include a sampled +token outside the top k, where chat keeps k entries and completions k + 1; the +greedy end-to-end parity tests never produce that case. +""" + +import pytest + +from tests.entrypoints.openai.chat_completion.test_serving_chat import ( + _build_mock_engine, + _build_serving_chat, +) +from tests.entrypoints.openai.completion.test_completion_error import ( + _build_serving_completion, +) +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.completion.protocol import CompletionRequest +from vllm.entrypoints.scale_out.token_in_token_out.protocol import GenerateLogProbs +from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens +from vllm.logprobs import append_logprobs_for_next_position +from vllm.renderers.online_derenderer import ( + _chat_top_logprobs_limit, + _completion_top_logprobs_limit, + _convert_chat_logprobs_to_completion_logprobs, + _resolve_logprobs, +) +from vllm.tokenizers import get_tokenizer +from vllm.tokenizers.detokenizer_utils import convert_ids_list_to_tokens + +MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM" + + +@pytest.fixture(scope="module") +def tokenizer(): + return get_tokenizer(MODEL_NAME) + + +@pytest.fixture(scope="module") +def serving_chat(): + return _build_serving_chat(_build_mock_engine()) + + +@pytest.fixture(scope="module") +def serving_completion(): + return _build_serving_completion(_build_mock_engine()) + + +@pytest.fixture(scope="module") +def vocab_ids(tokenizer): + """Ordinary word pieces (no byte-fallback tokens, which need the U+FFFD + context repair that is compared elsewhere).""" + ids = tokenizer.encode( + "the quick brown fox jumps over a lazy dog while seven cats watch " + "from an old wooden fence near the river bank", + add_special_tokens=False, + ) + pieces = tokenizer.convert_ids_to_tokens(ids) + return list(dict.fromkeys(i for i, p in zip(ids, pieces) if "<0x" not in p)) + + +def _positions(tokenizer, vocab_ids, k, sampled_ranks): + """Engine output for one completion: per position the sampled id and its + dict (sampled first, then ranks 1..k), with decoded tokens as the engine + fills them. k = -1 returns every candidate, as for the whole vocabulary.""" + token_ids, dicts = [], [] + for pos, sampled_rank in enumerate(sampled_ranks): + n_top = max(k, sampled_rank) if k != -1 else len(vocab_ids) - 1 + ranked = [vocab_ids[(pos + j) % len(vocab_ids)] for j in range(n_top + 1)] + sampled = ranked[sampled_rank - 1] + top = ranked[: len(ranked) - 1] if k == -1 else ranked[:k] + ids = [sampled, *top] + logprobs = [-0.25 * sampled_rank] + [-0.25 * (r + 1) for r in range(len(top))] + per_position: list = [] + append_logprobs_for_next_position( + per_position, + ids, + logprobs, + convert_ids_list_to_tokens(tokenizer, ids), + sampled_rank, + len(top), + ) + token_ids.append(sampled) + dicts.append(per_position[0]) + return token_ids, dicts + + +def _generate(token_ids, dicts) -> GenerateLogProbs: + out = ServingTokens._create_tokens_logprobs( + None, token_ids=token_ids, top_logprobs=dicts + ) + # Over the wire, as a derender client would send it back. + return GenerateLogProbs.model_validate_json(out.model_dump_json()) + + +# Sampled token rank per position: top-1, inside the top k, outside it. +RANKS = [1, 2, 1, 7, 3, 9, 1] + + +@pytest.mark.parametrize("k", [0, 1, 2, 5, -1]) +def test_chat_derender_matches_coupled_chat(tokenizer, vocab_ids, serving_chat, k): + token_ids, dicts = _positions(tokenizer, vocab_ids, k, RANKS) + coupled = serving_chat._create_chat_logprobs( + token_ids=token_ids, + top_logprobs=dicts, + tokenizer=tokenizer, + num_output_top_logprobs=k, + ) + request = ChatCompletionRequest( + model=MODEL_NAME, + messages=[{"role": "user", "content": "hi"}], + logprobs=True, + top_logprobs=k, + ) + derendered = _resolve_logprobs( + _generate(token_ids, dicts), + tokenizer, + top_limit=_chat_top_logprobs_limit(request), + ) + assert derendered.model_dump() == coupled.model_dump() + + +@pytest.mark.parametrize("k", [0, 1, 2, 5, -1]) +def test_completion_derender_matches_coupled_completion( + tokenizer, vocab_ids, serving_completion, k +): + token_ids, dicts = _positions(tokenizer, vocab_ids, k, RANKS) + coupled = serving_completion._create_completion_logprobs( + token_ids=token_ids, + top_logprobs=dicts, + num_output_top_logprobs=k, + tokenizer=tokenizer, + ) + request = CompletionRequest(model=MODEL_NAME, prompt="hi", logprobs=k) + generated = _generate(token_ids, dicts) + derendered = _convert_chat_logprobs_to_completion_logprobs( + _resolve_logprobs( + generated, + tokenizer, + top_limit=_completion_top_logprobs_limit(request), + ), + generated, + ) + assert derendered.model_dump() == coupled.model_dump() + if k > 0: + # The case #59513 is about: completions keeps k + 1 entries when the + # sampled token is outside the top k. + outside = [i for i, r in enumerate(RANKS) if r > k] + assert all(len(derendered.top_logprobs[i]) == k + 1 for i in outside) + if k == -1: + # /v1/completions returns every candidate for logprobs=-1 (#55717). + assert [len(top) for top in derendered.top_logprobs] == [len(d) for d in dicts] + + +@pytest.mark.parametrize( + "top_limit", + [ + _chat_top_logprobs_limit(None), + # A request that does not ask for logprobs says nothing about the cut. + _chat_top_logprobs_limit( + ChatCompletionRequest( + model=MODEL_NAME, messages=[{"role": "user", "content": "hi"}] + ) + ), + _completion_top_logprobs_limit(None), + _completion_top_logprobs_limit( + CompletionRequest(model=MODEL_NAME, prompt="hi") + ), + ], +) +def test_without_a_logprobs_request_derender_keeps_every_candidate( + tokenizer, vocab_ids, top_limit +): + token_ids, dicts = _positions(tokenizer, vocab_ids, 2, RANKS) + derendered = _resolve_logprobs( + _generate(token_ids, dicts), tokenizer, top_limit=top_limit + ) + assert [len(e.top_logprobs) for e in derendered.content] == [len(d) for d in dicts] + + +@pytest.mark.parametrize("k", [-1, 0, 2]) +def test_completion_position_without_logprobs_stays_none( + tokenizer, vocab_ids, serving_completion, k +): + """A position the engine returned no logprobs for is `None` in + `/v1/completions`, not `{}`; derender keeps it that way.""" + token_ids, dicts = _positions(tokenizer, vocab_ids, k, RANKS) + dicts[3] = None + coupled = serving_completion._create_completion_logprobs( + token_ids=token_ids, + top_logprobs=dicts, + num_output_top_logprobs=k, + tokenizer=tokenizer, + ) + request = CompletionRequest(model=MODEL_NAME, prompt="hi", logprobs=k) + generated = _generate(token_ids, dicts) + derendered = _convert_chat_logprobs_to_completion_logprobs( + _resolve_logprobs( + generated, + tokenizer, + top_limit=_completion_top_logprobs_limit(request), + ), + generated, + ) + assert derendered.top_logprobs == coupled.top_logprobs + assert derendered.top_logprobs[3] is None diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_tokens_logprobs.py b/tests/entrypoints/scale_out/token_in_token_out/test_tokens_logprobs.py index 05dba167fe..ebbf6ac33c 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_tokens_logprobs.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_tokens_logprobs.py @@ -14,11 +14,11 @@ from vllm.logprobs import Logprob def test_top_logprobs_alternatives_have_own_token_ids(): """Each top_logprobs alternative must carry its own integer token id.""" + # Engine dict at logprobs=2 with the sampled token 262 also top-1. result = ServingTokens._create_tokens_logprobs( None, token_ids=[262], - top_logprobs=[{262: Logprob(-0.1), 257: Logprob(-1.2), 428: Logprob(-2.3)}], - num_output_top_logprobs=2, + top_logprobs=[{262: Logprob(-0.1), 257: Logprob(-1.2)}], ) token_ids = {e.token_id for e in result.content[0].top_logprobs} assert token_ids == {262, 257}, f"got {token_ids}" @@ -32,7 +32,6 @@ def test_sampled_entry_carries_token_id_and_rank(): top_logprobs=[ {262: Logprob(-0.1, rank=1), 257: Logprob(-1.2, rank=2)}, ], - num_output_top_logprobs=2, ) entry = result.content[0] assert entry.token_id == 262 @@ -50,7 +49,6 @@ def test_sampled_token_absent_from_topk_uses_sentinel(): None, token_ids=[5], top_logprobs=[{7: Logprob(-0.9)}], - num_output_top_logprobs=1, ) entry = result.content[0] assert entry.token_id == 5 @@ -60,14 +58,13 @@ def test_sampled_token_absent_from_topk_uses_sentinel(): def test_logprobs_zero_emits_sampled_token(): - """logprobs=0 must still emit 1 entry (the sampled token).""" + """logprobs=0: the engine returns only the sampled token, and it is kept.""" result = ServingTokens._create_tokens_logprobs( None, token_ids=[7], - top_logprobs=[{7: Logprob(-0.9), 8: Logprob(-1.1)}], - num_output_top_logprobs=0, + top_logprobs=[{7: Logprob(-0.9)}], ) - assert len(result.content[0].top_logprobs) == 1 + assert [t.token_id for t in result.content[0].top_logprobs] == [7] def test_logprobs_minus_one_emits_all_tokens(): @@ -75,7 +72,6 @@ def test_logprobs_minus_one_emits_all_tokens(): None, token_ids=[7], top_logprobs=[{7: Logprob(-0.9), 8: Logprob(-1.1)}], - num_output_top_logprobs=-1, ) assert len(result.content[0].top_logprobs) == 2 @@ -86,7 +82,6 @@ def test_text_logprobs_without_tokenizer_are_placeholders_without_bytes(): None, token_ids=[7], top_logprobs=[{7: Logprob(-0.9), 8: Logprob(-1.1)}], - num_output_top_logprobs=2, ) entry = result.content[0] assert entry.token == "token_id:7" @@ -102,7 +97,6 @@ def test_resolved_logprobs_use_engine_decoded_tokens_and_bytes(): top_logprobs=[ {7: Logprob(-0.9, decoded_token="é"), 8: Logprob(-1.1, decoded_token="e")} ], - num_output_top_logprobs=2, tokenizer=tokenizer, ) entry = result.content[0] @@ -121,7 +115,6 @@ def test_resolved_logprobs_decode_tokens_the_engine_did_not(): None, token_ids=[7, 9], top_logprobs=[{7: Logprob(-0.9)}, None], - num_output_top_logprobs=1, tokenizer=tokenizer, ) assert [e.token for e in result.content] == ["<7>", "<9>"] @@ -130,9 +123,9 @@ def test_resolved_logprobs_decode_tokens_the_engine_did_not(): def test_sampled_token_outside_topk_comes_first(): - """The engine puts the sampled token first, then ranks 1..k. When the - sampled token is outside the top k it takes one of the k slots, so rank k - is left out (same as the OpenAI endpoints).""" + """The engine puts the sampled token first, then ranks 1..k. Generate keeps + all k + 1 candidates when the sampled token is outside the top k; derender + applies each endpoint's cut (#59513).""" result = ServingTokens._create_tokens_logprobs( None, token_ids=[50], @@ -143,9 +136,8 @@ def test_sampled_token_outside_topk_comes_first(): 20: Logprob(-1.0, rank=2), } ], - num_output_top_logprobs=2, ) - assert [t.rank for t in result.content[0].top_logprobs] == [5, 1] + assert [t.rank for t in result.content[0].top_logprobs] == [5, 1, 2] def test_logprob_is_required_on_the_wire(): @@ -160,7 +152,6 @@ def test_rank_zero_from_a_nan_logprob_is_none(): None, token_ids=[7], top_logprobs=[{7: Logprob(float("nan"), rank=0), 8: Logprob(-0.2, rank=1)}], - num_output_top_logprobs=1, ) entry = result.content[0] assert entry.rank is None @@ -176,7 +167,6 @@ def test_text_logprobs_clamp_a_nan_logprob(): None, token_ids=[7], top_logprobs=[{7: Logprob(float("nan"), rank=0, decoded_token="a")}], - num_output_top_logprobs=1, tokenizer=MagicMock(), ) entry = result.content[0] diff --git a/tests/kernels/attention/test_attention.py b/tests/kernels/attention/test_attention.py index e3c3ac0b24..a81c973852 100644 --- a/tests/kernels/attention/test_attention.py +++ b/tests/kernels/attention/test_attention.py @@ -122,6 +122,7 @@ def ref_single_query_cached_kv_attention( @pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE) @pytest.mark.parametrize("seed", SEEDS) @pytest.mark.parametrize("device", CUDA_DEVICES) +@pytest.mark.parametrize("use_interleaved_v_cache", [False, True]) def test_paged_attention( kv_cache_factory, num_seqs: int, @@ -133,9 +134,14 @@ def test_paged_attention( kv_cache_dtype: str, seed: int, device: str, + use_interleaved_v_cache: bool, ) -> None: if current_platform.is_navi() and ( - kv_cache_dtype == "fp8" or head_size != 128 or block_size != 16 or use_alibi + kv_cache_dtype == "fp8" + or head_size != 128 + or block_size != 16 + or use_alibi + or use_interleaved_v_cache ): pytest.skip() @@ -201,6 +207,20 @@ def test_paged_attention( last_block_idx = block_tables_lst[seq_idx][seq_len // block_size] value_cache[last_block_idx, :, :, padding_start:] = padding_nan + standard_value_cache = value_cache + if use_interleaved_v_cache: + # Hand the kernel V in the interleaved layout (same 4D shape); the + # reference below gets the standard layout back. + x_v = 16 // value_cache.element_size() + value_cache = ( + value_cache.view( + NUM_BLOCKS, num_kv_heads, head_size, block_size // x_v, x_v + ) + .permute(0, 1, 3, 2, 4) + .contiguous() + .view(NUM_BLOCKS, num_kv_heads, head_size, block_size) + ) + # Call the paged attention kernel. output = torch.empty_like(query) num_partitions = (max_seq_len + PARTITION_SIZE_ROCM - 1) // PARTITION_SIZE_ROCM @@ -234,6 +254,7 @@ def test_paged_attention( kv_cache_dtype, k_scale, v_scale, + use_interleaved_v_cache=use_interleaved_v_cache, ) opcheck( @@ -259,10 +280,13 @@ def test_paged_attention( v_scale, None, "f16", + use_interleaved_v_cache, ), cond=(head_size == 64 and block_size == BLOCK_SIZES[0]), ) + value_cache = standard_value_cache + # Run the reference implementation. if kv_cache_dtype == "fp8": # Convert cache data back to dtype. diff --git a/tests/kernels/attention/test_prefix_prefill.py b/tests/kernels/attention/test_prefix_prefill.py index 3f675e15e6..c93bb9259d 100644 --- a/tests/kernels/attention/test_prefix_prefill.py +++ b/tests/kernels/attention/test_prefix_prefill.py @@ -109,6 +109,7 @@ def create_alibi_causal_mask( @pytest.mark.parametrize("device", CUDA_DEVICES) @pytest.mark.parametrize("sliding_window", SLIDING_WINDOW) @pytest.mark.parametrize("op", OPS) +@pytest.mark.parametrize("interleave_v", [False, True]) @torch.inference_mode() def test_contexted_kv_attention( num_heads: int, @@ -119,6 +120,7 @@ def test_contexted_kv_attention( kv_cache_dtype: str, device: str, op: Callable, + interleave_v: bool, block_size: int = 32, ) -> None: if "fp8" in kv_cache_dtype and not current_platform.has_device_capability(89): @@ -225,6 +227,23 @@ def test_contexted_kv_attention( ) k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device) + extra_op_kwargs = {} + if interleave_v: + # Rewrite the V cache into the interleaved layout (same 4D shape). Both + # context_attention_fwd (prefill) and kernel_paged_attention_2d (decode) + # read it back interleaved. + x_v = 16 // v_cache.element_size() + v_cache = ( + v_cache.view(cache_size, num_kv_heads, head_size, block_size // x_v, x_v) + .permute(0, 1, 3, 2, 4) + .contiguous() + .view(cache_size, num_kv_heads, head_size, block_size) + ) + if op is context_attention_fwd: + extra_op_kwargs["interleaved_v_pack_factor"] = x_v + else: + extra_op_kwargs["use_interleaved_v_cache"] = True + # Warm up the Triton kernel by calling it once before actually measuring # generation time op( @@ -243,6 +262,7 @@ def test_contexted_kv_attention( k_scale, v_scale, sliding_window=sliding_window, + **extra_op_kwargs, ) torch.accelerator.synchronize() start_time = time.time() @@ -262,6 +282,7 @@ def test_contexted_kv_attention( k_scale, v_scale, sliding_window=sliding_window, + **extra_op_kwargs, ) torch.accelerator.synchronize() end_time = time.time() @@ -386,6 +407,7 @@ def test_rocm_backend_sliding_window_includes_boundary_token( @pytest.mark.parametrize("dtype", DTYPES) @pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES) @pytest.mark.parametrize("device", CUDA_DEVICES) +@pytest.mark.parametrize("interleave_v", [False, True]) @torch.inference_mode() def test_contexted_kv_attention_cached_kv( num_heads: int, @@ -394,6 +416,7 @@ def test_contexted_kv_attention_cached_kv( dtype: torch.dtype, kv_cache_dtype: str, device: str, + interleave_v: bool, block_size: int = 32, ) -> None: # Exercises the KV_FROM_CACHE path of context_attention_fwd: the current @@ -487,6 +510,18 @@ def test_contexted_kv_attention_cached_kv( ) k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device) + interleaved_v_pack_factor = 0 + if interleave_v: + # Rewrite the V cache into the interleaved layout (same 4D shape). + interleaved_v_pack_factor = 16 // v_cache.element_size() + x_v = interleaved_v_pack_factor + v_cache = ( + v_cache.view(cache_size, num_kv_heads, head_size, block_size // x_v, x_v) + .permute(0, 1, 3, 2, 4) + .contiguous() + .view(cache_size, num_kv_heads, head_size, block_size) + ) + # Cached-K/V path: current-chunk k and v are None. context_attention_fwd( query, @@ -504,6 +539,7 @@ def test_contexted_kv_attention_cached_kv( k_scale, v_scale, sliding_window=0, + interleaved_v_pack_factor=interleaved_v_pack_factor, ) torch.accelerator.synchronize() @@ -544,8 +580,11 @@ def test_contexted_kv_attention_cached_kv( @pytest.mark.parametrize("device", CUDA_DEVICES) +@pytest.mark.parametrize("interleave_v", [False, True]) @torch.inference_mode() -def test_contexted_kv_attention_cached_kv_block_table_boundary(device: str) -> None: +def test_contexted_kv_attention_cached_kv_block_table_boundary( + device: str, interleave_v: bool +) -> None: # Boundary guard for the KV_FROM_CACHE block-table load. With an # exact-sized block table (row length == number of blocks for the # sequence) and a query that ends the sequence, the last K/V tile has @@ -611,6 +650,18 @@ def test_contexted_kv_attention_cached_kv_block_table_boundary(device: str) -> N ) k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device) + interleaved_v_pack_factor = 0 + if interleave_v: + # Rewrite the V cache into the interleaved layout (same 4D shape). + interleaved_v_pack_factor = 16 // v_cache.element_size() + x_v = interleaved_v_pack_factor + v_cache = ( + v_cache.view(num_blocks, num_kv_heads, head_size, block_size // x_v, x_v) + .permute(0, 1, 3, 2, 4) + .contiguous() + .view(num_blocks, num_kv_heads, head_size, block_size) + ) + # Cached-K/V path: current-chunk k and v are None. context_attention_fwd( query, @@ -628,6 +679,7 @@ def test_contexted_kv_attention_cached_kv_block_table_boundary(device: str) -> N k_scale, v_scale, sliding_window=0, + interleaved_v_pack_factor=interleaved_v_pack_factor, ) torch.accelerator.synchronize() @@ -723,6 +775,7 @@ def test_contexted_kv_attention_cached_kv_alibi_unsupported(device: str) -> None @pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES) @pytest.mark.parametrize("device", CUDA_DEVICES) @pytest.mark.parametrize("op", OPS) +@pytest.mark.parametrize("interleave_v", [False, True]) @torch.inference_mode() def test_contexted_kv_attention_alibi( num_heads: int, @@ -732,6 +785,7 @@ def test_contexted_kv_attention_alibi( kv_cache_dtype: str, device: str, op: Callable, + interleave_v: bool, block_size: int = 32, ) -> None: if "fp8" in kv_cache_dtype and not current_platform.has_device_capability(89): @@ -860,6 +914,23 @@ def test_contexted_kv_attention_alibi( ) k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device) + extra_op_kwargs = {} + if interleave_v: + # Rewrite the V cache into the interleaved layout (same 4D shape). Both + # context_attention_fwd (prefill) and kernel_paged_attention_2d (decode) + # read it back interleaved. + x_v = 16 // v_cache.element_size() + v_cache = ( + v_cache.view(cache_size, num_kv_heads, head_size, block_size // x_v, x_v) + .permute(0, 1, 3, 2, 4) + .contiguous() + .view(cache_size, num_kv_heads, head_size, block_size) + ) + if op is context_attention_fwd: + extra_op_kwargs["interleaved_v_pack_factor"] = x_v + else: + extra_op_kwargs["use_interleaved_v_cache"] = True + # Warm up the Triton kernel by calling it once before actually measuring # generation time op( @@ -878,6 +949,7 @@ def test_contexted_kv_attention_alibi( k_scale, v_scale, alibi_slopes=alibi_slopes, + **extra_op_kwargs, ) torch.accelerator.synchronize() start_time = time.time() @@ -897,6 +969,7 @@ def test_contexted_kv_attention_alibi( k_scale, v_scale, alibi_slopes=alibi_slopes, + **extra_op_kwargs, ) torch.accelerator.synchronize() end_time = time.time() @@ -992,6 +1065,7 @@ def test_contexted_kv_attention_alibi( @pytest.mark.parametrize("device", CUDA_DEVICES) @pytest.mark.parametrize("sliding_window", SLIDING_WINDOW) @pytest.mark.parametrize("op", OPS) +@pytest.mark.parametrize("interleave_v", [False, True]) @torch.inference_mode() def test_contexted_kv_attention_f32( num_heads: int, @@ -1002,6 +1076,7 @@ def test_contexted_kv_attention_f32( kv_cache_dtype: str, device: str, op: Callable, + interleave_v: bool, ) -> None: test_contexted_kv_attention( num_heads, @@ -1012,6 +1087,7 @@ def test_contexted_kv_attention_f32( kv_cache_dtype, device, op, + interleave_v=interleave_v, ) @@ -1023,6 +1099,7 @@ def test_contexted_kv_attention_f32( @pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES) @pytest.mark.parametrize("device", CUDA_DEVICES) @pytest.mark.parametrize("op", OPS) +@pytest.mark.parametrize("interleave_v", [False, True]) @torch.inference_mode() def test_contexted_kv_attention_alibi_f32( num_heads: int, @@ -1032,9 +1109,17 @@ def test_contexted_kv_attention_alibi_f32( kv_cache_dtype: str, device: str, op: Callable, + interleave_v: bool, ) -> None: test_contexted_kv_attention_alibi( - num_heads, num_queries_per_kv, head_size, dtype, kv_cache_dtype, device, op + num_heads, + num_queries_per_kv, + head_size, + dtype, + kv_cache_dtype, + device, + op, + interleave_v=interleave_v, ) @@ -1052,6 +1137,7 @@ NONSTANDARD_BLOCK_SIZE_SHAPES = [ @pytest.mark.parametrize("dtype", DTYPES) @pytest.mark.parametrize("device", CUDA_DEVICES) @pytest.mark.parametrize("op", OPS) +@pytest.mark.parametrize("interleave_v", [False, True]) @torch.inference_mode() def test_qwen3_nonstandard_block_size( num_heads: int, @@ -1061,6 +1147,7 @@ def test_qwen3_nonstandard_block_size( dtype: torch.dtype, device: str, op: Callable, + interleave_v: bool, ) -> None: """Non-power-of-2 pages must match, even when a tile straddles a page.""" if not current_platform.is_rocm(): @@ -1076,4 +1163,5 @@ def test_qwen3_nonstandard_block_size( kv_cache_dtype="auto", device=device, op=op, + interleave_v=interleave_v, ) diff --git a/tests/kernels/attention/test_rocm_aiter_fa.py b/tests/kernels/attention/test_rocm_aiter_fa.py index c3b2c0754f..f6b3fb92cd 100644 --- a/tests/kernels/attention/test_rocm_aiter_fa.py +++ b/tests/kernels/attention/test_rocm_aiter_fa.py @@ -1061,3 +1061,25 @@ def test_aiter_mha_varlen_fp8_kv(dtype): rtol=rtol, ) torch.testing.assert_close(output, ref, atol=atol, rtol=rtol) + + +def test_asm_kv_scale_tables_shared_across_builders(): + """Builders share one AsmKvScaleTables; equal (shape, scale) requests must + return the same dense table, a different scale a different one.""" + from vllm.v1.attention.backends.rocm_aiter_fa import ( + AiterFlashAttentionMetadataBuilder, + AsmKvScaleTables, + ) + + tables = AsmKvScaleTables() + device = torch.device("cuda") + table = tables.get(device, 8, 2, 16, 0.5) + assert tables.get(device, 8, 2, 16, 0.5) is table + assert tables.get(device, 8, 2, 16, 1.0) is not table + assert table.shape == (8, 2, 16) + assert table.dtype == torch.float32 + assert torch.all(table == 0.5) + # The builder class holds the single process-wide instance. + assert isinstance( + AiterFlashAttentionMetadataBuilder.asm_kv_scale_tables, AsmKvScaleTables + ) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 97d8cbaa3e..66414f723e 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -1989,8 +1989,8 @@ def test_dsv41_adaptive_mla_swa_metadata_graph_replay(monkeypatch) -> None: """Replay V4.1 ROCm decode after device-only query reallocation.""" from tests.v1.attention.utils import create_vllm_config from vllm.models.deepseek_v41.amd.rocm import ( - DeepseekV4ROCMAiterSparseSWAMetadataBuilder, DeepseekV41ROCMAiterMLASparseMetadataBuilder, + DeepseekV41ROCMAiterSparseSWAMetadataBuilder, compute_global_topk_ragged_indices_and_indptr, ) from vllm.v1.attention.backend import AttentionCGSupport, CommonAttentionMetadata @@ -2044,7 +2044,7 @@ def test_dsv41_adaptive_mla_swa_metadata_graph_replay(monkeypatch) -> None: mla_builder = DeepseekV41ROCMAiterMLASparseMetadataBuilder( mla_spec, ["c2a"], vllm_config, device ) - swa_builder = DeepseekV4ROCMAiterSparseSWAMetadataBuilder( + swa_builder = DeepseekV41ROCMAiterSparseSWAMetadataBuilder( swa_spec, ["c2a"], vllm_config, device ) assert ( diff --git a/tests/kernels/moe/test_deep_gemm_mega_moe.py b/tests/kernels/moe/test_deep_gemm_mega_moe.py new file mode 100644 index 0000000000..9382ec1c08 --- /dev/null +++ b/tests/kernels/moe/test_deep_gemm_mega_moe.py @@ -0,0 +1,1274 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepGEMM MegaMoE experts: kernel-layout weight loading, shared-expert fusion, +input staging, and numerics against a dequantized reference for every routed +format, activation and shared-expert fusion in the feature matrix.""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe import ( + MEGA_MOE_BF16, + MEGA_MOE_MXFP4, + MEGA_MOE_NVFP4, + MEGA_MOE_WEIGHT_FORMATS, + DeepGemmMegaMoEExperts, + MegaMoEActivation, + MegaMoEWeightSpec, +) +from vllm.model_executor.layers.fused_moe.prepare_megamoe import ( + prepare_megamoe_inputs, + prepare_megamoe_nvfp4_inputs, +) +from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( + bind_routed_experts_capturer, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp8Dynamic, + kNvfp4DynamicToken, +) +from vllm.platforms import current_platform +from vllm.utils.deep_gemm import is_deep_gemm_supported + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda(), reason="DeepGEMM MegaMoE requires CUDA" +) +requires_sm100 = pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="MegaMoE input staging only feeds the SM100 kernels", +) + + +@pytest.fixture(autouse=True) +def sm100_mega_moe_layout(monkeypatch): + """Build the SM100 kernel layout on any GPU; kernels are gated per test.""" + monkeypatch.setattr( + "vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe." + "check_deep_gemm_mega_moe_supported", + lambda *args: None, + ) + + +def _make_experts(**kwargs) -> DeepGemmMegaMoEExperts: + vllm_config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=4), + compilation_config=SimpleNamespace(static_forward_context={}), + ) + kwargs.setdefault("activation", MegaMoEActivation()) + kwargs.setdefault("weight_format", MEGA_MOE_MXFP4) + return DeepGemmMegaMoEExperts(vllm_config, **kwargs) + + +def test_deep_gemm_mega_moe_capture_precedes_eplb(monkeypatch): + experts = DeepGemmMegaMoEExperts.__new__(DeepGemmMegaMoEExperts) + torch.nn.Module.__init__(experts) + experts._synchronize_first_launch = lambda: None + experts.prefix = "model.layers.3.ffn.experts" + experts.max_num_tokens = 4 + experts.num_shared_experts = 0 + experts.capture_fn = None + experts.eplb_state = SimpleNamespace( + logical_to_physical_map=torch.empty(1), + expert_load_view=torch.empty(1), + logical_replica_count=torch.empty(1), + should_record_tensor=torch.empty(1), + num_unpadded_tokens_tensors=None, + ) + + topk_ids = torch.tensor([[1, 2], [3, 4]]) + captured: list[tuple[int, torch.Tensor]] = [] + bind_routed_experts_capturer( + SimpleNamespace(modules=lambda: [experts]), + SimpleNamespace(capture=lambda layer_id, ids: captured.append((layer_id, ids))), + ) + + class MappingReached(Exception): + pass + + def map_ids(**kwargs): + assert captured == [(3, topk_ids)] + raise MappingReached + + monkeypatch.setattr( + f"{DeepGemmMegaMoEExperts.__module__}.eplb_map_to_physical_and_record", + map_ids, + ) + + with pytest.raises(MappingReached): + experts(torch.empty(2, 8), torch.empty(2, 2), topk_ids) + + +def test_deep_gemm_mega_moe_weight_loader_uses_ep_expert_ownership(): + experts = _make_experts( + num_experts=4, + num_local_experts=2, + experts_start_idx=2, + top_k=2, + hidden_size=128, + intermediate_size=128, + ) + + nonlocal_weight = torch.ones(128, 64, dtype=torch.uint8) + assert ( + experts.weight_loader( + experts.w13_weight, + nonlocal_weight, + "experts.w13_weight", + shard_id="w1", + expert_id=1, + return_success=True, + ) + is False + ) + + w1 = torch.full((128, 64), 3, dtype=torch.uint8) + w3 = torch.full((128, 64), 7, dtype=torch.uint8) + w2 = torch.full((128, 64), 11, dtype=torch.uint8) + + assert experts.weight_loader( + experts.w13_weight, + w1, + "experts.w13_weight", + shard_id="w1", + expert_id=2, + return_success=True, + ) + assert experts.weight_loader( + experts.w13_weight, + w3, + "experts.w13_weight", + shard_id="w3", + expert_id=2, + return_success=True, + ) + assert experts.weight_loader( + experts.w2_weight, + w2, + "experts.w2_weight", + shard_id="w2", + expert_id=2, + return_success=True, + ) + + interleaved = experts.w13_weight[0].view(torch.uint8).view(16, 2, 8, 64) + assert torch.equal(interleaved[:, 0], w1.view(16, 8, 64)) + assert torch.equal(interleaved[:, 1], w3.view(16, 8, 64)) + assert torch.equal(experts.w2_weight[0].view(torch.uint8), w2) + assert torch.count_nonzero(experts.w13_weight[1]) == 0 + + +def _reference_mega_moe_layout(w13, s13, w2, s2): + """CPU copy of DeepGEMM's transform_weights_for_mega_moe on packed scales.""" + + def interleave(t): + g, n, k = t.shape + halves = t.view(g, 2, n // 16, 8, k) + return torch.stack([halves[:, 0], halves[:, 1]], dim=2).reshape(g, n, k) + + def utccp(sf): + g, mn, k = sf.shape + return sf.reshape(g, -1, 4, 32, k).transpose(2, 3).reshape(g, mn, k) + + return ( + interleave(w13), + utccp(interleave(s13.view(torch.int32))), + w2, + utccp(s2.view(torch.int32)), + ) + + +def _make_mega_moe_experts( + num_local_experts, num_logical_experts, weight_format=MEGA_MOE_MXFP4 +): + return _make_experts( + num_experts=num_local_experts, + num_local_experts=num_local_experts, + experts_start_idx=0, + top_k=1, + hidden_size=256, + intermediate_size=384, + num_logical_experts=num_logical_experts, + weight_format=weight_format, + ) + + +def _load_mega_moe_checkpoint_shuffled(experts, num_logical_experts, seed): + hidden, inter = experts.hidden_size, experts.intermediate_size + routed = experts.weight_format.routed + group = routed.scale_group_k + shapes = { + "w13_weight": (num_logical_experts, 2 * inter, hidden // 2), + "w13_weight_scale": (num_logical_experts, 2 * inter, hidden // group), + "w2_weight": (num_logical_experts, hidden, inter // 2), + "w2_weight_scale": (num_logical_experts, hidden, inter // group), + } + generator = torch.Generator().manual_seed(seed) + ckpt = { + name: torch.randint(0, 256, shape, dtype=torch.uint8, generator=generator) + for name, shape in shapes.items() + } + if routed.has_global_scale: + for name in ("w13_weight_scale_2", "w2_weight_scale_2", "w2_input_scale"): + ckpt[name] = torch.rand(num_logical_experts, generator=generator) + 0.5 + shards = [] + for expert_id in range(num_logical_experts): + for name, tensor in ckpt.items(): + if name.startswith("w2_"): + shards.append((name, tensor[expert_id], "w2", expert_id)) + elif tensor.dim() == 1: + # w1 and w3 share one global scale. + shards.append((name, tensor[expert_id], "w1", expert_id)) + shards.append((name, tensor[expert_id], "w3", expert_id)) + else: + shards.append((name, tensor[expert_id, :inter], "w1", expert_id)) + shards.append((name, tensor[expert_id, inter:], "w3", expert_id)) + for i in torch.randperm(len(shards), generator=generator).tolist(): + name, loaded_weight, shard_id, expert_id = shards[i] + if name.endswith("weight_scale"): + loaded_weight = loaded_weight.view(routed.scale_dtype) + assert experts.weight_loader( + getattr(experts, name), + loaded_weight, + f"experts.{name}", + shard_id=shard_id, + expert_id=expert_id, + return_success=True, + ) + return ckpt + + +def _expected_mega_moe_layout(ckpt, num_local_experts, num_logical_experts): + physical = torch.arange(num_local_experts) % num_logical_experts + names = ("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale") + ckpt = {name: ckpt[name][physical] for name in names} + expected = _reference_mega_moe_layout(*ckpt.values()) + if is_deep_gemm_supported(): + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + w13, s13, w2, s2 = (t.cuda() for t in ckpt.values()) + hidden, inter = w2.shape[1], w2.shape[2] * 2 + if s13.shape[-1] * 16 == hidden: + # NVFP4: E4M3 bytes packed four per int32, MN-major. + scales = [ + sf.view(torch.int32).transpose(1, 2).contiguous().transpose(1, 2) + for sf in (s13, s2) + ] + else: + scales = [ + deep_gemm.transform_sf_into_required_layout( + (sf.to(torch.int32) << 23).view(torch.float32), + mn, + k, + (1, 32), + num_local_experts, + ) + for sf, mn, k in ((s13, 2 * inter, hidden), (s2, hidden, inter)) + ] + (l1, l1_sf), (l2, l2_sf) = deep_gemm.transform_weights_for_mega_moe( + (w13.view(torch.int8), scales[0]), (w2.view(torch.int8), scales[1]) + ) + for actual, reference in zip((l1, l1_sf, l2, l2_sf), expected): + assert torch.equal(actual.cpu().view(reference.dtype), reference) + return expected + + +def _assert_mega_moe_layout(experts, expected): + actual = (*experts.l1_weights, *experts.l2_weights) + for tensor, reference in zip(actual, expected): + assert torch.equal(tensor.view(reference.dtype), reference) + + +@pytest.mark.parametrize("weight_format", [MEGA_MOE_MXFP4, MEGA_MOE_NVFP4]) +def test_deep_gemm_mega_moe_loads_shards_into_kernel_layout(weight_format): + """Shards in any order, including EPLB replicas, land in DeepGEMM's layout.""" + experts = _make_mega_moe_experts(4, 3, weight_format) + ckpt = _load_mega_moe_checkpoint_shuffled(experts, 3, seed=0) + _assert_mega_moe_layout(experts, _expected_mega_moe_layout(ckpt, 4, 3)) + # DeepGEMM's MN-major scale strides. + group_k = 4 * weight_format.routed.scale_group_k + assert experts.w13_weight_scale.stride() == (768 * 256 // group_k, 1, 768) + assert experts.w2_weight_scale.stride() == (256 * 384 // group_k, 1, 256) + if weight_format.routed.has_global_scale: + physical = torch.arange(4) % 3 + for name in ("w13_weight_scale_2", "w2_weight_scale_2", "w2_input_scale"): + assert torch.equal(getattr(experts, name), ckpt[name][physical]) + assert experts.l2_activation_scale == ckpt["w2_input_scale"].max().item() + + +@pytest.mark.parametrize("weight_format", [MEGA_MOE_MXFP4, MEGA_MOE_NVFP4]) +def test_deep_gemm_mega_moe_reload_overwrites_kernel_layout_in_place(weight_format): + """A second load must update the live kernel tensors without reallocating.""" + experts = _make_mega_moe_experts(2, 2, weight_format) + _load_mega_moe_checkpoint_shuffled(experts, 2, seed=0) + data_ptrs = [param.data_ptr() for param in experts.parameters()] + ckpt = _load_mega_moe_checkpoint_shuffled(experts, 2, seed=1) + assert [param.data_ptr() for param in experts.parameters()] == data_ptrs + _assert_mega_moe_layout(experts, _expected_mega_moe_layout(ckpt, 2, 2)) + + +def test_deep_gemm_mega_moe_nvfp4_l2_scale_covers_every_expert(): + """Every rank quantizes L2 inputs with the max input scale over all experts, + including those it does not own, so EPLB can move experts between ranks.""" + experts = _make_experts( + num_experts=4, + num_local_experts=2, + experts_start_idx=2, + top_k=2, + hidden_size=128, + intermediate_size=128, + weight_format=MEGA_MOE_NVFP4, + ) + assert experts.l2_activation_scale == 1.0 + for expert_id, scale in enumerate((0.5, 4.0, 1.0, 2.0)): + loaded = experts.weight_loader( + experts.w2_input_scale, + torch.tensor(scale), + "experts.w2_input_scale", + shard_id="w2", + expert_id=expert_id, + return_success=True, + ) + assert loaded == (expert_id >= 2) + assert experts.l2_activation_scale == 4.0 + assert experts.w2_input_scale.tolist() == [1.0, 2.0] + # A reload replaces the expert's scale rather than keeping the old max. + experts.weight_loader( + experts.w2_input_scale, + torch.tensor(0.25), + "experts.w2_input_scale", + shard_id="w2", + expert_id=1, + ) + assert experts.l2_activation_scale == 2.0 + # Global scales are expert weights that EPLB moves. + assert len(experts.get_expert_weights()) == 7 + + +def test_deep_gemm_mega_moe_nvfp4_rejects_mismatched_l1_scales(): + """The kernel takes one L1 alpha per expert, so w1 and w3 must agree.""" + experts = _make_mega_moe_experts(1, 1, MEGA_MOE_NVFP4) + + def load(shard_id, scale): + experts.weight_loader( + experts.w13_weight_scale_2, + torch.tensor(scale), + "experts.w13_weight_scale_2", + shard_id=shard_id, + expert_id=0, + ) + + load("w1", 0.5) + load("w3", 0.5) + # A reload may change the scale while the old pair is still live. + load("w3", 2.0) + load("w1", 2.0) + load("w1", 1.0) + with pytest.raises(ValueError, match="one L1 scale per expert"): + load("w3", 4.0) + + +def test_deep_gemm_mega_moe_nvfp4_meta_load_skips_host_scales(): + """Reload dry runs load meta stand-ins, which must not reach host state.""" + experts = _make_mega_moe_experts(1, 1, MEGA_MOE_NVFP4) + for name, shard_id in ( + ("w2_input_scale", "w2"), + ("w13_weight_scale_2", "w1"), + ("w13_weight_scale_2", "w3"), + ): + assert experts.weight_loader( + getattr(experts, name).to("meta"), + torch.empty((), device="meta"), + f"experts.{name}", + shard_id=shard_id, + expert_id=0, + return_success=True, + ) + assert experts.l2_activation_scale == 1.0 + + +@pytest.mark.parametrize("intermediate_size", [512, 2304]) +def test_deep_gemm_mega_moe_preserves_checkpoint_dimensions(intermediate_size): + """Keep native widths so V4.1 routed and shared experts can fuse.""" + experts = _make_experts( + num_experts=2, + num_local_experts=2, + experts_start_idx=0, + top_k=1, + hidden_size=128, + intermediate_size=intermediate_size, + ) + + assert experts.intermediate_size == intermediate_size + assert experts.w13_weight.shape == (2, 2 * intermediate_size, 64) + assert experts.w13_weight_scale.shape == (2, 2 * intermediate_size, 1) + assert experts.w2_weight.shape == (2, 128, intermediate_size // 2) + assert experts.w2_weight_scale.shape == (2, 128, intermediate_size // 128) + transformed = (*experts.l1_weights, *experts.l2_weights) + for actual, param in zip(transformed, experts.parameters()): + assert actual.data_ptr() == param.data_ptr() + + +def _shared_mlp( + hidden_size, + shared_size, + weight_dtype=torch.float8_e4m3fn, + scale_dtype=torch.float8_e8m0fnu, + scale_block=(1, 32), +): + def linear(n, k): + scale = None + if scale_dtype is not None: + scale = torch.empty(n // scale_block[0], k // scale_block[1]).to( + scale_dtype + ) + return SimpleNamespace( + weight=torch.empty(n, k, dtype=weight_dtype), weight_scale_inv=scale + ) + + return SimpleNamespace( + gate_up_proj=linear(2 * shared_size, hidden_size), + down_proj=linear(hidden_size, shared_size), + ) + + +def test_deep_gemm_mega_moe_loads_shared_experts_into_kernel_layout(monkeypatch): + """Shared expert shards land directly in the MegaMoE kernel layout.""" + monkeypatch.setattr( + DeepGemmMegaMoEExperts, "_kernel_fuses_shared_experts", lambda self, spec: True + ) + experts = _make_experts( + num_experts=2, + num_local_experts=2, + experts_start_idx=0, + top_k=1, + hidden_size=256, + intermediate_size=384, + shared_experts=_shared_mlp(256, 384), + ) + assert experts.has_fused_shared_experts + hidden, inter = 256, 384 + generator = torch.Generator().manual_seed(42) + + def rand(*shape): + return torch.randint(0, 256, shape, dtype=torch.uint8, generator=generator) + + w1, w3, w2 = rand(inter, hidden), rand(inter, hidden), rand(hidden, inter) + s1, s3, s2 = ( + rand(inter, hidden // 32), + rand(inter, hidden // 32), + rand(hidden, inter // 32), + ) + fp8, e8m0 = torch.float8_e4m3fn, torch.float8_e8m0fnu + # The down shard arrives as `.w2.` (MTP, DSpark) or `.down_proj.` (mapped). + for name, data in [ + ("shared_experts.w1.weight", w1.view(fp8)), + ("shared_experts.w3.weight", w3.view(fp8)), + ("shared_experts.down_proj.weight", w2.view(fp8)), + ("shared_experts.w1.weight_scale", s1.view(e8m0)), + ("shared_experts.w3.weight_scale", s3.view(e8m0)), + ("shared_experts.w2.weight_scale", s2.view(e8m0)), + ]: + loaded = experts.load_shared_expert_weight(f"ffn.{name}", data) + # Weight tracking checks this name against the registered parameters. + assert loaded in {f"ffn.experts.{n}" for n, _ in experts.named_parameters()} + + def l1_halves(): + return experts.shared_l1_weight.view(torch.uint8).view(-1, 2, 8, hidden) + + # Every tensor matches DeepGEMM's batched transform of the checkpoint. + expected = _reference_mega_moe_layout( + *(t.unsqueeze(0) for t in (torch.cat([w1, w3]), torch.cat([s1, s3]), w2, s2)) + ) + actual = ( + experts.shared_l1_weight, + experts.shared_l1_weight_scale, + experts.shared_l2_weight, + experts.shared_l2_weight_scale, + ) + for tensor, reference in zip(actual, expected): + assert torch.equal(tensor.view(reference.dtype), reference[0]) + # A reload of one shard overwrites in place and leaves the other half alone. + ptr = experts.shared_l1_weight.data_ptr() + w1b = rand(inter, hidden) + experts.load_shared_expert_weight("shared_experts.w1.weight", w1b.view(fp8)) + assert experts.shared_l1_weight.data_ptr() == ptr + assert torch.equal(l1_halves()[:, 0], w1b.view(-1, 8, hidden)) + assert torch.equal(l1_halves()[:, 1], w3.view(-1, 8, hidden)) + with pytest.raises(ValueError): + experts.load_shared_expert_weight("shared_experts.gate_up_proj.weight", w1) + + +def _fake_deep_gemm(with_shared: bool): + def old_buffer(group): ... + + def old_kernel(y): ... + + def new_buffer(group, num_shared_experts=0, shared_dtype=None): ... + + def new_kernel(y, shared_l1_weights=None, shared_l2_weights=None): ... + + kernel = new_kernel if with_shared else old_kernel + buffer = new_buffer if with_shared else old_buffer + return SimpleNamespace( + get_symm_buffer_for_mega_moe=buffer, + NVFP4SymmBuffer=buffer, + fp8_fp4_mega_moe=kernel, + bf16_mega_moe=kernel, + nvfp4_mega_moe=kernel, + get_block_m_for_mega_moe=lambda *args: 0, + get_block_m_for_nvfp4_mega_moe=lambda *args: 0, + ) + + +@pytest.mark.parametrize( + "weight_format,weight_dtype,scale_dtype,scale_block,shared_size,api,fused", + [ + ( + MEGA_MOE_MXFP4, + torch.float8_e4m3fn, + torch.float8_e8m0fnu, + (128, 128), + 128, + True, + True, + ), + (MEGA_MOE_MXFP4, torch.float8_e4m3fn, torch.uint8, (1, 32), 256, True, True), + ( + MEGA_MOE_MXFP4, + torch.float8_e4m3fn, + torch.float8_e8m0fnu, + (128, 128), + 128, + False, + False, + ), + (MEGA_MOE_MXFP4, torch.bfloat16, None, None, 128, True, False), + ( + MEGA_MOE_MXFP4, + torch.float8_e4m3fn, + torch.float32, + (128, 128), + 128, + True, + False, + ), + (MEGA_MOE_MXFP4, torch.float8_e4m3fn, torch.uint8, (1, 16), 128, True, False), + # A TP-sharded shared MLP is not a whole multiple of the routed expert. + (MEGA_MOE_MXFP4, torch.float8_e4m3fn, torch.uint8, (1, 32), 64, True, False), + (MEGA_MOE_BF16, torch.bfloat16, None, None, 128, True, True), + (MEGA_MOE_BF16, torch.float8_e4m3fn, torch.uint8, (1, 32), 128, True, False), + (MEGA_MOE_NVFP4, torch.float8_e4m3fn, torch.uint8, (1, 32), 256, True, True), + (MEGA_MOE_NVFP4, torch.bfloat16, None, None, 128, True, True), + (MEGA_MOE_NVFP4, torch.bfloat16, None, None, 128, False, False), + # NVFP4-style E4M3 block scales are not MXFP8. + ( + MEGA_MOE_NVFP4, + torch.float8_e4m3fn, + torch.float8_e4m3fn, + (1, 16), + 128, + True, + False, + ), + ], +) +def test_deep_gemm_mega_moe_decides_shared_fusion_at_construction( + monkeypatch, + weight_format, + weight_dtype, + scale_dtype, + scale_block, + shared_size, + api, + fused, +): + """Fuse only shared experts in the format's shared precision, shaped as whole + routed experts, and only when the installed kernel takes them.""" + monkeypatch.setattr( + "vllm.utils.deep_gemm._import_deep_gemm", lambda: _fake_deep_gemm(api) + ) + experts = _make_experts( + num_experts=2, + num_local_experts=2, + experts_start_idx=0, + top_k=1, + hidden_size=256, + intermediate_size=128, + weight_format=weight_format, + shared_experts=_shared_mlp( + 256, shared_size, weight_dtype, scale_dtype, scale_block + ), + ) + + assert experts.num_shared_experts == (shared_size // 128 if fused else 0) + assert hasattr(experts, "shared_l1_weight") == fused + if fused: + assert experts.shared_spec in weight_format.shared + assert experts.shared_l1_weight.dtype == weight_dtype + + +def test_deep_gemm_mega_moe_rejects_unsupported_activation(): + with pytest.raises(NotImplementedError, match="situ"): + _make_experts( + num_experts=2, + num_local_experts=2, + experts_start_idx=0, + top_k=1, + hidden_size=128, + intermediate_size=128, + weight_format=MEGA_MOE_BF16, + activation=MegaMoEActivation("situ"), + ) + with pytest.raises(ValueError, match="SiTU takes no clamp"): + MegaMoEActivation("situ", clamp=10.0) + + +@pytest.mark.parametrize("fused", [False, True]) +def test_deep_gemm_mega_moe_adds_unfused_shared_expert_once(monkeypatch, fused): + """The fusion decided at init fixes the launch; forward adds only an + unfused shared expert, and refuses one the layer already fused.""" + monkeypatch.setattr( + DeepGemmMegaMoEExperts, "_kernel_fuses_shared_experts", lambda self, spec: True + ) + launches = [] + + def kernel(y, l1, l2, symm_buffer, **kwargs): + launches.append(kwargs) + y.fill_(1) + + monkeypatch.setattr( + "vllm.utils.deep_gemm._import_deep_gemm", + lambda: SimpleNamespace( + fp8_fp4_mega_moe=kernel, get_block_m_for_mega_moe=lambda *args: 8 + ), + ) + module = "vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe" + monkeypatch.setattr(f"{module}.prepare_megamoe_inputs", lambda *a, **k: None) + monkeypatch.setattr(f"{module}.get_ep_group", lambda: SimpleNamespace(world_size=1)) + shared = _shared_mlp(128, 128 if fused else 64) + experts = _make_experts( + num_experts=2, + num_local_experts=2, + experts_start_idx=0, + top_k=1, + hidden_size=128, + intermediate_size=128, + shared_experts=shared, + ) + assert experts.has_fused_shared_experts == fused + symm_buffer = SimpleNamespace( + x=torch.empty(4, 128), + x_sf=torch.empty(4, 1), + topk_idx=torch.empty(4, 1), + topk_weights=torch.empty(4, 1), + shared_l1_acts_sf=torch.empty(128, 1), + num_max_tokens_per_rank=4, + ) + experts.get_symm_buffer = lambda: symm_buffer + experts._synchronize_first_launch = lambda: None + + def shared_experts(x): + return torch.full_like(x, 2) + + hidden_states = torch.zeros(2, 128, dtype=torch.bfloat16) + topk = (torch.ones(2, 1), torch.zeros(2, 1, dtype=torch.int64)) + if fused: + with pytest.raises(ValueError, match="pass shared_experts=None"): + experts(hidden_states, *topk, shared_experts) + output = experts(hidden_states, *topk) + else: + output = experts(hidden_states, *topk, shared_experts) + + assert torch.all(output == (1 if fused else 3)) + assert ("shared_l1_weights" in launches[-1]) == fused + + +@requires_sm100 +@pytest.mark.parametrize("num_tokens", [7, 63, 64, 65, 127, 128, 135, 434, 16384]) +@pytest.mark.parametrize("hidden_size", [256, 5120]) +def test_prepare_megamoe_inputs_is_bitwise_exact(num_tokens, hidden_size): + from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8 + + device = torch.device("cuda") + top_k = 8 + + generator = torch.Generator(device=device) + generator.manual_seed(0) + hidden_states = ( + torch.randn( + num_tokens, + hidden_size, + device=device, + dtype=torch.float32, + generator=generator, + ) + * 17.0 + ).to(torch.bfloat16) + hidden_states[0, :32] = 0 + hidden_states[1, 32:64] = 1.0e-6 + hidden_states[2, 64:96] = -1.0e-6 + + topk_ids = torch.randint( + 0, + 256, + (num_tokens, top_k), + device=device, + dtype=torch.int32, + generator=generator, + ) + topk_weights = torch.randn( + num_tokens, + top_k, + device=device, + dtype=torch.float32, + generator=generator, + ) + + ref_x, ref_x_sf = per_token_cast_to_fp8( + hidden_states, + use_ue8m0=True, + gran_k=32, + use_packed_ue8m0=True, + ) + ref_topk_idx = topk_ids.to(torch.int64) + ref_topk_weights = topk_weights.clone() + + fused_x = torch.empty_like(ref_x) + fused_x_sf = torch.empty_like(ref_x_sf) + fused_topk_idx = torch.empty_like(ref_topk_idx) + fused_topk_weights = torch.empty_like(ref_topk_weights) + + prepare_megamoe_inputs( + hidden_states, + topk_weights, + topk_ids, + fused_x, + fused_x_sf, + fused_topk_idx, + fused_topk_weights, + ) + torch.accelerator.synchronize() + + assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8)) + assert torch.equal(fused_x_sf, ref_x_sf) + assert torch.equal(fused_topk_idx, ref_topk_idx) + assert torch.equal( + fused_topk_weights.view(torch.uint8), + ref_topk_weights.view(torch.uint8), + ) + + +def _shared_sf_tma_layout(sf: torch.Tensor, block_m: int) -> torch.Tensor: + """Scatter per-token scales into DeepGEMM's padded MN-major UTCCP pages.""" + aligned_block_m = (block_m + 127) // 128 * 128 + num_rows = (sf.shape[0] + block_m - 1) // block_m * aligned_block_m + m = torch.arange(sf.shape[0], device=sf.device) + local = m % block_m + rows = ( + m // block_m * aligned_block_m + + local // 128 * 128 + + local % 32 * 4 + + local % 128 // 32 + ) + out = sf.new_zeros(num_rows, sf.shape[1]) + out[rows] = sf + return out + + +@requires_sm100 +@pytest.mark.parametrize("shared_block_m", [8, 32, 96, 128, 192]) +def test_prepare_megamoe_inputs_stages_shared_scale_tma_layout(shared_block_m): + from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8 + + device = torch.device("cuda") + num_tokens = shared_block_m + 7 + hidden_size = 256 + top_k = 8 + generator = torch.Generator(device=device) + generator.manual_seed(shared_block_m) + hidden_states = torch.randn( + num_tokens, + hidden_size, + device=device, + dtype=torch.bfloat16, + generator=generator, + ) + topk_ids = torch.randint( + 0, + 256, + (num_tokens, top_k), + device=device, + dtype=torch.int32, + generator=generator, + ) + topk_weights = torch.randn( + num_tokens, + top_k, + device=device, + dtype=torch.float32, + generator=generator, + ) + + ref_x, ref_x_sf = per_token_cast_to_fp8( + hidden_states, + use_ue8m0=True, + gran_k=32, + use_packed_ue8m0=True, + ) + ref_shared_x_sf = _shared_sf_tma_layout(ref_x_sf, shared_block_m) + num_shared_rows = ref_shared_x_sf.shape[0] + + fused_x = torch.empty_like(ref_x) + fused_x_sf = torch.empty_like(ref_x_sf) + fused_shared_storage = torch.full( + (hidden_size // 128, num_shared_rows), + -1, + dtype=torch.int32, + device=device, + ) + fused_shared_x_sf = fused_shared_storage.t() + fused_topk_idx = torch.empty_like(topk_ids, dtype=torch.int64) + fused_topk_weights = torch.empty_like(topk_weights) + + prepare_megamoe_inputs( + hidden_states, + topk_weights, + topk_ids, + fused_x, + fused_x_sf, + fused_topk_idx, + fused_topk_weights, + shared_x_sf=fused_shared_x_sf, + shared_block_m=shared_block_m, + ) + torch.accelerator.synchronize() + + populated = ref_shared_x_sf != 0 + assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8)) + assert torch.equal(fused_x_sf, ref_x_sf) + assert torch.equal(fused_shared_x_sf[populated], ref_shared_x_sf[populated]) + assert torch.all(fused_shared_x_sf[~populated] == -1) + + +@requires_sm100 +@pytest.mark.parametrize("quantize", [True, False]) +@pytest.mark.parametrize("nonfinite_padding", [False, True]) +@pytest.mark.parametrize("num_tokens", [7, 65, 135]) +def test_prepare_megamoe_inputs_masks_padding(nonfinite_padding, num_tokens, quantize): + """Padded tokens are unrouted; BF16 dispatch copies the hidden states.""" + from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8 + + device = torch.device("cuda") + hidden_size = 256 + top_k = 8 + + generator = torch.Generator(device=device) + generator.manual_seed(1) + hidden_states = torch.randn( + num_tokens, + hidden_size, + device=device, + dtype=torch.bfloat16, + generator=generator, + ) + topk_ids = torch.randint( + 0, + 256, + (num_tokens, top_k), + device=device, + dtype=torch.int32, + generator=generator, + ) + topk_weights = torch.randn( + num_tokens, + top_k, + device=device, + dtype=torch.float32, + generator=generator, + ) + is_padding = torch.zeros(num_tokens, device=device, dtype=torch.bool) + is_padding[1::3] = True + is_padding[-1] = True + if nonfinite_padding: + topk_weights[is_padding] = float("nan") + + if quantize: + ref_x, ref_x_sf = per_token_cast_to_fp8( + hidden_states, + use_ue8m0=True, + gran_k=32, + use_packed_ue8m0=True, + ) + else: + ref_x, ref_x_sf = hidden_states, None + ref_topk_idx = topk_ids.to(torch.int64) + ref_topk_idx[is_padding] = -1 + ref_topk_weights = topk_weights.clone() + ref_topk_weights[is_padding] = 0.0 + + fused_x = torch.empty_like(ref_x) + fused_x_sf = None if ref_x_sf is None else torch.empty_like(ref_x_sf) + fused_topk_idx = torch.empty_like(ref_topk_idx) + fused_topk_weights = torch.empty_like(ref_topk_weights) + + prepare_megamoe_inputs( + hidden_states, + topk_weights, + topk_ids, + fused_x, + fused_x_sf, + fused_topk_idx, + fused_topk_weights, + is_padding=is_padding, + hidden_quant=kMxfp8Dynamic if quantize else None, + ) + torch.accelerator.synchronize() + + assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8)) + if quantize: + assert torch.equal(fused_x_sf, ref_x_sf) + assert torch.equal(fused_topk_idx, ref_topk_idx) + assert torch.equal( + fused_topk_weights.view(torch.uint8), + ref_topk_weights.view(torch.uint8), + ) + + +@requires_sm100 +@pytest.mark.parametrize("shared", [None, "mxfp8", "bf16"]) +@pytest.mark.parametrize("num_tokens", [0, 7, 65, 135]) +def test_prepare_megamoe_nvfp4_inputs_is_bitwise_exact(num_tokens, shared): + """Per-token NVFP4 dispatch, shared-expert inputs, unrouted padding and the + L2 alphas match DeepGEMM's reference quantizers bit for bit.""" + from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8 + from vllm.third_party.deep_gemm.utils.nvfp4 import per_token_cast_to_nvfp4 + + device = torch.device("cuda") + hidden_size, top_k, shared_block_m = 5120, 8, 32 + generator = torch.Generator(device=device).manual_seed(num_tokens) + # Rows span four decades, which one tensor-wide scale cannot cover. + row_scale = torch.logspace(-3, 1, max(num_tokens, 1), device=device) + hidden_states = ( + torch.randn(num_tokens, hidden_size, device=device, generator=generator) + * row_scale[:num_tokens, None] + ).bfloat16() + if num_tokens: + hidden_states[0] = 0 + topk_ids = torch.randint( + 0, 256, (num_tokens, top_k), device=device, generator=generator + ).int() + topk_weights = torch.rand(num_tokens, top_k, device=device, generator=generator) + is_padding = torch.zeros(num_tokens, device=device, dtype=torch.bool) + is_padding[1::3] = True + w2_global_scale = torch.rand(6, device=device, generator=generator) + + x = torch.empty(num_tokens, hidden_size // 2, dtype=torch.int8, device=device) + x_sf = torch.empty(num_tokens, hidden_size // 64, dtype=torch.int32, device=device) + x_scales = torch.empty(num_tokens, device=device) + topk_idx_out = torch.empty(num_tokens, top_k, dtype=torch.int64, device=device) + topk_weights_out = torch.empty_like(topk_weights) + l2_alpha = torch.empty_like(w2_global_scale) + shared_x = shared_x_sf = None + if shared is not None: + dtype = torch.float8_e4m3fn if shared == "mxfp8" else torch.bfloat16 + shared_x = torch.empty(num_tokens, hidden_size, dtype=dtype, device=device) + if shared == "mxfp8": + ref_shared_x, ref_shared_x_sf = per_token_cast_to_fp8( + hidden_states, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True + ) + ref_shared_x_sf = _shared_sf_tma_layout(ref_shared_x_sf, shared_block_m) + shared_x_sf = torch.full_like(ref_shared_x_sf.t(), -1).t() + + prepare_megamoe_nvfp4_inputs( + hidden_states, + topk_weights, + topk_ids, + x, + x_sf, + x_scales, + topk_idx_out, + topk_weights_out, + is_padding=is_padding, + shared_x=shared_x, + shared_x_sf=shared_x_sf, + shared_block_m=shared_block_m if shared == "mxfp8" else None, + w2_global_scale=w2_global_scale, + l2_alpha=l2_alpha, + l2_activation_scale=0.37, + ) + torch.accelerator.synchronize() + + # Written even without local tokens: other ranks route to these experts. + assert torch.equal(l2_alpha, w2_global_scale * 0.37) + ref_x_scales = hidden_states.float().abs().amax(-1).clamp_min(1e-12) * (1 / 2688) + ref_x, ref_x_sf = per_token_cast_to_nvfp4( + hidden_states.float() / ref_x_scales[:, None] + ) + assert torch.equal(x_scales, ref_x_scales) + assert torch.equal(x, ref_x) + assert torch.equal(x_sf, ref_x_sf) + ref_topk_idx = topk_ids.long().masked_fill(is_padding[:, None], -1) + assert torch.equal(topk_idx_out, ref_topk_idx) + ref_weights = topk_weights.masked_fill(is_padding[:, None], 0.0) + assert torch.equal(topk_weights_out, ref_weights) + if shared == "bf16": + assert torch.equal(shared_x, hidden_states) + elif shared == "mxfp8": + assert shared_x is not None and shared_x_sf is not None + assert torch.equal(shared_x.view(torch.uint8), ref_shared_x.view(torch.uint8)) + populated = ref_shared_x_sf != 0 + assert torch.equal(shared_x_sf[populated], ref_shared_x_sf[populated]) + assert torch.all(shared_x_sf[~populated] == -1) + + +NUM_EXPERTS, TOP_K, HIDDEN, INTER, NUM_TOKENS = 8, 2, 512, 256, 48 + +ACTIVATIONS = { + "swiglu": MegaMoEActivation(), + "swiglu_clamp": MegaMoEActivation(clamp=0.5), + "swiglu_alpha_beta": MegaMoEActivation(alpha=1.702, beta=1.0), + "situ": MegaMoEActivation("situ", alpha=2.0, beta=3.0), +} + + +def _quantize(w: torch.Tensor, spec: MegaMoEWeightSpec): + """Return checkpoint ``(weight, scale, global_scale)`` and the dequantized + weight; NVFP4 takes ``global_scale`` as ``weight_scale_2``.""" + from vllm.third_party.deep_gemm.utils.math import ( + cast_back_from_fp4, + cast_back_from_fp8, + per_token_cast_to_fp4, + per_token_cast_to_fp8, + ) + from vllm.third_party.deep_gemm.utils.nvfp4 import ( + cast_back_from_nvfp4, + per_token_cast_to_nvfp4, + ) + + if spec.scale_group_k is None: + return w, None, None, w.float() + if spec.has_global_scale: + global_scale = w.float().abs().amax().item() / 2688 + q, sf = per_token_cast_to_nvfp4(w, global_scale) + deq = cast_back_from_nvfp4(q, sf, global_scale) + return q, sf.view(torch.float8_e4m3fn), torch.tensor(global_scale), deq + if spec.values_per_element == 2: + q, sf = per_token_cast_to_fp4(w, use_ue8m0=True, gran_k=32) + deq = cast_back_from_fp4(q, sf, gran_k=32) + else: + q, sf = per_token_cast_to_fp8(w, use_ue8m0=True, gran_k=32) + deq = cast_back_from_fp8(q, sf, gran_k=32) + return q, (sf.view(torch.int32) >> 23).to(torch.uint8), None, deq + + +def _round_trip(x: torch.Tensor, quant, nvfp4_scale: float | None = None): + """What the kernel computes on: x after ``quant`` (None keeps BF16).""" + from vllm.third_party.deep_gemm.utils.math import ( + cast_back_from_fp8, + per_token_cast_to_fp8, + ) + from vllm.third_party.deep_gemm.utils.nvfp4 import ( + cast_back_from_nvfp4, + per_token_cast_to_nvfp4, + ) + + if quant is None: + return x.bfloat16().float() + if quant == kNvfp4DynamicToken: + if nvfp4_scale is not None: + q = per_token_cast_to_nvfp4(x, nvfp4_scale) + return cast_back_from_nvfp4(*q, nvfp4_scale) + x = x.bfloat16().float() + scales = x.abs().amax(-1, keepdim=True).clamp_min(1e-12) / 2688 + return cast_back_from_nvfp4(*per_token_cast_to_nvfp4(x / scales)) * scales + return cast_back_from_fp8( + *per_token_cast_to_fp8(x.bfloat16(), use_ue8m0=True, gran_k=32), gran_k=32 + ) + + +def _activate(h: torch.Tensor, act: MegaMoEActivation) -> torch.Tensor: + gate, up = h.bfloat16().float().chunk(2, dim=-1) + if act.name == "situ": + gate = act.alpha * torch.tanh(gate / act.alpha) * torch.sigmoid(gate) + up = act.beta * torch.tanh(up / act.beta) if act.beta else up + return gate * up + if act.clamp is not None: + gate, up = gate.clamp(max=act.clamp), up.clamp(-act.clamp, act.clamp) + return gate * torch.sigmoid(act.alpha * gate) * (up + act.beta) + + +def _expert(x, w13, w2, weight, fmt, act, l2_scale): + """Reference routed expert. MX and BF16 weight the activation before the + L2 dispatch; NVFP4 quantizes the unweighted activation with the global + ``l2_scale`` and weights the BF16 output.""" + a = _activate(x @ w13.t(), act) + if fmt.hidden_quant == kNvfp4DynamicToken: + a = _round_trip(a, fmt.hidden_quant, l2_scale) + return (a @ w2.t()).bfloat16().float() * weight[:, None] + return _round_trip(a * weight[:, None], fmt.hidden_quant) @ w2.t() + + +class _ReferenceSharedExpert(torch.nn.Module): + def __init__(self, w13, w2, spec: MegaMoEWeightSpec, act): + super().__init__() + self.w13, self.w2, self.act = w13, w2, act + self.quant = kMxfp8Dynamic if spec.scale_group_k else None + + def forward(self, x): + a = _activate(_round_trip(x, self.quant) @ self.w13.t(), self.act) + return (_round_trip(a, self.quant) @ self.w2.t()).bfloat16() + + +def _calc_diff(x: torch.Tensor, y: torch.Tensor) -> float: + x, y = x.double(), y.double() + return (1 - 2 * (x * y).sum() / (x * x + y * y).sum()).item() + + +@pytest.mark.skipif( + not ( + current_platform.is_device_capability_family(100) and is_deep_gemm_supported() + ), + reason="DeepGEMM MegaMoE kernels require SM100 and DeepGEMM", +) +@pytest.mark.parametrize( + "fmt_name,shared,act_name", + [ + pytest.param( + name, + shared, + act_name, + id=f"{act_name}-{name}-{str(shared).removeprefix('torch.')}", + ) + for act_name, act in ACTIVATIONS.items() + for name, fmt in MEGA_MOE_WEIGHT_FORMATS.items() + if act.name in fmt.activations + for shared in (None, *(spec.dtype for spec in fmt.shared)) + ], +) +def test_deep_gemm_mega_moe_matches_reference(dist_init, fmt_name, shared, act_name): + """``shared`` is the fused shared-expert dtype; None runs it unfused.""" + fmt, act = MEGA_MOE_WEIGHT_FORMATS[fmt_name], ACTIVATIONS[act_name] + shared_spec = next( + (spec for spec in fmt.shared if spec.dtype == shared), fmt.shared[0] + ) + fuse_shared = shared is not None + torch.manual_seed(0) + device = torch.device("cuda") + vllm_config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=NUM_TOKENS), + compilation_config=SimpleNamespace(static_forward_context={}), + ) + + def randn(*shape, fan_in): + return (torch.randn(*shape, device=device) / fan_in**0.5).bfloat16() + + shared_ckpt = [ + _quantize(randn(*shape, fan_in=k), shared_spec) + for shape, k in (((INTER, HIDDEN), HIDDEN),) * 2 + (((HIDDEN, INTER), INTER),) + ] + shared_mlp = SimpleNamespace( + gate_up_proj=SimpleNamespace( + weight=torch.cat([shared_ckpt[0][0], shared_ckpt[1][0]]), + weight_scale_inv=None + if shared_ckpt[0][1] is None + else torch.cat([shared_ckpt[0][1], shared_ckpt[1][1]]), + ), + down_proj=SimpleNamespace( + weight=shared_ckpt[2][0], weight_scale_inv=shared_ckpt[2][1] + ), + ) + + with torch.device(device): + experts = DeepGemmMegaMoEExperts( + vllm_config, + num_experts=NUM_EXPERTS, + num_local_experts=NUM_EXPERTS, + experts_start_idx=0, + top_k=TOP_K, + hidden_size=HIDDEN, + intermediate_size=INTER, + activation=act, + weight_format=fmt, + shared_experts=shared_mlp if fuse_shared else None, + prefix=f"model.layers.0.{fmt_name}.{act_name}.{shared}", + ) + assert experts.has_fused_shared_experts == fuse_shared + assert experts.shared_spec == (shared_spec if fuse_shared else None) + + w13 = torch.empty(NUM_EXPERTS, 2 * INTER, HIDDEN, device=device) + w2 = torch.empty(NUM_EXPERTS, HIDDEN, INTER, device=device) + # NVFP4 L2 input scales; the kernel uses their max. + l2_scales = [2e-3 * (e + 1) for e in range(NUM_EXPERTS)] + for e in range(NUM_EXPERTS): + # w1 and w3 are quantized together: NVFP4 needs one global scale. + w13_ckpt = _quantize(randn(2 * INTER, HIDDEN, fan_in=HIDDEN), fmt.routed) + w2_ckpt = _quantize(randn(HIDDEN, INTER, fan_in=INTER), fmt.routed) + w13[e], w2[e] = w13_ckpt[3], w2_ckpt[3] + for shard_id, param, rows, (weight, scale, global_scale, _) in ( + ("w1", "w13", slice(0, INTER), w13_ckpt), + ("w3", "w13", slice(INTER, None), w13_ckpt), + ("w2", "w2", slice(None), w2_ckpt), + ): + tensors = { + "weight": weight[rows], + "weight_scale": None if scale is None else scale[rows], + "weight_scale_2": global_scale, + } + if shard_id == "w2" and global_scale is not None: + tensors["input_scale"] = torch.tensor(l2_scales[e]) + for suffix, tensor in tensors.items(): + if tensor is None: + continue + name = f"{param}_{suffix}" + assert experts.weight_loader( + getattr(experts, name), + tensor, + f"experts.{name}", + shard_id=shard_id, + expert_id=e, + return_success=True, + ) + + shared_w13 = torch.cat([shared_ckpt[0][3], shared_ckpt[1][3]]) + reference_shared = _ReferenceSharedExpert( + shared_w13, shared_ckpt[2][3], shared_spec, act + ) + if fuse_shared: + for (weight, scale, _, _), shard in zip(shared_ckpt, ("w1", "w3", "w2")): + experts.load_shared_expert_weight(f"shared_experts.{shard}.weight", weight) + if scale is not None: + experts.load_shared_expert_weight( + f"shared_experts.{shard}.weight_scale", scale + ) + + hidden_states = torch.randn(NUM_TOKENS, HIDDEN, device=device).bfloat16() + scores = torch.randn(NUM_TOKENS, NUM_EXPERTS, device=device) + topk_weights, topk_ids = torch.topk(scores.softmax(-1), TOP_K, dim=-1) + + output = experts( + hidden_states, + topk_weights, + topk_ids, + None if fuse_shared else reference_shared, + ) + + x = _round_trip(hidden_states.float(), fmt.hidden_quant) + l2_scale = experts.l2_activation_scale + expected = reference_shared(hidden_states).float() + for e in range(NUM_EXPERTS): + token, slot = (topk_ids == e).nonzero(as_tuple=True) + expected.index_add_( + 0, + token, + _expert( + x[token], w13[e], w2[e], topk_weights[token, slot], fmt, act, l2_scale + ), + ) + + assert output.dtype == torch.bfloat16 + assert _calc_diff(output, expected) < 1e-3 + assert F.cosine_similarity(output.float(), expected, dim=-1).min() > 0.99 diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 629759de40..900b281fbf 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -52,6 +52,7 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( ) from vllm.model_executor.layers.quantization.utils.fp8_utils import input_to_float8 from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + dequant_mxfp8_to_bf16, mxfp8_e4m3_quantize, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( @@ -65,7 +66,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( ) from vllm.model_executor.models.llama4 import Llama4MoE from vllm.platforms import current_platform -from vllm.utils.math_utils import next_power_of_2 +from vllm.utils.math_utils import next_power_of_2, round_up from vllm.utils.torch_utils import set_random_seed try: @@ -788,7 +789,7 @@ def test_trtllm_fp8_swiglu_clamp_support( def _make_mxfp8_moe_weights( - e: int, n: int, k: int + e: int, n: int, k: int, is_gated: bool = True ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: def quantize(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: q, s = zip( @@ -796,8 +797,9 @@ def _make_mxfp8_moe_weights( ) return torch.stack(q), torch.stack(s) + w1_rows = 2 * n if is_gated else n w1, w1_scale = quantize( - torch.randn(e, 2 * n, k, device="cuda", dtype=torch.bfloat16) / 10 + torch.randn(e, w1_rows, k, device="cuda", dtype=torch.bfloat16) / 10 ) w2, w2_scale = quantize( torch.randn(e, k, n, device="cuda", dtype=torch.bfloat16) / 10 @@ -806,16 +808,20 @@ def _make_mxfp8_moe_weights( @pytest.mark.parametrize("block_shape", [[128, 128], [1, 32]], ids=["block", "mxfp8"]) +@pytest.mark.parametrize("transposed", [True, False], ids=["transposed", "row-major"]) @pytest.mark.parametrize("m", [1, 16]) @pytest.mark.skipif( not current_platform.is_device_capability_family(100), reason="Requires TRTLLM-Gen FP8 MoE (SM100)", ) def test_trtllm_fp8_block_moe_deferred_finalize( - m: int, block_shape: list[int], workspace_init + m: int, block_shape: list[int], transposed: bool, workspace_init ): """TRTLLM-Gen block-FP8 and MXFP8 modular experts can leave the top-k - finalize to the caller.""" + finalize to the caller. For fp8-block, ``transposed`` toggles the oracle's + is_scale_transposed column-major activation-scale layout.""" + if block_shape != [128, 128] and transposed: + pytest.skip("column-major scales apply to fp8 [128, 128] block only") e, topk, n, k = 32, 4, 1024, 1024 set_random_seed(7) with set_current_vllm_config(vllm_config): @@ -849,6 +855,9 @@ def test_trtllm_fp8_block_moe_deferred_finalize( a2_scale=None, block_shape=block_shape, ) + is_block = block_shape == [128, 128] + if is_block and not transposed: + quant_config = replace(quant_config, is_scale_transposed=False) # One rank of a TP group, which deferral needs. moe_config = FusedMoEConfig( num_experts=e, @@ -866,19 +875,27 @@ def test_trtllm_fp8_block_moe_deferred_finalize( routing_method=RoutingMethodType.TopK, max_num_tokens=next_power_of_2(m), ) - kernel = mk.FusedMoEKernel( - maybe_make_prepare_finalize( - moe=moe_config, quant_config=quant_config, allow_new_interface=True - ), - TrtLlmFp8ExpertsModular(moe_config=moe_config, quant_config=quant_config), - ) + + def make_kernel(quant_config: FusedMoEQuantConfig) -> mk.FusedMoEKernel: + return mk.FusedMoEKernel( + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + ), + TrtLlmFp8ExpertsModular( + moe_config=moe_config, quant_config=quant_config + ), + ) + + kernel = make_kernel(quant_config) a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10 score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16) topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False) - check_deferred_moe_finalize( - moe_config, - lambda: kernel.apply( + + def apply(k_: mk.FusedMoEKernel) -> torch.Tensor: + return k_.apply( hidden_states=a, w1=w1, w2=w2, @@ -888,7 +905,34 @@ def test_trtllm_fp8_block_moe_deferred_finalize( global_num_experts=e, expert_map=None, apply_router_weight_on_input=False, - ), + ) + + if is_block and transposed: + # The oracle gates column-major scales on TRTLLM fp8-block. + assert quant_config.is_scale_transposed + assert not make_fp8_moe_quant_config( + Fp8MoeBackend.DEEPGEMM, + w1_scale, + w2_scale, + a1_scale=None, + a2_scale=None, + block_shape=block_shape, + ).is_scale_transposed + # no_dp_ep emits column-major [M, G] scales (physical [G, M]), so + # the experts' a1q_scale.t().contiguous() is a no-op. + _, a1q_scale, _, _, _ = kernel.prepare_finalize.prepare( + a, topk_weights, topk_ids, e, None, False, quant_config + ) + assert a1q_scale.t().is_contiguous() + # Bitwise-equal to running with row-major scales. + row_major = replace(quant_config, is_scale_transposed=False) + torch.testing.assert_close( + apply(kernel), apply(make_kernel(row_major)), atol=0, rtol=0 + ) + + check_deferred_moe_finalize( + moe_config, + lambda: apply(kernel), router_weights=topk_weights, ) @@ -986,3 +1030,123 @@ def test_trtllm_mxfp8_minimax2_routing_applies_routed_scale(m: int, workspace_in ) torch.testing.assert_close(fused, reference, atol=2e-2, rtol=2e-2) + + +@pytest.mark.parametrize( + ("activation", "n", "k"), + [ + (MoEActivation.SILU, 1024, 1024), + (MoEActivation.SILU, 928, 1024), + (MoEActivation.SILU, 1024, 2688), + (MoEActivation.RELU2_NO_MUL, 1024, 1024), + # Nemotron-3.5 Lightning's non-gated MoE: intermediate 1856, hidden 2688. + (MoEActivation.RELU2_NO_MUL, 1856, 2688), + ], +) +@pytest.mark.parametrize("m", [1, 32]) +@pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="Requires TRTLLM-Gen FP8 MoE (SM100)", +) +def test_trtllm_mxfp8_moe_pads_intermediate( + m: int, activation: MoEActivation, n: int, k: int, workspace_init +): + """TRTLLM-Gen MXFP8 experts need the intermediate size to be a multiple + of 128 and the hidden size a multiple of 256. Misaligned experts, gated or + not, are padded at load and, with the activations padded and the output + cropped as the MoE runner does, must still match the unquantized reference + on the dequantized weights.""" + e, topk = 8, 2 + is_gated = activation.is_gated + block_shape = [1, 32] + set_random_seed(7) + with set_current_vllm_config(vllm_config): + w1, w2, w1_scale, w2_scale = _make_mxfp8_moe_weights(e, n, k, is_gated) + w1_ref = dequant_mxfp8_to_bf16(w1, w1_scale) + w2_ref = dequant_mxfp8_to_bf16(w2, w2_scale) + + layer = SimpleNamespace( + weight_block_size=block_shape, + moe_config=SimpleNamespace( + is_act_and_mul=is_gated, intermediate_size_per_partition=n + ), + activation=activation, + ) + w1, w2, w1_scale, w2_scale = convert_to_fp8_moe_kernel_format( + Fp8MoeBackend.FLASHINFER_TRTLLM, + layer, + w1, + w2, + w1_scale, + w2_scale, + w13_input_scale=None, + w2_input_scale=None, + ) + padded_n = layer.moe_config.intermediate_size_per_partition + padded_k = layer.moe_config.hidden_dim + assert padded_n == round_up(n, 128) + assert padded_k == round_up(k, 256) + + quant_config = make_fp8_moe_quant_config( + Fp8MoeBackend.FLASHINFER_TRTLLM, + w1_scale, + w2_scale, + a1_scale=None, + a2_scale=None, + block_shape=block_shape, + ) + moe_config = FusedMoEConfig( + num_experts=e, + experts_per_token=topk, + hidden_dim=padded_k, + intermediate_size=padded_n, + num_local_experts=e, + num_logical_experts=e, + activation=activation, + device="cuda", + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + in_dtype=torch.bfloat16, + routing_method=RoutingMethodType.Renormalize, + max_num_tokens=next_power_of_2(m), + ) + # The monolithic experts pass the activation type to the kernel; the + # modular ones always run it gated. + kernel = mk.FusedMoEKernel( + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=True, + ), + TrtLlmFp8ExpertsMonolithic( + moe_config=moe_config, quant_config=quant_config + ), + ) + + a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10 + logits = torch.randn((m, e), device="cuda", dtype=torch.float32) + output = kernel.apply_monolithic( + hidden_states=torch.nn.functional.pad(a, (0, padded_k - k)), + w1=w1, + w2=w2, + router_logits=logits, + activation=activation, + global_num_experts=e, + expert_map=None, + apply_router_weight_on_input=False, + routed_scaling_factor=1.0, + routing_sink=None, + )[:, :k] + topk_weights, topk_ids, _ = fused_topk(a, logits, topk, renormalize=True) + reference = fused_experts( + a, + w1_ref, + w2_ref, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=activation, + global_num_experts=e, + expert_map=None, + ) + + check_accuracy(ref_output=reference, actual_output=output) diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index 0a8f5c29d0..968781a188 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -42,13 +42,15 @@ ROCM_AVAILABLE = current_platform.is_rocm() ROCM_TRITON_KERNELS_AVAILABLE = False ROCM_AITER_AVAILABLE = is_aiter_found() ROCM_GFX950 = False +ROCM_GFX1250 = False if ROCM_AVAILABLE: - from vllm.platforms.rocm import on_gfx950 + from vllm.platforms.rocm import on_gfx950, on_gfx1250 from vllm.utils.import_utils import has_triton_kernels ROCM_TRITON_KERNELS_AVAILABLE = has_triton_kernels() ROCM_GFX950 = on_gfx950() + ROCM_GFX1250 = on_gfx1250() if ROCM_AITER_AVAILABLE: from aiter.ops.triton.quant import dynamic_mxfp4_quant @@ -1568,6 +1570,15 @@ ROCM_BACKEND_CONFIGS = { "requires_aiter": True, "requires_gfx950": True, }, + "AITER_TRITON_MXFP4_MXFP4": { + "activation": "SILU", + "act_type": "mxfp4", + "rtol": 1.0, + "percent": 0.8, + "requires_aiter": True, + "requires_gfx950": False, + "requires_gfx1250": True, + }, } @@ -1607,6 +1618,8 @@ def test_rocm_mxfp4_moe_oracle( pytest.skip(f"Backend {backend_name} requires AITER") if config["requires_gfx950"] and not ROCM_GFX950: pytest.skip(f"Backend {backend_name} requires GFX950") + if config.get("requires_gfx1250") and not ROCM_GFX1250: + pytest.skip(f"Backend {backend_name} requires GFX1250") import vllm.distributed.parallel_state as ps from vllm.config import VllmConfig, set_current_vllm_config diff --git a/tests/kernels/moe/test_rocm_aiter_topk.py b/tests/kernels/moe/test_rocm_aiter_topk.py index 11f9cf28ae..827ebe7c04 100644 --- a/tests/kernels/moe/test_rocm_aiter_topk.py +++ b/tests/kernels/moe/test_rocm_aiter_topk.py @@ -219,3 +219,169 @@ def test_rocm_aiter_grouped_topk_torch_compile_compatibility(): topk_weights_original, topk_weights_compiled, rtol=1e-2, atol=1e-2 ) assert torch.allclose(topk_ids_original, topk_ids_compiled) + + +def test_rocm_aiter_topk_gating_custom_op_registration(): + assert callable(torch.ops.vllm.rocm_aiter_topk_gating) + + +def _sort_routing(weights: torch.Tensor, ids: torch.Tensor): + ids_sorted, order = torch.sort(ids, dim=-1) + weights_sorted = torch.gather(weights, 1, order) + return weights_sorted, ids_sorted + + +@pytest.mark.parametrize("num_tokens", [1, 4, 16, 64]) +@pytest.mark.parametrize("num_experts", [256, 512]) +@pytest.mark.parametrize("topk", [8, 10]) +@pytest.mark.parametrize("renormalize", [True, False]) +def test_rocm_aiter_topk_gating_matches_softmax_reference( + num_tokens: int, + num_experts: int, + topk: int, + renormalize: bool, +): + """Numerical check vs softmax + topk, including Qwen3.8 (E=512, k=10).""" + torch.manual_seed(0) + gating_output = torch.randn( + (num_tokens, num_experts), dtype=torch.bfloat16, device="cuda" + ) + scores = torch.softmax(gating_output.float(), dim=-1) + cutoff = torch.topk(scores, k=topk, dim=-1).values[:, -1:] + + topk_weights = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda") + topk_ids = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda") + token_expert_indices = torch.empty( + (num_tokens, topk), dtype=torch.int32, device="cuda" + ) + torch.ops.vllm.rocm_aiter_topk_gating( + topk_weights, + topk_ids, + token_expert_indices, + gating_output, + renormalize, + 0, + "", + ) + + # BF16 logits can tie at the top-k boundary, so a different tied expert is + # valid. Check selection against the cutoff and weights against the scores + # of the experts the kernel actually selected. + assert ((topk_ids >= 0) & (topk_ids < num_experts)).all() + sorted_ids = torch.sort(topk_ids, dim=-1).values + assert (sorted_ids[:, 1:] != sorted_ids[:, :-1]).all() + selected_scores = torch.gather(scores, 1, topk_ids.long()) + assert (selected_scores >= cutoff).all() + expected_weights = selected_scores + if renormalize: + expected_weights = expected_weights / expected_weights.sum(dim=-1, keepdim=True) + torch.testing.assert_close(topk_weights, expected_weights, atol=2e-2, rtol=2e-2) + + +def test_rocm_aiter_topk_gating_shared_expert_suffix(): + """Fused shared expert: softmax top-k on the prefix, sigmoid appended. + + Matches the Qwen3.8 FSE launch: 512 routed experts, top-10, one shared + column, weights wide enough for the suffix, and the shared id prefilled. + """ + torch.manual_seed(0) + num_tokens, num_experts, topk, num_shared = 4, 512, 10, 1 + routed = torch.randn((num_tokens, num_experts), dtype=torch.bfloat16, device="cuda") + shared = torch.randn((num_tokens, num_shared), dtype=torch.bfloat16, device="cuda") + shared[0, 0] = 40 + gating_output = torch.cat([routed, shared], dim=-1).contiguous() + width = topk + num_shared + shared_ids = torch.arange( + num_experts, num_experts + num_shared, device="cuda", dtype=torch.int32 + ) + topk_ids = torch.full((num_tokens, width), -1, dtype=torch.int32, device="cuda") + topk_ids[:, topk:] = shared_ids + topk_weights = torch.empty((num_tokens, width), dtype=torch.float32, device="cuda") + token_expert_indices = torch.empty( + (num_tokens, topk), dtype=torch.int32, device="cuda" + ) + + torch.ops.vllm.rocm_aiter_topk_gating( + topk_weights, + topk_ids[:, :topk], + token_expert_indices, + gating_output, + True, + num_shared, + "sigmoid", + ) + + scores = torch.softmax(routed.float(), dim=-1) + cutoff = torch.topk(scores, k=topk, dim=-1).values[:, -1:] + routed_ids = topk_ids[:, :topk] + assert ((routed_ids >= 0) & (routed_ids < num_experts)).all() + sorted_ids = torch.sort(routed_ids, dim=-1).values + assert (sorted_ids[:, 1:] != sorted_ids[:, :-1]).all() + selected = torch.gather(scores, 1, routed_ids.long()) + assert (selected >= cutoff).all() + expected = selected / selected.sum(dim=-1, keepdim=True) + torch.testing.assert_close(topk_weights[:, :topk], expected, atol=2e-2, rtol=2e-2) + torch.testing.assert_close( + topk_weights[:, topk:], + torch.sigmoid(shared.float()), + atol=2e-2, + rtol=2e-2, + ) + assert torch.equal(topk_ids[:, topk:], shared_ids.expand_as(topk_ids[:, topk:])) + + +def test_rocm_aiter_topk_gating_torch_compile_compatibility(): + token, expert, topk, renormalize = 4, 512, 10, True + gating_output = torch.randn((token, expert), dtype=torch.bfloat16, device="cuda") + device = gating_output.device + + def gating_fn(gating_output, topk_weights, topk_ids, token_expert_indices): + return torch.ops.vllm.rocm_aiter_topk_gating( + topk_weights, + topk_ids, + token_expert_indices, + gating_output, + renormalize, + 0, + "", + ) + + torch.library.opcheck( + torch.ops.vllm.rocm_aiter_topk_gating, + ( + torch.empty((token, topk), dtype=torch.float32, device=device), + torch.empty((token, topk), dtype=torch.int32, device=device), + torch.empty((token, topk), dtype=torch.int32, device=device), + gating_output, + ), + kwargs={ + "renormalize": renormalize, + "num_shared_experts": 0, + "shared_expert_scoring_func": "", + }, + test_utils=("test_faketensor"), + ) + + compiled_fn = torch.compile( + gating_fn, + fullgraph=True, + backend="inductor", + mode="reduce-overhead", + dynamic=False, + ) + + def alloc(): + return ( + torch.empty((token, topk), dtype=torch.float32, device=device), + torch.empty((token, topk), dtype=torch.int32, device=device), + torch.empty((token, topk), dtype=torch.int32, device=device), + ) + + w0, i0, t0 = alloc() + w1, i1, t1 = alloc() + gating_fn(gating_output, w0, i0, t0) + compiled_fn(gating_output, w1, i1, t1) + w0, i0 = _sort_routing(w0, i0) + w1, i1 = _sort_routing(w1, i1) + assert torch.allclose(w0, w1, rtol=1e-2, atol=1e-2) + assert torch.equal(i0, i1) diff --git a/tests/kernels/moe/test_topk_softplus_sqrt.py b/tests/kernels/moe/test_topk_softplus_sqrt.py index 0d92fe6eac..c9abdd36b1 100644 --- a/tests/kernels/moe/test_topk_softplus_sqrt.py +++ b/tests/kernels/moe/test_topk_softplus_sqrt.py @@ -81,32 +81,11 @@ def _torch_topk_softplus_sqrt( return topk_weights.to(torch.float32), topk_ids.to(torch.int32) -@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm backend selection") -@pytest.mark.parametrize("has_hash_routing", [False, True]) -@pytest.mark.parametrize( - "backend,expected_experts", - [ - ("aiter_triton_mxfp4_bf16", "AiterW4A16ExpertsMonolithic"), - ("aiter", "AiterMxfp4Experts"), - ("triton_unfused", "UnfusedOAITritonExperts"), - ], -) -def test_hash_routing_backend_selection( - dist_init, - default_vllm_config, - monkeypatch, - has_hash_routing, - backend, - expected_experts, -): - """Hash tables must reach backend selection; modular AITER remains eligible.""" +def _dsv4_moe_config(monkeypatch, default_vllm_config, has_hash_routing, backend): from dataclasses import replace from vllm._aiter_ops import rocm_aiter_ops from vllm.model_executor.layers.fused_moe.layer import FusedMoEFactory - from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( - select_deepseek_v4_mxfp4_moe_backend, - ) from vllm.platforms import rocm # Exercise both native MX and monolithic AITER selection on any ROCm device. @@ -130,14 +109,63 @@ def test_hash_routing_backend_selection( ), prefix="hash_routing_selection", ) - config = replace(layer.routed_experts.moe_config, moe_backend=backend) + return layer, replace(layer.routed_experts.moe_config, moe_backend=backend) - if has_hash_routing and backend == "aiter_triton_mxfp4_bf16": - with pytest.raises(ValueError, match="hash routing"): - select_deepseek_v4_mxfp4_moe_backend(config) - else: - _, experts_cls = select_deepseek_v4_mxfp4_moe_backend(config) - assert experts_cls.__name__ == expected_experts + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm backend selection") +@pytest.mark.parametrize("has_hash_routing", [False, True]) +@pytest.mark.parametrize( + "backend,expected_experts", + [ + ("aiter_triton_mxfp4_bf16", "AiterW4A16ExpertsMonolithic"), + ("aiter", "AiterMxfp4Experts"), + ("triton_unfused", "UnfusedOAITritonExperts"), + ], +) +def test_hash_routing_backend_selection( + dist_init, + default_vllm_config, + monkeypatch, + has_hash_routing, + backend, + expected_experts, +): + """Hash tables reach backend selection and hash-capable kernels.""" + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + select_deepseek_v4_mxfp4_moe_backend, + ) + + layer, config = _dsv4_moe_config( + monkeypatch, default_vllm_config, has_hash_routing, backend + ) + _, experts_cls = select_deepseek_v4_mxfp4_moe_backend(config) + assert experts_cls.__name__ == expected_experts + # Monolithic kernels route internally, so they read the table from the layer. + assert (layer.routed_experts.hash_indices_table is not None) == has_hash_routing + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm backend selection") +def test_hash_routing_rejects_monolithic_without_support( + dist_init, default_vllm_config, monkeypatch +): + """A monolithic kernel that does not implement hash routing is rejected.""" + from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a16_moe import ( + AiterW4A16ExpertsMonolithic, + ) + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + select_deepseek_v4_mxfp4_moe_backend, + ) + + monkeypatch.setattr( + AiterW4A16ExpertsMonolithic, + "_supports_hash_routing", + staticmethod(lambda: False), + ) + _, config = _dsv4_moe_config( + monkeypatch, default_vllm_config, True, "aiter_triton_mxfp4_bf16" + ) + with pytest.raises(ValueError, match="hash routing"): + select_deepseek_v4_mxfp4_moe_backend(config) def test_torch_topk_softplus_sqrt_breaks_ties_by_expert_id(): diff --git a/tests/kernels/quantization/test_rocm_mxfp4.py b/tests/kernels/quantization/test_rocm_mxfp4.py index fa39c641ba..f83445b39c 100644 --- a/tests/kernels/quantization/test_rocm_mxfp4.py +++ b/tests/kernels/quantization/test_rocm_mxfp4.py @@ -50,7 +50,7 @@ SKINNY_GEMM_PASS_RATES = { } PRESHUFFLED_SHAPES = [ (64, 4096, 8192), - (32, 8192, 16384), + (32, 16384, 16384), ] @@ -421,8 +421,12 @@ def test_aiter_fp4_gemm_preshuffled_tuned_shapes(shape): M, K, N = shape assert M <= 64 - # The predicate takes K as packed bytes, matching `weight.shape[1]`. + # The predicate takes K as packed bytes, matching `weight.shape[1]`, and + # must agree with the config the kernel itself resolves. assert rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K // 2) + from aiter.ops.triton.utils.gemm_config_utils import get_gemm_config + + assert get_gemm_config("GEMM-AFP4WFP4_PRESHUFFLED", M, N, K)[1] A = torch.randn(M, K, dtype=torch.bfloat16) B = torch.randn(N, K, dtype=torch.bfloat16) @@ -464,6 +468,18 @@ def test_aiter_fp4_gemm_preshuffled_tuned_shapes(shape): _assert_deterministic(run_preshuffled, n_runs=3) +@pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only") +def test_aiter_fp4_preshuffled_tuned_gate_uses_logical_k(): + """AITER ships N=16384-K=16384 but not N=16384-K=8192, so a K=8192 layer + has no tuned preshuffled config even though twice its K does.""" + _assert_aiter_supported() + from vllm._aiter_ops import rocm_aiter_ops + + rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned.cache_clear() + assert rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(16384, 16384 // 2) + assert not rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(16384, 8192 // 2) + + @pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only") def test_aiter_fp4_gemm_a4w4_determinism(): """AITER FP4 A4W4 GEMM is bitwise deterministic across repeated runs. diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index 237c552fc1..3774d14b0c 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -14,6 +14,7 @@ These tests cover: import math from types import SimpleNamespace +import numpy as np import pytest import torch @@ -157,6 +158,104 @@ def test_dspark_context_kv_matches_query_insert( torch.testing.assert_close(kv, before, rtol=0, atol=0) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +@pytest.mark.parametrize("kv_mxfp8", [False, True]) +def test_v41_pcp_swa_cache_matches_non_pcp_writer(kv_mxfp8, monkeypatch): + """PCP SWA writes must preserve the established CUDA cache bytes.""" + from vllm.models.deepseek_v41 import attention as attention_module + from vllm.models.deepseek_v41.attention import DeepseekV4Attention + + torch.manual_seed(43) + block_size = 64 + num_local_tokens = 5 + num_tokens = 2 * num_local_tokens + row_size = 528 if kv_mxfp8 else 584 + num_blocks = 2 + page_stride = math.ceil(block_size * row_size / 576) * 576 + 576 + backing = torch.full( + (num_blocks, page_stride), 165, dtype=torch.uint8, device="cuda" + ) + cache = backing.as_strided( + (num_blocks, block_size, row_size), (page_stride, row_size, 1) + ) + reference_backing = backing.clone() + reference = reference_backing.as_strided(cache.shape, cache.stride()) + + local_kv = torch.randn(num_local_tokens, 512, dtype=torch.bfloat16, device="cuda") + remote_kv = torch.randn_like(local_kv) + gathered_kv = torch.cat((local_kv, remote_kv)) + local_positions = torch.arange(num_local_tokens, device="cuda") + 3 + cache_positions = torch.cat((local_positions, local_positions + 11)) + cache_slots = torch.randperm( + num_blocks * block_size, dtype=torch.int64, device="cuda" + )[:num_tokens] + cache_slots[2] = -1 + angles = torch.randn(32, 32, device="cuda") + cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1) + num_heads = 8 + q = torch.randn( + num_local_tokens, + num_heads, + 512, + dtype=torch.bfloat16, + device="cuda", + ) + + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + torch.zeros(num_tokens, num_heads, 512, dtype=torch.bfloat16, device="cuda"), + gathered_kv, + reference.view(num_blocks, -1), + cache_slots, + cache_positions, + cos_sin, + num_heads, + 1e-20, + block_size, + False, + kv_mxfp8, + ) + + monkeypatch.setattr( + attention_module, + "get_pcp_group", + lambda: SimpleNamespace( + all_gather=lambda tensor, dim: torch.cat((tensor, remote_kv), dim=dim) + ), + ) + metadata = SimpleNamespace( + slot_mapping=torch.full( + (num_local_tokens,), -1, dtype=torch.int64, device="cuda" + ), + cache_slot_mapping=cache_slots, + cache_positions=cache_positions, + block_size=block_size, + ) + attention = SimpleNamespace( + use_pcp=True, + swa_cache_layer=SimpleNamespace( + prefix="swa", kv_cache=cache, block_size=block_size + ), + rotary_emb=SimpleNamespace(cos_sin_cache=cos_sin), + _flashinfer_fp8_kv_scale=None, + kv_mxfp8=kv_mxfp8, + accepts_unnormed_unroped_query=False, + n_local_heads=num_heads, + padded_heads=num_heads, + eps=1e-20, + head_dim=512, + ) + DeepseekV4Attention._fused_qnorm_rope_kv_insert( + attention, q, local_kv, local_positions, {"swa": metadata} + ) + + torch.testing.assert_close( + backing.view(torch.uint8), + reference_backing.view(torch.uint8), + rtol=0, + atol=0, + ) + + @pytest.mark.skipif( not current_platform.is_cuda_alike(), reason="graph capture coverage" ) @@ -360,7 +459,7 @@ def _mxfp8_record_reference( @pytest.mark.parametrize("compress_ratio", [1, 2]) def test_v41_rope_insert_mxfp8_record(compress_ratio: int): """The 528-byte V4.1 record quantizes the RoPE dims instead of keeping them - in bf16, so a page is [block_size x 512 fp8][block_size x 16 UE8M0]. + in bf16, so a token is [512 fp8][16 UE8M0]. Byte-exact against the reference quantizer, including the rows and page padding the kernel must leave alone. @@ -398,9 +497,8 @@ def test_v41_rope_insert_mxfp8_record(compress_ratio: int): values, scales = _mxfp8_record_reference( latent[t], cos_sin[pos // compress_ratio * compress_ratio] ) - expected[page, row * 512 : (row + 1) * 512] = values - scale_offset = cache_block * 512 + row * 16 - expected[page, scale_offset : scale_offset + 16] = scales + expected[page, row * 528 : row * 528 + 512] = values + expected[page, row * 528 + 512 : (row + 1) * 528] = scales torch.testing.assert_close(cache_backing, expected, rtol=0, atol=0) @@ -510,14 +608,13 @@ def test_v41_mxfp8_cache_round_trip(): values, scales = _ue8m0_reference(k[t].float(), 32, 448.0) page, row = divmod(t, block_size) torch.testing.assert_close( - cache[page, row * 512 : (row + 1) * 512], + cache[page, row * 528 : row * 528 + 512], values.view(torch.uint8), rtol=0, atol=0, ) - scale_offset = block_size * 512 + row * 16 torch.testing.assert_close( - cache[page, scale_offset : scale_offset + 16], + cache[page, row * 528 + 512 : (row + 1) * 528], (scales.log2() + 127).to(torch.uint8), rtol=0, atol=0, @@ -564,7 +661,8 @@ def _decode_nvfp4_row(packed: torch.Tensor, scales: torch.Tensor) -> torch.Tenso @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") @pytest.mark.parametrize("compress_ratio", [1, 2]) def test_v41_rope_insert_nvfp4_record(compress_ratio: int): - """The 288-byte V4.1 NVFP4 record: 256 B of e2m1 pairs then 32 e4m3 scales. + """The 288-byte V4.1 NVFP4 record: 256 B of e2m1 pairs then 32 e4m3 scales + per token. Scales are byte-exact against the reference (``amax / 6`` clamped to the e4m3 range); values are checked as round-to-nearest on the e2m1 grid, whose @@ -604,17 +702,13 @@ def test_v41_rope_insert_nvfp4_record(compress_ratio: int): amax = rotated.view(32, 16).abs().amax(-1) scale = (amax / 6.0).clamp(2.0**-9, 448.0).to(torch.float8_e4m3fn) - got_scales = cache_backing[page, cache_block * 256 + row * 32 :][:32] + record = cache_backing[page, row * 288 : (row + 1) * 288] + got_scales = record[256:] torch.testing.assert_close(got_scales, scale.view(torch.uint8), rtol=0, atol=0) - decoded = _decode_nvfp4_row( - cache_backing[page, row * 256 : (row + 1) * 256], got_scales - ) + decoded = _decode_nvfp4_row(record[:256], got_scales) step = scale.float().repeat_interleave(16) assert ((decoded - rotated).abs() <= step).all() - untouched[page, row * 256 : (row + 1) * 256] = cache_backing[ - page, row * 256 : (row + 1) * 256 - ] - untouched[page, cache_block * 256 + row * 32 :][:32] = got_scales + untouched[page, row * 288 : (row + 1) * 288] = record # Rows the kernel must not have touched, page padding included. torch.testing.assert_close(cache_backing, untouched, rtol=0, atol=0) @@ -670,10 +764,8 @@ def test_v41_nvfp4_gather_matches_insert(num_reqs, num_tokens, partial): offset=offset, ) - expected = _decode_nvfp4_row( - backing[:, : block_size * 256].reshape(-1, 256), - backing[:, block_size * 256 : block_size * 288].reshape(-1, 32), - ) + records = backing[:, : block_size * 288].reshape(-1, 288) + expected = _decode_nvfp4_row(records[:, :256], records[:, 256:]) def check(shorter_by): assert (storage[:, num_tokens + offset :] == -123).all() @@ -702,6 +794,72 @@ def test_v41_nvfp4_gather_matches_insert(num_reqs, num_tokens, partial): check(1) +@pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="The V4.1 records are SM100-only", +) +@pytest.mark.parametrize("record", ["swa_fp8", "compressed_fp8", "compressed_fp4"]) +def test_v41_cache_writers_match_flashmla_decode(record: str): + """FlashMLA and vLLM's gathers read back the token a V4.1 writer put mid + page in a padded block, so all agree on the interleaved record layout.""" + import vllm.v1.attention.ops.flashmla as fm + from vllm.models.deepseek_v41.common.ops.cache_utils import ( + dequantize_and_gather_k_cache, + dequantize_and_gather_k_cache_triton, + ) + from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import ( + rope_quant_insert, + ) + + device, block_size, topk = "cuda", 64, 64 + slot = block_size + 37 + bytes_per_token = 288 if record == "compressed_fp4" else 528 + torch.manual_seed(0) + kv = torch.randn(1, 512, dtype=torch.bfloat16, device=device) + positions = torch.tensor([slot], device=device) + # cos = 1, sin = 0: RoPE is the identity, so the cache should hold `kv`. + cos_sin = torch.zeros(slot + 1, 64, device=device) + cos_sin[:, :32] = 1 + # Padded like a packed block: a multiple of the record and of 512 B. + cache = torch.zeros(2, block_size + 32, bytes_per_token, dtype=torch.uint8) + cache = cache.to(device)[:, :block_size] + if record == "swa_fp8": + torch.ops._C.fused_deepseek_v4_kv_rope_insert( + kv, cache, positions, positions, cos_sin, block_size, None, True + ) + else: + rope_quant_insert(kv, positions, cos_sin, cache, positions, 1) + + indices = torch.full((1, 1, topk), -1, dtype=torch.int32, device=device) + indices[..., 0] = slot + extra = None + if record == "compressed_fp4": + cache, extra = torch.zeros(2, block_size, 528).to(cache), cache + out, _ = fm.flash_mla_with_kvcache( + q=torch.randn(1, 1, 64, 512, dtype=torch.bfloat16, device=device), + k_cache=cache.unsqueeze(2), + block_table=None, + cache_seqlens=None, + head_dim_v=512, + tile_scheduler_metadata=fm.get_mla_metadata()[0], + is_fp8_kvcache=True, + indices=indices if extra is None else torch.full_like(indices, -1), + extra_k_cache=None if extra is None else extra.unsqueeze(2), + extra_indices_in_kvcache=None if extra is None else indices, + ) + for gather in (dequantize_and_gather_k_cache, dequantize_and_gather_k_cache_triton): + rows = torch.zeros(1, slot + 1, 512, dtype=torch.bfloat16, device=device) + block_table = torch.arange(2, dtype=torch.int32, device=device)[None] + seq_lens = torch.tensor([slot + 1], dtype=torch.int32, device=device) + k_cache = cache if extra is None else extra + gather(rows, k_cache, seq_lens, None, block_table, block_size, offset=0) + torch.testing.assert_close( + out[0, 0], rows[0, slot].expand(64, -1), rtol=0, atol=0 + ) + cos = torch.nn.functional.cosine_similarity(rows[0, slot].float(), kv[0].float(), 0) + assert cos > (0.98 if extra is not None else 0.999) + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") @pytest.mark.parametrize("compress_ratio", [1, 2]) @pytest.mark.parametrize("store_fp8", [False, True]) @@ -796,6 +954,7 @@ def test_v41_compressor_metadata_maps_tokens_to_their_ring(): capacity = 8 vllm_config = MagicMock() vllm_config.scheduler_config.max_num_batched_tokens = 16 + vllm_config.parallel_config.prefill_context_parallel_size = 1 spec = CircularBufferSpec( block_size=capacity, num_kv_heads=1, @@ -832,6 +991,67 @@ def test_v41_compressor_metadata_maps_tokens_to_their_ring(): assert metadata.token_to_req_indices.tolist() == [0, 0, 0, 1, 1, 2, 2] +def test_v41_pcp_compressor_latent_is_gathered_for_cache_writes(monkeypatch): + """PCP compresses locally and publishes the gathered latent to both caches.""" + from vllm.forward_context import ForwardContext, override_forward_context + from vllm.models.deepseek_v41 import attention as attention_module + from vllm.models.deepseek_v41.attention import DeepseekV4Attention + + local_latent = torch.arange(8, dtype=torch.float32).view(2, 4) + cache_positions = torch.tensor([0, 1, 10, 11]) + observed = {} + + class Compressor: + def __call__(self, kv_score, positions): + return local_latent + + def insert_cache(self, latent, positions, rotary_emb): + observed["insert"] = (latent, positions) + + def prepare_indexer(qr, latent, weights, positions, rotary, qr_scale, **kwargs): + observed["indexer"] = (latent, kwargs) + return None, None, None + + monkeypatch.setattr( + attention_module, + "get_pcp_group", + lambda: SimpleNamespace( + all_gather=lambda tensor, dim: torch.cat((tensor, tensor + 10), dim=dim) + ), + ) + metadata = {"swa": SimpleNamespace(cache_positions=cache_positions)} + attention = SimpleNamespace( + use_pcp=True, + compressor=Compressor(), + indexer=prepare_indexer, + aux_stream_list=None, + ln_events=[None, None], + swa_cache_layer=SimpleNamespace(prefix="swa"), + rotary_emb=SimpleNamespace(), + indexer_rotary_emb=SimpleNamespace(), + n_local_heads=1, + head_dim=4, + _wq_b_proj=lambda qr, scale: qr, + _fused_qnorm_rope_kv_insert=lambda q, kv, pos, meta: q, + _sparse_indexer_and_attn=lambda *args: None, + ) + rows = torch.zeros(2, 4) + positions = torch.arange(2) + with override_forward_context(ForwardContext({}, metadata, {})): + DeepseekV4Attention._prepare_and_attn( + attention, rows, rows, rows, None, rows, rows, positions, rows + ) + + expected_global = torch.cat((local_latent, local_latent + 10)) + insert_latent, insert_positions = observed["insert"] + torch.testing.assert_close(insert_latent, expected_global) + torch.testing.assert_close(insert_positions, cache_positions) + indexer_latent, indexer_kwargs = observed["indexer"] + torch.testing.assert_close(indexer_latent, local_latent) + torch.testing.assert_close(indexer_kwargs["cache_latent"], expected_global) + torch.testing.assert_close(indexer_kwargs["cache_positions"], cache_positions) + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA stream coverage") @pytest.mark.parametrize( "use_aux,use_graph", [(False, False), (True, False), (True, True)] @@ -871,6 +1091,7 @@ def test_v41_attention_joins_cache_writes_before_consumption(use_aux, use_graph) context = ForwardContext({}, metadata, {}) compressor = DeepseekCompressor.__new__(DeepseekCompressor) torch.nn.Module.__init__(compressor) + compressor.pcp_size = 1 compressor.head_dim, compressor.rope_head_dim, compressor.compress_ratio = ( 512, 64, @@ -896,7 +1117,7 @@ def test_v41_attention_joins_cache_writes_before_consumption(use_aux, use_graph) use_fp4_kv=False, ) - def prepare_indexer(qr, latent, weights, positions, rotary, qr_scale): + def prepare_indexer(qr, latent, weights, positions, rotary, qr_scale, **kwargs): DeepseekV4Indexer._produce_k(indexer, latent, positions, rotary) return None, None, None @@ -906,6 +1127,7 @@ def test_v41_attention_joins_cache_writes_before_consumption(use_aux, use_graph) auxiliary = [torch.cuda.Stream()] if use_aux else None attention = SimpleNamespace( + use_pcp=False, compressor=compressor, indexer=prepare_indexer, aux_stream_list=auxiliary, @@ -963,6 +1185,291 @@ def test_v41_attention_joins_cache_writes_before_consumption(use_aux, use_graph) previous = expected +@pytest.mark.parametrize("pcp_size", [2, 4]) +@pytest.mark.parametrize( + "lengths,prefilling,starts", + [ + ([9, 1], [True, False], [1, 4]), + ([1, 1], [False, False], [16, 3]), + ([3, 1], [True, False], [1, 0]), + ], +) +def test_v41_pcp_ring_links_match_request_positions( + pcp_size, lengths, prefilling, starts, monkeypatch +): + """Cross-chunk predecessors and open ring writes must belong to the request.""" + from vllm.models.deepseek_v41 import compressor as compressor_module + from vllm.models.deepseek_v41.compressor import ( + CompressorMetadata, + CompressorMetadataBuilder, + ) + from vllm.v1.worker.gpu import pcp_manager as pcp_manager_module + from vllm.v1.worker.gpu.pcp_manager import PCPManager + + monkeypatch.setattr( + pcp_manager_module, + "async_tensor_h2d", + lambda array, **kwargs: torch.from_numpy(array), + ) + manager = PCPManager(pcp_size, 0, torch.device("cpu"), shard_decode_requests=True) + query_start = np.array([0, *np.cumsum(lengths)], dtype=np.int32) + segments, counts = manager._build_batch_layout( + np.array(lengths), np.array(starts), np.array(prefilling), query_start + ) + width = max(counts) + capacity = 8 + blocks = [1, (1 << 20) + 2] + num_rows = min(width, 2 * len(lengths)) + metadata_by_rank = [] + positions_by_rank = [] + gathered_keys = [] + for local_segments in segments: + positions = torch.zeros(width, dtype=torch.int64) + slots = torch.full((width,), -1, dtype=torch.int64) + local_query_start = [0] + keys = [-1] * num_rows + for row, segment in enumerate(local_segments): + req = segment.global_batch_req_idx + begin = starts[req] + segment.global_batch_slice.start - query_start[req] + local = segment.rank_local_batch_slice + positions[local] = torch.arange(begin, begin + segment.num_tokens) + slots[local] = blocks[req] * capacity + positions[local] % capacity + local_query_start.append(local.stop) + keys[row] = (blocks[req] << 32) | int(positions[local.stop - 1]) + if not local_segments: + local_query_start.append(0) + metadata_by_rank.append( + CompressorMetadata( + slot_mapping=slots, + query_start_loc=torch.tensor(local_query_start, dtype=torch.int32), + token_to_req_indices=torch.empty(0, dtype=torch.int32), + ) + ) + positions_by_rank.append(positions) + gathered_keys.extend(keys) + + keys = torch.tensor(gathered_keys, dtype=torch.int64) + contributed_keys = [] + + def all_gather(tensor, dim): + contributed_keys.append(tensor.clone()) + return keys + + monkeypatch.setattr( + compressor_module, + "get_pcp_group", + lambda: SimpleNamespace(all_gather=all_gather), + ) + builder = SimpleNamespace( + capacity=capacity, + device=torch.device("cpu"), + vllm_config=SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_seqs=len(lengths)) + ), + ) + expected_rows = { + block: max(key for key in gathered_keys if key >> 32 == block) + for block in blocks + } + expected_slots = [ + block * capacity + (key & 0xFFFFFFFF) % capacity + for block, key in expected_rows.items() + if key % 2 == 0 + ] + for metadata, positions in zip(metadata_by_rank, positions_by_rank): + num_reqs = metadata.query_start_loc.numel() - 1 + CompressorMetadataBuilder._link_pcp_chunks( + builder, metadata, num_reqs, positions + ) + assert metadata.prev_row_indices is not None + assert metadata.gather_tokens is not None + assert metadata.ring_write_slots is not None + for row, token in enumerate(metadata.gather_tokens.tolist()): + key = int(contributed_keys[-1][row]) + if key >= 0: + assert int(metadata.slot_mapping[token]) // capacity == key >> 32 + assert int(positions[token]) == key & 0xFFFFFFFF + for start, prev in zip( + metadata.query_start_loc[:-1].tolist(), metadata.prev_row_indices.tolist() + ): + position = int(positions[start]) + block = int(metadata.slot_mapping[start]) // capacity + wanted = (block << 32) | (position - 1) + if position % 2 == 1 and wanted in gathered_keys: + assert prev >= 0 and gathered_keys[prev] == wanted + else: + assert prev == -1 + writes = metadata.ring_write_slots.tolist() + assert sorted(slot for slot in writes if slot >= 0) == sorted(expected_slots) + for row, slot in enumerate(writes): + if slot >= 0: + assert gathered_keys[row] == expected_rows[slot // capacity] + assert (gathered_keys[row] & 0xFFFFFFFF) % capacity == slot % capacity + torch.testing.assert_close(torch.cat(contributed_keys), keys) + + +def _build_pcp_ring_metadata(ranks, all_gather, spec, config, device, monkeypatch): + from vllm.models.deepseek_v41 import compressor as compressor_module + from vllm.models.deepseek_v41.compressor import CompressorMetadataBuilder + + monkeypatch.setattr( + compressor_module, + "get_pcp_group", + lambda: SimpleNamespace(all_gather=all_gather), + ) + for local in ranks: + builder = CompressorMetadataBuilder(spec, ["ring"], config, device) + local.metadata = builder.build(0, local.common) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Triton kernel") +@pytest.mark.parametrize("pcp_size", [2, 4]) +def test_v41_pcp_compressor_preserves_pairs_and_decode_state(pcp_size, monkeypatch): + """Pairs split across PCP chunks and decode owner changes keep exact latents.""" + from vllm.forward_context import ForwardContext, override_forward_context + from vllm.models.deepseek_v41 import compressor as compressor_module + from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import ( + fused_save_compress_norm, + ) + from vllm.models.deepseek_v41.compressor import DeepseekCompressor + from vllm.v1.attention.backend import CommonAttentionMetadata + from vllm.v1.kv_cache_interface import CircularBufferSpec + from vllm.v1.worker.gpu.pcp_manager import PCPManager + + torch.manual_seed(42) + device = torch.device("cuda") + norm = torch.randn(512, dtype=torch.bfloat16, device=device) + # Block 0 is reserved as the null block; requests use physical blocks 1/2. + reference_state = torch.zeros(3, 8, 1024, device=device) + rank_states = [reference_state.clone() for _ in range(pcp_size)] + config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=32, max_num_seqs=4), + parallel_config=SimpleNamespace(prefill_context_parallel_size=pcp_size), + speculative_config=None, + ) + spec = CircularBufferSpec( + block_size=8, num_kv_heads=1, head_size=1024, head_size_v=0, dtype=torch.float32 + ) + starts = np.zeros(2, dtype=np.int32) + for lengths, prefilling in [ + ([7, 3], [True, True]), + ([9, 1], [True, False]), + ([1, 1], [False, False]), + ([1, 1], [False, False]), + ]: + lengths = np.array(lengths, dtype=np.int32) + query_start = np.concatenate(([0], np.cumsum(lengths))).astype(np.int32) + qsl = torch.from_numpy(query_start).to(device) + positions = torch.cat( + [torch.arange(s, s + n, device=device) for s, n in zip(starts, lengths)] + ) + req_ids = torch.repeat_interleave( + torch.arange(2, device=device, dtype=torch.int32), + torch.from_numpy(lengths).to(device), + ) + slots = (req_ids.long() + 1) * 8 + positions % 8 + raw = torch.randn(sum(lengths), 1024, device=device) + expected = torch.empty(sum(lengths), 512, dtype=torch.bfloat16, device=device) + fused_save_compress_norm( + raw, + positions, + reference_state, + slots, + qsl, + req_ids, + norm, + 1e-20, + 2, + expected, + ) + + ranks = [] + for rank in range(pcp_size): + manager = PCPManager(pcp_size, rank, device, shard_decode_requests=True) + segments, counts = manager._build_batch_layout( + lengths, starts, np.array(prefilling), query_start + ) + width = max(counts) + local_indices = manager._padded_gather_idx[ + rank * width : (rank + 1) * width + ] + local_lengths = [s.num_tokens for s in segments[rank]] or [0] + local_qsl = torch.tensor([0, *np.cumsum(local_lengths)], dtype=torch.int32) + local_reqs = [s.global_batch_req_idx for s in segments[rank]] or [0] + common = CommonAttentionMetadata( + query_start_loc=local_qsl.to(device), + query_start_loc_cpu=local_qsl, + seq_lens=torch.zeros(len(local_reqs), dtype=torch.int32, device=device), + num_reqs=len(local_reqs), + num_actual_tokens=counts[rank], + max_query_len=max(local_lengths), + max_seq_len=int(max(starts + lengths)), + block_table_tensor=torch.tensor( + np.array(local_reqs) + 1, dtype=torch.int32, device=device + ).unsqueeze(1), + slot_mapping=torch.full((width * pcp_size,), -1, device=device), + positions=positions[local_indices], + ) + ranks.append( + SimpleNamespace( + indices=local_indices[: counts[rank]], + raw=raw[local_indices], + positions=common.positions, + common=common, + ) + ) + + # Each rank's builder contributes its last-row candidates to a gather. + candidates: list[torch.Tensor] = [] + _build_pcp_ring_metadata( + ranks, + lambda t, dim, out=candidates: out.append(t.clone()) or t.repeat(pcp_size), + spec, + config, + device, + monkeypatch, + ) + _build_pcp_ring_metadata( + ranks, + lambda t, dim, out=candidates: torch.cat(out), + spec, + config, + device, + monkeypatch, + ) + gathered_rows = torch.cat([r.raw[r.metadata.gather_tokens] for r in ranks]) + monkeypatch.setattr( + compressor_module, + "get_pcp_group", + lambda rows=gathered_rows: SimpleNamespace(all_gather=lambda t, dim: rows), + ) + for rank, local in enumerate(ranks): + compressor = SimpleNamespace( + pcp_size=pcp_size, + head_dim=512, + compress_ratio=2, + rms_norm_eps=1e-20, + norm=SimpleNamespace(weight=norm), + state_cache=SimpleNamespace(prefix="ring", kv_cache=rank_states[rank]), + ) + context = ForwardContext({}, {"ring": local.metadata}, {}) + with override_forward_context(context): + actual = DeepseekCompressor.forward( + compressor, local.raw, local.positions + ) + assert isinstance(actual, torch.Tensor) + valid = (positions[local.indices] + 1) % 2 == 0 + torch.testing.assert_close( + actual[: local.indices.numel()][valid], + expected[local.indices][valid], + rtol=0, + atol=0, + ) + for state in rank_states[1:]: + torch.testing.assert_close(state, rank_states[0], rtol=0, atol=0) + starts += lengths + + def _on_gfx950() -> bool: if not current_platform.is_rocm(): return False diff --git a/tests/kernels/test_engram.py b/tests/kernels/test_engram.py index 3663108d9d..dd80f3d2cf 100644 --- a/tests/kernels/test_engram.py +++ b/tests/kernels/test_engram.py @@ -502,11 +502,12 @@ def test_lookback_window_reproduces_single_instance(runner, capture): @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") -def test_v2_model_state_gathers_lookback_window(): +@pytest.mark.parametrize("pcp", [False, True]) +def test_v2_model_state_gathers_lookback_window(pcp): """The window is gathered on device from the runner's token history.""" from vllm.models.deepseek_v41.nvidia.model_state import DeepseekV41ModelState - depth, max_num_reqs, max_model_len = 3, 4, 16 + depth, max_num_reqs, max_model_len = 3, 6, 16 state = DeepseekV41ModelState.__new__(DeepseekV41ModelState) state.rope_state = None state.supports_mm_inputs = False @@ -527,15 +528,23 @@ def test_v2_model_state_gathers_lookback_window(): # Batch rows -> request state rows: a request two decode steps in, a fresh # request, and one at the end of a 5-token prompt. input_batch = SimpleNamespace( - idx_mapping=torch.tensor([2, 0, 3], dtype=torch.int64, device="cuda") + idx_mapping=torch.tensor( + [2, 0, 3, 2] if pcp else [2, 0, 3], dtype=torch.int64, device="cuda" + ), + query_start_loc=torch.arange(5 if pcp else 4, dtype=torch.int32, device="cuda"), + positions=torch.tensor([8, 0, 6, 5] if pcp else [8, 0, 6], device="cuda"), ) window = state.prepare_inputs(input_batch, req_states)["lookback_token_ids"] - assert window.cpu().tolist() == [ + expected = [ [22, 21, 16], [-1, -1, -1], [0, 0, 0], - [-1, -1, -1], ] + if pcp: + # The second PCP chunk of request 2 starts at 5, not at its global 8. + expected.append([15, 14, 13]) + expected.extend([[-1] * depth] * (max_num_reqs - len(expected))) + assert window.cpu().tolist() == expected # Graph capture runs on the dummy inputs and replays read the same buffer. dummy = state.prepare_dummy_inputs(num_reqs=3, num_tokens=3)["lookback_token_ids"] assert dummy is window and torch.all(dummy == -1) diff --git a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py index c41af1e4b9..a14c889caa 100644 --- a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py +++ b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py @@ -525,9 +525,8 @@ def test_kv_path_mxfp8_record_matches_reference(num_tokens: int, block_size: int step = 32.0 * tile_scale.repeat_interleave(32, dim=-1) assert ((rec_fused - rec_ref).abs() <= step).all() # The NoPE dims see no rotation at all, so they must be byte-identical. - nope_bytes = slice(0, block_size * HEAD_DIM) - ref_rows = k_cache_ref[:, nope_bytes].view(num_blocks, block_size, HEAD_DIM) - fused_rows = k_cache_fused[:, nope_bytes].view(num_blocks, block_size, HEAD_DIM) + ref_rows = k_cache_ref.view(num_blocks, block_size, V41_HEAD_BYTES) + fused_rows = k_cache_fused.view(num_blocks, block_size, V41_HEAD_BYTES) torch.testing.assert_close( fused_rows[..., :NOPE_DIM], ref_rows[..., :NOPE_DIM], rtol=0, atol=0 ) diff --git a/tests/kernels/test_gate_linear_rocm_dispatch.py b/tests/kernels/test_gate_linear_rocm_dispatch.py index 06d69fc985..22fc355415 100644 --- a/tests/kernels/test_gate_linear_rocm_dispatch.py +++ b/tests/kernels/test_gate_linear_rocm_dispatch.py @@ -109,7 +109,7 @@ def test_rocm_set_out_dtype_respects_bias_guard(monkeypatch): @pytest.mark.parametrize( ("input_size", "output_size"), - [(3072, 256), (4096, 8), (4096, 192), (6144, 128), (6144, 256)], + [(3072, 256), (4096, 8), (6144, 128), (2048, 64), (2560, 512), (7168, 384)], ) def test_rocm_gfx950_enables_fp32_router_gemm( monkeypatch, input_size: int, output_size: int @@ -147,7 +147,6 @@ def test_rocm_fp32_router_gemm_is_replicated_across_ep_sizes( ("input_size", "output_size", "params_dtype", "bias", "on_gfx950"), [ pytest.param(6144, 128, torch.float32, False, False, id="gfx942"), - pytest.param(2048, 64, torch.float32, False, True, id="shape"), pytest.param(6144, 128, torch.bfloat16, False, True, id="bf16-weight"), pytest.param(6144, 128, torch.float32, True, True, id="bias"), ], diff --git a/tests/kernels/test_rocm_fp32_router_gemm.py b/tests/kernels/test_rocm_fp32_router_gemm.py index 899045d39c..efa6e6aeb5 100644 --- a/tests/kernels/test_rocm_fp32_router_gemm.py +++ b/tests/kernels/test_rocm_fp32_router_gemm.py @@ -5,6 +5,9 @@ import pytest import torch +from vllm.model_executor.layers.fused_moe.router import ( + rocm_fp32_router_gemm as rocm_router, +) from vllm.model_executor.layers.fused_moe.router.gate_linear import ( fp32_router_gemm_dispatch_impl, ) @@ -19,6 +22,11 @@ SHAPES = [ (4096, 192), (6144, 128), (6144, 256), + (2560, 512), + (2880, 128), + (5120, 16), + (7168, 384), + (1000, 100), ] MAX_TOKENS = 32 ATOL = 5e-4 @@ -80,7 +88,7 @@ def test_rocm_fp32_router_gemm_preserves_topk( x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=device) weight = torch.randn(num_experts, hidden_size, dtype=torch.float32, device=device) - top_k = 2 if (hidden_size, num_experts) == (4096, 8) else 8 + top_k = 2 if num_experts <= 16 else 8 actual = rocm_fp32_router_gemm(x, weight).topk(top_k, dim=-1).indices reference = _reference(x, weight) reference_values, expected = reference.topk(top_k, dim=-1) @@ -113,15 +121,22 @@ def test_rocm_fp32_router_gemm_rejects_invalid_inputs() -> None: weight, ) - with pytest.raises(ValueError, match="shape"): - rocm_fp32_router_gemm(x[:, :2048].contiguous(), weight[:, :2048].contiguous()) + with pytest.raises(ValueError, match="hidden sizes"): + rocm_fp32_router_gemm(x, weight[:, :2048].contiguous()) + + with pytest.raises(ValueError, match="num_experts"): + big_weight = torch.empty(1024, 7168, dtype=torch.float32, device=device) + rocm_fp32_router_gemm( + torch.empty(MAX_TOKENS, 7168, dtype=torch.bfloat16, device=device), + big_weight, + ) with pytest.raises(ValueError, match="contiguous"): wide_x = torch.randn(4, hidden_size * 2, dtype=torch.bfloat16, device=device) rocm_fp32_router_gemm(wide_x[:, ::2], weight) - with pytest.raises(ValueError, match="float32"): - rocm_fp32_router_gemm(x, weight.to(torch.bfloat16)) + with pytest.raises(ValueError, match="float32 or bfloat16"): + rocm_fp32_router_gemm(x, weight.to(torch.float16)) with pytest.raises(ValueError, match="bfloat16 or float32"): rocm_fp32_router_gemm(x.to(torch.float16), weight) @@ -226,3 +241,154 @@ def test_rocm_fp32_router_gemm_dynamic_compile_dispatch() -> None: ) output = compiled_dispatch(x) torch.testing.assert_close(output, _reference(x, weight), atol=ATOL, rtol=RTOL) + + +@pytest.mark.parametrize("hidden_size", [64, 100, 513, 1000, 2049, 3000]) +@pytest.mark.parametrize("num_tokens", [1, 5, 32]) +@torch.inference_mode() +def test_rocm_fp32_router_gemm_any_hidden_size(num_tokens: int, hidden_size: int): + """Hidden sizes that do not divide the K tile exercise the masked tail.""" + torch.manual_seed(6100 + num_tokens + hidden_size) + x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(64, hidden_size, dtype=torch.float32, device="cuda") + torch.testing.assert_close( + rocm_fp32_router_gemm(x, weight), _reference(x, weight), atol=ATOL, rtol=RTOL + ) + + +@pytest.mark.parametrize("case", ["large", "near_zero", "zero_input", "ones"]) +@torch.inference_mode() +def test_rocm_fp32_router_gemm_special_values(case: str) -> None: + torch.manual_seed(7100) + x = torch.randn(4, 5120, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(128, 5120, dtype=torch.float32, device="cuda") + if case in ("large", "near_zero"): + scale = 100.0 if case == "large" else 1e-4 + x = (x * scale).to(torch.bfloat16) + weight *= scale + elif case == "zero_input": + x.zero_() + else: + x.fill_(1.0) + output = rocm_fp32_router_gemm(x, weight) + expected = _reference(x, weight) + assert torch.isfinite(output).all() + if case == "zero_input": + assert torch.equal(output, torch.zeros_like(output)) + else: + assert ((output - expected).norm() / expected.norm()).item() < 1e-5 + + +@pytest.mark.parametrize("num_tokens", [1, 4, 16, 32]) +@torch.inference_mode() +def test_rocm_fp32_router_gemm_deterministic(num_tokens: int) -> None: + torch.manual_seed(8100 + num_tokens) + x = torch.randn(num_tokens, 7168, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(384, 7168, dtype=torch.float32, device="cuda") + first = rocm_fp32_router_gemm(x, weight) + assert torch.equal(first, rocm_fp32_router_gemm(x, weight)) + + +@pytest.mark.parametrize("num_tokens", [1, 4, 16, 32]) +@pytest.mark.parametrize(("hidden_size", "num_experts"), SHAPES) +@torch.inference_mode() +def test_rocm_fp32_router_gemm_bf16_weight( + num_tokens: int, hidden_size: int, num_experts: int +) -> None: + torch.manual_seed(9100 + num_tokens + hidden_size + num_experts) + x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(num_experts, hidden_size, dtype=torch.bfloat16, device="cuda") + # Past this size hipBLASLt beats the kernel for BF16 weights on MI355. + eligible = num_tokens * num_experts <= 8192 + assert rocm_router.can_use_rocm_fp32_router_gemm(x, weight) == eligible + if eligible: + torch.testing.assert_close( + rocm_fp32_router_gemm(x, weight), + _reference(x.float(), weight.float()), + atol=ATOL, + rtol=RTOL, + ) + + +def _make_gate_linear( + monkeypatch: pytest.MonkeyPatch, + input_size, + output_size, + params_dtype=torch.float32, +): + for module in ("layers.linear", "parameter"): + monkeypatch.setattr( + f"vllm.model_executor.{module}.get_tensor_model_parallel_rank", lambda: 0 + ) + monkeypatch.setattr( + f"vllm.model_executor.{module}.get_tensor_model_parallel_world_size", + lambda: 1, + ) + from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear + + gate = GateLinear( + input_size=input_size, + output_size=output_size, + bias=False, + out_dtype=torch.float32, + params_dtype=params_dtype, + ).cuda() + with torch.no_grad(): + gate.weight.normal_() + gate.quant_method.process_weights_after_loading(gate) + return gate + + +@pytest.mark.parametrize("x_dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("num_tokens", [1, 32, 33]) +@torch.inference_mode() +def test_gate_linear_uses_rocm_fp32_router_gemm( + monkeypatch: pytest.MonkeyPatch, num_tokens: int, x_dtype: torch.dtype +) -> None: + """A router shape outside the old tuned list takes the fast path for M<=32.""" + gate = _make_gate_linear(monkeypatch, 2560, 512) + assert gate.allow_fp32_router_gemm + calls = [] + kernel = rocm_router.rocm_fp32_router_gemm + + def spy(hidden_states, router_weight): + calls.append(hidden_states.shape[0]) + return kernel(hidden_states, router_weight) + + monkeypatch.setattr(rocm_router, "rocm_fp32_router_gemm", spy) + x = torch.randn(num_tokens, 2560, dtype=x_dtype, device="cuda") + output, bias = gate(x) + + assert bias is None + assert calls == ([num_tokens] if num_tokens <= MAX_TOKENS else []) + torch.testing.assert_close(output, _reference(x, gate.weight), atol=ATOL, rtol=RTOL) + + +@pytest.mark.parametrize( + ("num_tokens", "num_experts", "uses_kernel"), + [(1, 256, True), (32, 256, True), (33, 256, False), (32, 512, False)], +) +@torch.inference_mode() +def test_gate_linear_bf16_weight_uses_rocm_router_gemm( + monkeypatch: pytest.MonkeyPatch, num_tokens, num_experts, uses_kernel +) -> None: + """BF16-weight gates with FP32 logits use the kernel only where it beats + hipBLASLt; everything else keeps the torch.mm path.""" + gate = _make_gate_linear(monkeypatch, 7168, num_experts, torch.bfloat16) + assert gate.allow_cublas_router_gemm + calls = [] + kernel = rocm_router.rocm_fp32_router_gemm + + def spy(hidden_states, router_weight): + calls.append(hidden_states.shape[0]) + return kernel(hidden_states, router_weight) + + monkeypatch.setattr(rocm_router, "rocm_fp32_router_gemm", spy) + x = torch.randn(num_tokens, 7168, dtype=torch.bfloat16, device="cuda") + output, _ = gate(x) + + assert output.dtype == torch.float32 + assert calls == ([num_tokens] if uses_kernel else []) + torch.testing.assert_close( + output, _reference(x.float(), gate.weight.float()), atol=ATOL, rtol=RTOL + ) diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py index 7ca4cff404..8c071b6730 100644 --- a/tests/model_executor/layers/test_mla_short_prefill_indexer.py +++ b/tests/model_executor/layers/test_mla_short_prefill_indexer.py @@ -830,15 +830,19 @@ def test_sparse_mqa_logits_prefill_matches_dense_masked(capture: bool): _assert_same_selection(dense, out, ref) -def _build_fp4_paged_cache(req_lens: torch.Tensor, page_kv: int, head_dim: int): +def _build_fp4_paged_cache( + req_lens: torch.Tensor, page_kv: int, head_dim: int, page_pad: int = 0 +): """MXFP4 paged cache in DeepGEMM's fused layout: per page, all tokens' - fp4 data first, then their packed UE8M0 scales.""" + fp4 data first, then their packed UE8M0 scales; ``page_pad`` bytes pad + each page like the other layers' pages in a packed block.""" device = req_lens.device pages_per = (req_lens + page_kv - 1) // page_kv num_pages, max_pages = int(pages_per.sum()), int(pages_per.max()) + head_bytes = head_dim // 2 + 4 kv_cache = torch.zeros( - num_pages, page_kv, 1, head_dim // 2 + 4, dtype=torch.uint8, device=device - ) + num_pages, page_kv * head_bytes + page_pad, dtype=torch.uint8, device=device + )[:, : page_kv * head_bytes].view(num_pages, page_kv, 1, head_bytes) block_table = torch.zeros( len(req_lens), max_pages, dtype=torch.int32, device=device ) @@ -864,7 +868,10 @@ def _build_fp4_paged_cache(req_lens: torch.Tensor, page_kv: int, head_dim: int): @pytest.mark.parametrize("next_n", [1, 2]) @pytest.mark.parametrize("capture", [False, True]) -def test_sparse_mqa_logits_paged_matches_dense_masked(next_n: int, capture: bool): +@pytest.mark.parametrize("page_pad", [0, 256]) +def test_sparse_mqa_logits_paged_matches_dense_masked( + next_n: int, capture: bool, page_pad: int +): """Paged sparse-MQA decode selects exactly the values ``torch.topk`` picks from the dense bf16 masked logits, incl. MTP rows (next_n > 1 flattened to one query row each), CUDA graph capture/replay and schedule reuse by a @@ -886,8 +893,10 @@ def test_sparse_mqa_logits_paged_matches_dense_masked(next_n: int, capture: bool .clamp(min=0) .to(torch.int32) ) - kv_cache, block_table = _build_fp4_paged_cache(req_lens, page_kv, head_dim) - assert kv_cache.stride(0) % 512 == 0 # the paged sparse kernel requires it + kv_cache, block_table = _build_fp4_paged_cache( + req_lens, page_kv, head_dim, page_pad + ) + assert kv_cache.stride(0) % 512 == page_pad # padded pages are not 512B q_fp, q_sf = _quant_fp4( torch.randn(rows * num_heads, head_dim, device="cuda", dtype=torch.bfloat16) diff --git a/tests/model_executor/layers/test_rocm_unquantized_gemm.py b/tests/model_executor/layers/test_rocm_unquantized_gemm.py index 491cfa5435..db736c7438 100644 --- a/tests/model_executor/layers/test_rocm_unquantized_gemm.py +++ b/tests/model_executor/layers/test_rocm_unquantized_gemm.py @@ -17,6 +17,15 @@ if current_platform.is_cuda(): from vllm.model_executor.layers import utils +def _on_gfx11() -> bool: + """gfx11 check that is safe to evaluate at collection time on any platform.""" + if not current_platform.is_rocm(): + return False + from vllm.platforms.rocm import on_gfx11 + + return on_gfx11() + + @pytest.mark.parametrize("m", [1, 128]) def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch, m): x = torch.randn(1, 64, dtype=torch.float16) @@ -145,6 +154,101 @@ def test_rocm_unquantized_gemm_rejects_unsupported_skinny_layouts( assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3) +def _row_stride_padded(weight: torch.Tensor, pad: int) -> torch.Tensor: + """A weight whose row stride is offset by ``pad`` elements, as produced by + ``maybe_pad_weight_avoid_cache_aliasing``.""" + padded = torch.nn.functional.pad(weight, (0, pad))[..., :-pad] + assert padded.shape == weight.shape + assert padded.stride() == (weight.shape[1] + pad, 1) + return padded + + +def test_rocm_unquantized_gemm_wvsplitk_accepts_row_stride_padded_weight(monkeypatch): + # wvSplitK reads the weight row stride, so a row-stride-padded weight must + # still reach it instead of falling back to torch.nn.functional.linear. + x = torch.randn(4, 1024, dtype=torch.float16) + weight = _row_stride_padded(torch.randn(128, 1024, dtype=torch.float16), 64) + + monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False) + monkeypatch.setattr(utils.rocm_aiter_ops, "is_tgemm_enabled", lambda: False) + monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx11", lambda: True) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: False) + monkeypatch.setattr(utils, "num_compute_units", lambda: 120) + + wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t()) + monkeypatch.setattr(utils.ops, "wvSplitK", wvsplitk_mock) + + out = utils.rocm_unquantized_gemm_impl(x, weight, None) + ref = torch.nn.functional.linear(x, weight, None) + + wvsplitk_mock.assert_called_once() + assert wvsplitk_mock.call_args.args[0].stride() == weight.stride() + assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3) + + +@pytest.mark.parametrize( + "k,element_size,expected", + [ + (4096, 2, 64), # 8192 B row stride, aliases + (12288, 2, 64), # 24576 B row stride, aliases + (1024, 2, 64), # smallest aliasing bf16/fp16 K + (2560, 2, 0), # 5120 B, no aliasing + (11008, 2, 0), # 22016 B, no aliasing + (2048, 1, 128), # fp8: 2048 B, aliases, twice the pad + (512, 4, 32), # fp32: 2048 B, aliases, half the pad + ], +) +def test_cache_aliasing_pad_elems(k, element_size, expected): + assert utils.cache_aliasing_pad_elems(k, element_size) == expected + + +def test_maybe_pad_weight_avoid_cache_aliasing_skips_non_gfx11(monkeypatch): + # Padding is gfx11-only, even for a K whose row stride would alias there. + monkeypatch.setattr(utils.current_platform, "is_rocm", lambda: True) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx11", lambda: False) + + weight = torch.randn(128, 4096, dtype=torch.bfloat16) + assert utils.maybe_pad_weight_avoid_cache_aliasing(weight) is weight + + +def test_maybe_pad_weight_avoid_cache_aliasing(monkeypatch): + monkeypatch.setattr(utils.current_platform, "is_rocm", lambda: True) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx11", lambda: True) + + weight = torch.randn(128, 4096, dtype=torch.bfloat16) + padded = utils.maybe_pad_weight_avoid_cache_aliasing(weight) + + assert padded.shape == weight.shape + assert padded.stride() == (4096 + 64, 1) + assert torch.equal(padded, weight) + # Idempotent: the padded stride no longer aliases, so a second call is a no-op + # and preserves the pointer captured by any graph. + repadded = utils.maybe_pad_weight_avoid_cache_aliasing(padded) + assert repadded.data_ptr() == padded.data_ptr() + + +@pytest.mark.parametrize( + "weight", + [ + torch.randn(128, 2560, dtype=torch.bfloat16), # no aliasing + torch.randn(2560, 128, dtype=torch.bfloat16).t(), # not contiguous + torch.randn(2, 128, 4096, dtype=torch.bfloat16), # not 2D + ], +) +def test_maybe_pad_weight_avoid_cache_aliasing_skips(monkeypatch, weight): + monkeypatch.setattr(utils.current_platform, "is_rocm", lambda: True) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx11", lambda: True) + + assert ( + utils.maybe_pad_weight_avoid_cache_aliasing(weight).data_ptr() + == weight.data_ptr() + ) + + @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-only kernel test") @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) def test_rocm_unquantized_gemm_noncontiguous_activation_real_kernel(monkeypatch, dtype): @@ -164,6 +268,26 @@ def test_rocm_unquantized_gemm_noncontiguous_activation_real_kernel(monkeypatch, torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2) +@pytest.mark.skipif( + not _on_gfx11(), reason="cache-aliasing padding only applies on gfx11" +) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_rocm_unquantized_gemm_padded_weight_real_kernel(monkeypatch, dtype): + x = torch.randn(4, 4096, device="cuda", dtype=dtype) + weight = _row_stride_padded(torch.randn(128, 4096, device="cuda", dtype=dtype), 64) + + monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True) + original_wvsplitk = utils.ops.wvSplitK + wvsplitk_mock = MagicMock(side_effect=original_wvsplitk) + monkeypatch.setattr(utils.ops, "wvSplitK", wvsplitk_mock) + + out = utils.rocm_unquantized_gemm_impl(x, weight, None) + ref = torch.nn.functional.linear(x, weight, None) + + wvsplitk_mock.assert_called_once() + torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2) + + @pytest.mark.parametrize("m", [1, 128]) def test_rocm_unquantized_gemm_gfx1x_n_gt_5_falls_back(monkeypatch, m): # wvSplitK skinny GEMM handles n in [1, 5] (see PR #40687); n > 5 must diff --git a/tests/model_executor/test_deep_gemm_warmup.py b/tests/model_executor/test_deep_gemm_warmup.py index c2dc12a27d..26c410d1bf 100644 --- a/tests/model_executor/test_deep_gemm_warmup.py +++ b/tests/model_executor/test_deep_gemm_warmup.py @@ -123,19 +123,12 @@ def test_mega_moe_m_values_reach_mega_gate(monkeypatch) -> None: threshold (1 for 128 experts), so the forward actually reaches the kernel.""" dg = SimpleNamespace( get_bf16_mega_gate_config=lambda m, *_: {"num_sms": -(-m // 16)}, - get_block_m_for_mega_moe=lambda *_: 16, ) monkeypatch.setattr(deep_gemm_warmup, "_import_deep_gemm", lambda: dg) monkeypatch.setattr(deep_gemm_warmup, "_mega_moe_uses_mega_gate", lambda _: True) - monkeypatch.setattr( - deep_gemm_warmup, "get_ep_group", lambda: SimpleNamespace(world_size=8) - ) - buffer = SimpleNamespace(num_max_tokens_per_rank=64) module = SimpleNamespace( gate=SimpleNamespace(input_size=4096, output_size=128), - experts=SimpleNamespace( - top_k=6, num_experts=128, get_symm_buffer=lambda: buffer - ), + experts=SimpleNamespace(top_k=6, get_block_m=lambda num_tokens: 16), use_sequence_parallel=False, ) assert deep_gemm_warmup._get_fp8_fp4_mega_moe_m_values(module, 40) == [16, 32, 40] diff --git a/tests/model_executor/test_enabled_custom_ops.py b/tests/model_executor/test_enabled_custom_ops.py index 490284f439..27985bba65 100644 --- a/tests/model_executor/test_enabled_custom_ops.py +++ b/tests/model_executor/test_enabled_custom_ops.py @@ -139,6 +139,122 @@ def test_topk_softmax_dispatch(use_rocm_aiter: bool): assert topk_func == vllm_topk_softmax +def _topk_gating_launch( + num_tokens: int = 4, + num_experts: int = 512, + topk: int = 10, + num_shared_experts: int = 0, + weights_width: int | None = None, + ids_width: int | None = None, + contiguous: bool = True, +) -> dict: + """Tensors for one softmax top-k launch, as dispatch_topk_softmax_func sees them.""" + width = topk + num_shared_experts + weights = torch.empty(num_tokens, weights_width or width) + ids = torch.empty(num_tokens, ids_width or width, dtype=torch.int32) + ids = ids[:, :topk] + if contiguous: + gating = torch.empty(num_tokens, num_experts + num_shared_experts) + else: + gating = torch.empty(num_experts + num_shared_experts + 3, num_tokens).t() + return dict(topk_weights=weights, topk_indices=ids, gating_output=gating) + + +@pytest.mark.parametrize( + "launch, num_shared_experts, scoring_func, expect_gating", + [ + pytest.param({}, 0, "", True, id="plain-softmax"), + pytest.param({"num_tokens": 4096}, 0, "", True, id="max-tokens"), + pytest.param({"num_tokens": 4097}, 0, "", False, id="too-many-tokens"), + pytest.param({"contiguous": False}, 0, "", False, id="non-contiguous"), + pytest.param({}, 0, "sigmoid", False, id="scoring-without-shared"), + pytest.param({"num_shared_experts": 1}, 1, "sigmoid", True, id="one-shared"), + pytest.param({"num_shared_experts": 8}, 8, "sigmoid", True, id="eight-shared"), + pytest.param( + {"num_shared_experts": 3}, 3, "sigmoid", False, id="unsupported-count" + ), + pytest.param({"num_shared_experts": 1}, 1, "", False, id="no-shared-scoring"), + pytest.param( + {"num_shared_experts": 1}, 1, "softmax", False, id="wrong-shared-scoring" + ), + pytest.param( + {"num_shared_experts": 1, "weights_width": 10}, + 1, + "sigmoid", + False, + id="weights-too-narrow", + ), + pytest.param( + {"num_shared_experts": 1, "ids_width": 12}, + 1, + "sigmoid", + False, + id="row-stride-mismatch", + ), + pytest.param( + {"num_tokens": 4097, "num_shared_experts": 1}, + 1, + "sigmoid", + False, + id="shared-too-many-tokens", + ), + ], +) +@pytest.mark.parametrize("gating_enabled", [True, False]) +def test_topk_softmax_dispatch_aiter_topk_gating( + monkeypatch: pytest.MonkeyPatch, + launch: dict, + num_shared_experts: int, + scoring_func: str, + expect_gating: bool, + gating_enabled: bool, +): + """AITER launches use topk_gating only where it is enabled, and only when it + supports them. + + Every other AITER launch must reach the legacy topk_softmax, so this is the + one place the choice is made. + """ + monkeypatch.setattr( + rocm_aiter_ops, "is_topk_gating_enabled", lambda: gating_enabled + ) + topk_func = dispatch_topk_softmax_func( + True, + **_topk_gating_launch(**launch), + num_shared_experts=num_shared_experts, + shared_expert_scoring_func=scoring_func, + ) + if expect_gating and gating_enabled: + assert topk_func == rocm_aiter_ops.topk_gating + else: + assert topk_func == rocm_aiter_ops.topk_softmax + + +@pytest.mark.parametrize( + "cdna, fused_moe_enabled, expected", + [ + pytest.param(3, True, True, id="cdna3"), + pytest.param(4, True, True, id="cdna4"), + pytest.param(2, True, False, id="cdna2"), + pytest.param(3, False, False, id="cdna3-aiter-moe-off"), + pytest.param(4, False, False, id="cdna4-aiter-moe-off"), + ], +) +def test_topk_gating_enabled_by_cdna( + monkeypatch: pytest.MonkeyPatch, + cdna: int, + fused_moe_enabled: bool, + expected: bool, +): + """topk_gating is on for CDNA3 and CDNA4 with AITER MoE enabled, nowhere else.""" + monkeypatch.setattr("vllm._aiter_ops.is_aiter_found_and_supported", lambda: True) + monkeypatch.setattr( + rocm_aiter_ops, "is_fused_moe_enabled", lambda: fused_moe_enabled + ) + monkeypatch.setattr("vllm.platforms.rocm.get_cdna_version", lambda: cdna) + assert bool(rocm_aiter_ops.is_topk_gating_enabled()) is expected + + @pytest.mark.parametrize( "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False] ) diff --git a/tests/model_executor/test_routed_experts_capture.py b/tests/model_executor/test_routed_experts_capture.py index 5416fdc8ea..896a214a4a 100644 --- a/tests/model_executor/test_routed_experts_capture.py +++ b/tests/model_executor/test_routed_experts_capture.py @@ -7,9 +7,16 @@ from unittest.mock import Mock, patch import pytest import torch -from vllm.config import VllmConfig +from vllm.config import CUDAGraphMode, VllmConfig from vllm.config.compilation import CompilationMode from vllm.distributed.eplb.eplb_state import EplbLayerState +from vllm.forward_context import ( + BatchDescriptor, + DPMetadata, + ForwardContext, + get_forward_context, + set_forward_context, +) from vllm.model_executor.layers.fused_moe.config import RoutingMethodType from vllm.model_executor.layers.fused_moe.experts.cpu_int4_moe import CPUExpertsInt4 from vllm.model_executor.layers.fused_moe.modular_kernel import ( @@ -36,6 +43,12 @@ pytestmark = pytest.mark.cpu_test _REC_MODULE = "vllm.model_executor.layers.fused_moe.routed_experts_capturer" +def _forward_context(**kwargs): + return ForwardContext( + no_compile_layers={}, attn_metadata={}, slot_mapping={}, **kwargs + ) + + def _capturer_with_buffer( *, max_tokens: int = 8, @@ -319,13 +332,120 @@ def test_routed_experts_capturer_single_dp_no_metadata(): """dp_metadata is None: capture writes the full topk_ids rows.""" capturer = _capturer_with_buffer(dp_rank=0) topk = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.int32) - ctx = SimpleNamespace(dp_metadata=None) + ctx = _forward_context() with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx): capturer.capture(layer_id=0, topk_ids=topk) assert torch.equal(capturer.device_buffer[:3, 0, :], topk) assert capturer.device_buffer[3, 0, 0].item() == -1 +@pytest.mark.parametrize("tp_size, num_tokens", [(2, 6), (2, 3), (2, 1), (4, 2)]) +@pytest.mark.parametrize( + "cg_mode", [CUDAGraphMode.NONE, CUDAGraphMode.FULL, CUDAGraphMode.PIECEWISE] +) +def test_single_dp_ep_capture_restores_tp_shards(tp_size, num_tokens, cg_mode): + """Capture all current rows after model-level SP, including buffer reuse.""" + capturer = _capturer_with_buffer(tp_size=tp_size) + config = SimpleNamespace( + parallel_config=SimpleNamespace( + data_parallel_size=1, + tensor_parallel_size=tp_size, + enable_expert_parallel=True, + use_sequence_parallel_moe=False, + is_moe_model=True, + ), + compilation_config=SimpleNamespace( + fast_moe_cold_start=False, static_forward_context={} + ), + ) + tp_group = Mock() + shard_n = (num_tokens + tp_size - 1) // tp_size + with ( + set_forward_context( + None, + config, + num_tokens=num_tokens, + cudagraph_runtime_mode=cg_mode, + ), + patch(f"{_REC_MODULE}.get_tp_group", return_value=tp_group), + ): + assert get_forward_context().dp_metadata is None + for offset in (20, 1): + expected = torch.arange(num_tokens * 2, dtype=torch.int32).reshape(-1, 2) + expected += offset + padded = torch.full((shard_n * tp_size, 2), 99, dtype=torch.int32) + padded[:num_tokens] = expected + tp_group.all_gather.return_value = padded + capturer.capture(0, padded[:shard_n]) + assert torch.equal( + capturer.snapshot_routing_data(num_tokens)[:, 0].int(), expected + ) + assert (capturer.device_buffer[num_tokens:, 0] == -1).all() + if num_tokens == 1: + tp_group.all_gather.assert_not_called() + else: + assert tp_group.all_gather.call_count == 2 + + +@pytest.mark.parametrize("num_tokens, num_rows", [(6, 6), (3, 4), (0, 0)]) +def test_single_dp_tp_capture_does_not_regather_full_batch(num_tokens, num_rows): + capturer = _capturer_with_buffer(tp_size=2) + topk = torch.arange(num_rows * 2, dtype=torch.int32).reshape(-1, 2) + ctx = _forward_context(batch_descriptor=BatchDescriptor(num_tokens=num_tokens)) + with ( + patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx), + patch(f"{_REC_MODULE}.get_tp_group") as get_tp_group, + ): + capturer.capture(0, topk) + get_tp_group.assert_not_called() + assert torch.equal(capturer.device_buffer[:num_tokens, 0], topk[:num_tokens]) + assert (capturer.device_buffer[num_tokens:, 0] == -1).all() + + +@pytest.mark.parametrize("num_tokens, num_rows", [(6, 2), (0, 1)]) +def test_single_dp_tp_capture_rejects_unexpected_batch_without_writing( + num_tokens, num_rows +): + capturer = _capturer_with_buffer(tp_size=2) + topk = torch.ones(num_rows, 2, dtype=torch.int32) + ctx = _forward_context(batch_descriptor=BatchDescriptor(num_tokens=num_tokens)) + with ( + patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx), + pytest.raises(AssertionError, match="unexpected topk_ids batch dim"), + ): + capturer.capture(0, topk) + assert (capturer.device_buffer == -1).all() + + +@pytest.mark.parametrize("tp_rank", [0, 1]) +def test_capture_uses_pre_sp_size_after_model_shards_padding_mask(tp_rank): + """Model-level SP shortens the mask; capture still needs all three rows.""" + from vllm.models.common.ops import sequence_parallel as sp + + capturer = _capturer_with_buffer(tp_size=2) + expected = torch.arange(8, dtype=torch.int32).reshape(4, 2) + ctx = _forward_context( + is_padding=torch.zeros(3, dtype=torch.bool), + batch_descriptor=BatchDescriptor(num_tokens=3), + ) + with ( + patch.object(sp, "get_tensor_model_parallel_world_size", return_value=2), + patch.object(sp, "get_tensor_model_parallel_rank", return_value=tp_rank), + ): + ctx.is_padding = sp.sp_padding_mask(ctx.is_padding, torch.empty(3, 1)) + assert len(ctx.is_padding) == 2 + tp_group = Mock() + tp_group.all_gather.return_value = expected + with ( + patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx), + patch(f"{_REC_MODULE}.get_tp_group", return_value=tp_group), + ): + capturer.capture(0, expected[2 * tp_rank : 2 * (tp_rank + 1)]) + tp_group.all_gather.assert_called_once() + torch.testing.assert_close(capturer.device_buffer[:3, 0], expected[:3]) + assert (capturer.device_buffer[3:, 0] == -1).all() + + @pytest.mark.parametrize( ("num_experts", "dtype_name", "torch_dtype"), [(256, "uint8", torch.uint8), (257, "uint16", torch.uint16)], @@ -359,7 +479,7 @@ def test_routed_experts_capturer_narrows_snapshot(output_dtype): capturer = _capturer_with_buffer(dtype=torch.int32) capturer.output_dtype = output_dtype topk = torch.tensor([[1, 2], [254, 255]], dtype=torch.int64) - ctx = SimpleNamespace(dp_metadata=None) + ctx = _forward_context() with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx): capturer.capture(layer_id=0, topk_ids=topk) @@ -375,8 +495,8 @@ def test_routed_experts_capturer_dp_naive_concatenated_all_ranks(): """N == sum(num_tokens_dp): slice this rank's segment from concatenated topk.""" capturer = _capturer_with_buffer(dp_rank=1) num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32) - ctx = SimpleNamespace( - dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp) + ctx = _forward_context( + dp_metadata=DPMetadata(num_tokens_across_dp_cpu=num_tokens_dp) ) # Concatenated order: rank0 rows then rank1 rows. topk = torch.tensor( @@ -392,8 +512,8 @@ def test_routed_experts_capturer_dp_modular_local_tokens(): """N == token_num_per_dp: topk is already local to this DP rank.""" capturer = _capturer_with_buffer(dp_rank=1) num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32) - ctx = SimpleNamespace( - dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp) + ctx = _forward_context( + dp_metadata=DPMetadata(num_tokens_across_dp_cpu=num_tokens_dp) ) topk = torch.tensor([[10, 11], [12, 13], [14, 15]], dtype=torch.int32) with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx): @@ -404,8 +524,8 @@ def test_routed_experts_capturer_dp_modular_local_tokens(): def test_routed_experts_capturer_dp_ep_gathered_shards(): capturer = _capturer_with_buffer(dp_rank=1, tp_size=2) num_tokens_dp = torch.tensor([3, 2], dtype=torch.int32) - ctx = SimpleNamespace( - dp_metadata=SimpleNamespace( + ctx = _forward_context( + dp_metadata=DPMetadata( num_tokens_across_dp_cpu=num_tokens_dp, local_sizes=[2, 2, 1, 1], ) @@ -421,13 +541,17 @@ def test_routed_experts_capturer_dp_ep_gathered_shards(): assert torch.equal(capturer.device_buffer[:2, 0], topk[4:]) -def test_routed_experts_capturer_sp_modular_gathers_tp_shards(): - capturer = _capturer_with_buffer(dp_rank=1, tp_size=2) +@pytest.mark.parametrize("single_dp", [False, True]) +def test_routed_experts_capturer_sp_modular_gathers_tp_shards(single_dp): + capturer = _capturer_with_buffer(dp_rank=0 if single_dp else 1, tp_size=2) num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32) - ctx = SimpleNamespace( - dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp) + ctx = _forward_context( + dp_metadata=( + None if single_dp else DPMetadata(num_tokens_across_dp_cpu=num_tokens_dp) + ), + batch_descriptor=BatchDescriptor(num_tokens=3), ) - local_shard = torch.tensor([[10, 11], [12, 13]], dtype=torch.int32) + local_shard = torch.tensor([[10, 12], [11, 13]], dtype=torch.int32).t() gathered = torch.tensor([[10, 11], [12, 13], [14, 15], [-1, -1]], dtype=torch.int32) tp_group = Mock() tp_group.all_gather.return_value = gathered @@ -438,16 +562,83 @@ def test_routed_experts_capturer_sp_modular_gathers_tp_shards(): ): capturer.capture(layer_id=0, topk_ids=local_shard) - tp_group.all_gather.assert_called_once_with(local_shard, dim=0) + tp_group.all_gather.assert_called_once() + gather_input = tp_group.all_gather.call_args.args[0] + assert gather_input.is_contiguous() + torch.testing.assert_close(gather_input, local_shard) + assert tp_group.all_gather.call_args.kwargs == {"dim": 0} assert torch.equal(capturer.device_buffer[:3, 0], gathered[:3]) +def test_routed_experts_capturer_rejects_incomplete_dp_shard_groups(): + """Invalid gather metadata must not silently select another rank's rows.""" + capturer = _capturer_with_buffer(dp_rank=1, tp_size=2) + ctx = _forward_context( + dp_metadata=DPMetadata( + num_tokens_across_dp_cpu=torch.tensor([3, 2]), + local_sizes=[2, 2, 2], + ) + ) + with ( + patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx), + pytest.raises(ValueError, match="not divisible by DP rank count"), + ): + capturer.capture(0, torch.ones(6, 2, dtype=torch.int32)) + assert (capturer.device_buffer == -1).all() + + +@pytest.mark.parametrize("layer_id", [-1, 4]) +def test_routed_experts_capturer_rejects_layer_before_restoring_rows(layer_id): + capturer = _capturer_with_buffer(tp_size=2) + ctx = _forward_context(batch_descriptor=BatchDescriptor(num_tokens=3)) + with ( + patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx), + patch(f"{_REC_MODULE}.get_tp_group") as get_tp_group, + pytest.raises(IndexError, match="exceeds capture buffer"), + ): + capturer.capture(layer_id, torch.ones(2, 2, dtype=torch.int32)) + get_tp_group.assert_not_called() + assert (capturer.device_buffer == -1).all() + + +@pytest.mark.parametrize( + ("counts", "local_sizes", "start", "stop"), + [ + ([3, 0], [2, 2, 0, 0], 3, 3), + ([0, 3], None, 0, 3), + ([0, 0], [0, 0, 0, 0], 0, 0), + ([2, 2], [1, 1, 1, 1], 2, 4), + ], + ids=["idle-rank", "only-active-rank", "all-idle", "unpadded-gather-tie"], +) +def test_routed_experts_capturer_dp_preserves_empty_and_unpadded_spans( + counts, local_sizes, start, stop +): + capturer = _capturer_with_buffer(dp_rank=1, tp_size=2) + topk = torch.arange(sum(counts) * 2, dtype=torch.int32).reshape(-1, 2) + ctx = _forward_context( + dp_metadata=DPMetadata( + num_tokens_across_dp_cpu=torch.tensor(counts), local_sizes=local_sizes + ) + ) + with ( + patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx), + patch(f"{_REC_MODULE}.get_tp_group") as get_tp_group, + ): + capturer.capture(0, topk) + get_tp_group.assert_not_called() + torch.testing.assert_close( + capturer.device_buffer[: stop - start, 0], topk[start:stop] + ) + assert (capturer.device_buffer[stop - start :, 0] == -1).all() + + def test_routed_experts_capturer_dp_unexpected_batch_raises(): """Mismatch between topk batch dim and DP layout: fail fast.""" capturer = _capturer_with_buffer(dp_rank=0) num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32) - ctx = SimpleNamespace( - dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp) + ctx = _forward_context( + dp_metadata=DPMetadata(num_tokens_across_dp_cpu=num_tokens_dp) ) # total=5, local=2: n=1 matches neither naive (5) nor modular (2). topk = torch.tensor([[1, 2]], dtype=torch.int32) diff --git a/tests/models/kimi_k3/test_amd_kda_decode.py b/tests/models/kimi_k3/test_amd_kda_decode.py index 647c982c7d..fc6039ee54 100644 --- a/tests/models/kimi_k3/test_amd_kda_decode.py +++ b/tests/models/kimi_k3/test_amd_kda_decode.py @@ -1,11 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The ROCm fused KDA decode kernel must match the Triton chain it replaces. +"""The ROCm fused KDA decode kernels must match the Triton chain they replace. -The fused kernel folds the packed causal conv1d update, the gated delta-rule -recurrence and the gated output RMSNorm into one launch, and updates both the -conv state and the recurrent state in place. Every one of those outputs is -compared against the three-kernel path the AMD layer falls back to. +Two fused paths fold the packed causal conv1d update, the gated delta-rule +recurrence and the gated output RMSNorm into one launch, updating the conv +state and the recurrent state in place: + +* the vLLM HIP kernel ``ops.fused_kda_decode`` (gfx942 / gfx950), and +* AITER's Gluon kernels behind ``rocm_aiter_ops`` (CDNA4+): the packed + non-spec decode and the speculative-decode ``fused_recurrent_kda``. + +Every output of each is compared against the unfused path the AMD layer falls +back to, with the argument layout the layer itself uses. """ import pytest @@ -14,7 +20,7 @@ import torch from vllm.platforms import current_platform -def _on_supported_arch() -> bool: +def _hip_kernel_arch() -> bool: if not current_platform.is_rocm(): return False from vllm.platforms.rocm import on_gfx942, on_gfx950 @@ -22,10 +28,26 @@ def _on_supported_arch() -> bool: return on_gfx950() or on_gfx942() +def _aiter_gluon_kda_available() -> bool: + if not current_platform.is_rocm(): + return False + from vllm._aiter_ops import rocm_aiter_ops + from vllm.platforms.rocm import get_cdna_version + + return get_cdna_version() >= 4 and rocm_aiter_ops._kda_gluon_decode_importable() + + pytestmark = pytest.mark.skipif( - not _on_supported_arch(), + not current_platform.is_rocm(), reason="ROCm-only KDA decode kernels" +) +requires_hip_kernel_arch = pytest.mark.skipif( + not _hip_kernel_arch(), reason="The fused KDA decode kernel is only built for gfx942 / gfx950", ) +requires_aiter_gluon_kda = pytest.mark.skipif( + not _aiter_gluon_kda_available(), + reason="AITER Gluon KDA decode needs CDNA4+ and an AITER build that ships it", +) # Kimi-K3 KDA: 96 heads x 128, conv width 4, gate_lower_bound -5.0. HEAD_DIM = 128 @@ -41,7 +63,11 @@ def _requires_kernel() -> None: class KdaDecodeInputs: - """One decode step of a Kimi-K3 KDA layer, in the layer's own layout.""" + """One decode step of a Kimi-K3 KDA layer, in the layer's own layout. + + ``num_spec`` > 0 sizes the conv state for speculative decoding + (``conv_width - 1 + num_spec`` taps, as ``kda_state_shape`` does). + """ def __init__( self, @@ -49,6 +75,7 @@ class KdaDecodeInputs: num_heads: int, num_slots: int, seed: int = 0, + num_spec: int = 0, ) -> None: torch.manual_seed(seed) device = "cuda" @@ -70,9 +97,15 @@ class KdaDecodeInputs: for i in range(3) ] ) - # SD cache layout: [slots, width - 1, 3 * dim]; the layer transposes it. + # SD cache layout: [slots, taps, 3 * dim]; the layer transposes it. self.conv_state = ( - torch.randn(num_slots, CONV_WIDTH - 1, 3 * dim, device=device, dtype=DTYPE) + torch.randn( + num_slots, + CONV_WIDTH - 1 + num_spec, + 3 * dim, + device=device, + dtype=DTYPE, + ) * 0.5 ) self.recurrent_state = ( @@ -112,6 +145,12 @@ class KdaDecodeInputs: def conv_state_view(self, state: torch.Tensor) -> torch.Tensor: return state.transpose(-1, -2) + def empty_out(self) -> torch.Tensor: + num_tokens = self.mixed_qkv.shape[0] + return torch.empty( + 1, num_tokens, self.num_heads, HEAD_DIM, device="cuda", dtype=DTYPE + ) + def _gated_rmsnorm( x: torch.Tensor, gate: torch.Tensor, weight: torch.Tensor, eps: float @@ -196,6 +235,7 @@ def _run_fused( return out, conv_state, recurrent_state +@requires_hip_kernel_arch @pytest.mark.parametrize("num_heads", [12, 24, 96]) @pytest.mark.parametrize("num_tokens", [1, 7, 128]) @torch.inference_mode() @@ -214,6 +254,7 @@ def test_fused_kda_decode_matches_triton_chain(num_heads: int, num_tokens: int) torch.testing.assert_close(actual_state, expected_state, atol=2e-3, rtol=2e-3) +@requires_hip_kernel_arch @torch.inference_mode() def test_fused_kda_decode_leaves_untouched_slots_alone() -> None: """Only the slots named by state_indices may be written.""" @@ -234,6 +275,7 @@ def test_fused_kda_decode_leaves_untouched_slots_alone() -> None: ) +@requires_hip_kernel_arch @torch.inference_mode() def test_fused_kda_decode_without_output_norm() -> None: """Omitting the gate/norm pair returns the raw recurrent output.""" @@ -265,6 +307,7 @@ def test_fused_kda_decode_without_output_norm() -> None: assert torch.isfinite(out).all() +@requires_hip_kernel_arch @torch.inference_mode() def test_fused_kda_decode_skips_null_block_padding() -> None: """A CUDA-graph decode batch is padded with NULL_BLOCK_ID (0). @@ -291,6 +334,7 @@ def test_fused_kda_decode_skips_null_block_padding() -> None: torch.testing.assert_close(actual_state, expected_state, atol=2e-3, rtol=2e-3) +@requires_hip_kernel_arch def test_decode_weights_staged_after_loading(dist_init) -> None: """The fused-decode weight copies must not depend on the weight loaders. @@ -323,11 +367,15 @@ def test_decode_weights_staged_after_loading(dist_init) -> None: config, vllm_config, prefix="model.layers.0.self_attn" ) assert layer.decode_conv1d_weight is not None + # The AITER Gluon decode path takes o_norm.weight directly, so only the HIP + # kernel gets an fp32 norm copy. + assert (layer.decode_norm_weight is None) == layer.use_aiter_decode with torch.no_grad(): for param in layer.parameters(): param.normal_() layer.decode_conv1d_weight.fill_(float("nan")) - layer.decode_norm_weight.fill_(float("nan")) + if layer.decode_norm_weight is not None: + layer.decode_norm_weight.fill_(float("nan")) process_weights_after_loading( torch.nn.ModuleList([layer]), vllm_config.model_config, torch.device("cuda") @@ -341,7 +389,271 @@ def test_decode_weights_staged_after_loading(dist_init) -> None: torch.testing.assert_close( layer.decode_conv1d_weight, expected_conv, atol=0, rtol=0 ) + if layer.decode_norm_weight is not None: + torch.testing.assert_close( + layer.decode_norm_weight, layer.o_norm.weight.float(), atol=0, rtol=0 + ) + assert not layer.decode_norm_weight.requires_grad + assert not layer.decode_conv1d_weight.requires_grad + + +# -------------------------------------------------------------------------- +# non-spec decode: fused_recurrent_kda_packed_decode +# -------------------------------------------------------------------------- + + +def _run_aiter_packed_decode( + inp: KdaDecodeInputs, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Exactly the call KimiK3DeltaAttention._forward makes for decode batches.""" + from vllm._aiter_ops import rocm_aiter_ops + + conv_state = inp.conv_state.clone() + recurrent_state = inp.recurrent_state.clone() + num_tokens = inp.mixed_qkv.shape[0] + core_attn_out = inp.empty_out() + rocm_aiter_ops.kda_fused_recurrent_packed_decode( + mixed_qkv=inp.mixed_qkv, + g=inp.g1[0], + beta=inp.beta[0], + A_log=inp.A_log, + dt_bias=inp.dt_bias, + lower_bound=GATE_LOWER_BOUND, + initial_state=recurrent_state, + ssm_state_indices=inp.state_indices[:num_tokens], + conv_state=inp.conv_state_view(conv_state), + conv_weight=inp.decode_conv1d_weight, + out_gate=inp.g2[:num_tokens], + norm_weight=inp.norm_weight_bf16, + norm_eps=NORM_EPS, + out=core_attn_out[0, :num_tokens], + ) + return core_attn_out, conv_state, recurrent_state + + +@requires_aiter_gluon_kda +@pytest.mark.parametrize("num_heads", [12, 24]) +@pytest.mark.parametrize("num_tokens", [1, 7, 128]) +@torch.inference_mode() +def test_aiter_packed_decode_matches_triton_chain( + num_heads: int, num_tokens: int +) -> None: + inp = KdaDecodeInputs(num_tokens, num_heads, num_slots=max(num_tokens, 4) + 3) + + expected_out, expected_conv, expected_state = _run_triton_chain(inp) + actual_out, actual_conv, actual_state = _run_aiter_packed_decode(inp) + + # The fused kernel keeps the recurrent output in fp32 through the norm, + # while the Triton chain rounds it to BF16 in between. + torch.testing.assert_close(actual_out, expected_out, atol=3e-2, rtol=3e-2) + # The conv state is a pure shift-and-append: it must be bit-exact. + torch.testing.assert_close(actual_conv, expected_conv, atol=0, rtol=0) + torch.testing.assert_close(actual_state, expected_state, atol=2e-3, rtol=2e-3) + + +@requires_aiter_gluon_kda +@torch.inference_mode() +def test_aiter_packed_decode_leaves_untouched_slots_alone() -> None: + """Only the slots named by the state indices may be written.""" + num_tokens, num_heads, num_slots = 4, 12, 9 + inp = KdaDecodeInputs(num_tokens, num_heads, num_slots, seed=3) + + _, actual_conv, actual_state = _run_aiter_packed_decode(inp) + + touched = set(inp.state_indices.tolist()) + untouched = [slot for slot in range(num_slots) if slot not in touched] + assert untouched, "test needs at least one unused slot" torch.testing.assert_close( - layer.decode_norm_weight, layer.o_norm.weight.float(), atol=0, rtol=0 + actual_conv[untouched], inp.conv_state[untouched], atol=0, rtol=0 + ) + torch.testing.assert_close( + actual_state[untouched], inp.recurrent_state[untouched], atol=0, rtol=0 + ) + + +@requires_aiter_gluon_kda +@torch.inference_mode() +def test_aiter_packed_decode_skips_null_block_padding() -> None: + """A CUDA-graph decode batch is padded with NULL_BLOCK_ID (0). + + gdn_attn.py fills the tail of ``non_spec_state_indices_tensor`` with 0 and + the Triton chain leaves slot 0 alone; the fused kernel's pad_slot_guard + must do the same rather than read-modify-write slot 0 once per padded row. + """ + num_real, num_padded, num_heads = 3, 5, 12 + inp = KdaDecodeInputs(num_real + num_padded, num_heads, num_slots=12, seed=11) + inp.state_indices[num_real:] = 0 + + expected_out, expected_conv, expected_state = _run_triton_chain(inp) + actual_out, actual_conv, actual_state = _run_aiter_packed_decode(inp) + + torch.testing.assert_close( + actual_out[0, :num_real], expected_out[0, :num_real], atol=3e-2, rtol=3e-2 + ) + assert torch.isfinite(actual_out).all() + # Slot 0 must be untouched by both paths. + torch.testing.assert_close(actual_conv[0], inp.conv_state[0], atol=0, rtol=0) + torch.testing.assert_close(actual_state[0], inp.recurrent_state[0], atol=0, rtol=0) + torch.testing.assert_close(actual_conv, expected_conv, atol=0, rtol=0) + torch.testing.assert_close(actual_state, expected_state, atol=2e-3, rtol=2e-3) + + +# -------------------------------------------------------------------------- +# spec-only decode: fused_recurrent_kda with conv / gate / norm fused in +# -------------------------------------------------------------------------- + + +class SpecBatch: + """A spec-only KDA batch: ``num_spec + 1`` query tokens per request at most. + + ``seq_lens`` are the per-request query lengths, ``num_accepted`` the tokens + verified in the previous step; ``state_indices`` is the [N, num_spec + 1] + per-token slot table the GDN metadata hands the layer. + """ + + def __init__( + self, + seq_lens: list[int], + num_accepted: list[int], + num_heads: int, + num_spec: int, + seed: int = 0, + ) -> None: + assert len(seq_lens) == len(num_accepted) + assert max(seq_lens) <= num_spec + 1 + n_seqs = len(seq_lens) + slots_per_seq = num_spec + 1 + num_slots = n_seqs * slots_per_seq + 2 + self.inp = KdaDecodeInputs( + sum(seq_lens), num_heads, num_slots, seed=seed, num_spec=num_spec + ) + # Distinct slots per (request, token), shuffled, none at slot 0. + perm = torch.randperm(num_slots - 1, device="cuda")[: n_seqs * slots_per_seq] + self.state_indices = (perm + 1).to(torch.int32).view(n_seqs, slots_per_seq) + self.query_start_loc = torch.tensor( + [0] + torch.tensor(seq_lens).cumsum(0).tolist(), + dtype=torch.int32, + device="cuda", + ) + self.num_accepted_tokens = torch.tensor( + num_accepted, dtype=torch.int32, device="cuda" + ) + self.max_query_len = slots_per_seq + + +def _run_spec_triton_chain( + sb: SpecBatch, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """The layer's unfused spec path: conv update -> recurrent KDA -> norm.""" + from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( + causal_conv1d_update, + ) + from vllm.models.kimi_k3.amd.ops.third_party.kda import fused_recurrent_kda + + inp = sb.inp + conv_state = inp.conv_state.clone() + recurrent_state = inp.recurrent_state.clone() + conv_out = causal_conv1d_update( + inp.mixed_qkv, + inp.conv_state_view(conv_state), + inp.conv_weights, + None, + activation="silu", + conv_state_indices=sb.state_indices[:, 0], + num_accepted_tokens=sb.num_accepted_tokens, + query_start_loc=sb.query_start_loc, + max_query_len=sb.max_query_len, + validate_data=False, + out=torch.empty_like(inp.mixed_qkv), + ) + q, k, v = ( + x.view(1, -1, inp.num_heads, HEAD_DIM) for x in conv_out.split(inp.dim, dim=-1) + ) + core_attn_out, _ = fused_recurrent_kda( + q=q, + k=k, + v=v, + raw_g=inp.g1, + raw_beta=inp.beta, + A_log=inp.A_log, + dt_bias=inp.dt_bias, + lower_bound=GATE_LOWER_BOUND, + initial_state=recurrent_state, + cu_seqlens=sb.query_start_loc, + ssm_state_indices=sb.state_indices, + num_accepted_tokens=sb.num_accepted_tokens, + ) + out = _gated_rmsnorm(core_attn_out, inp.g2, inp.norm_weight_bf16, NORM_EPS) + return out, conv_state, recurrent_state + + +def _run_aiter_spec( + sb: SpecBatch, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Exactly the call KimiK3DeltaAttention._forward makes for spec-only batches.""" + from vllm._aiter_ops import rocm_aiter_ops + + inp = sb.inp + conv_state = inp.conv_state.clone() + recurrent_state = inp.recurrent_state.clone() + q, k, v = ( + x.view(1, -1, inp.num_heads, HEAD_DIM) + for x in inp.mixed_qkv.split(inp.dim, dim=-1) + ) + n_spec = q.shape[1] + core_attn_out = inp.empty_out() + rocm_aiter_ops.kda_fused_recurrent_spec_decode( + q=q, + k=k, + v=v, + g=inp.g1, + beta=inp.beta, + A_log=inp.A_log, + dt_bias=inp.dt_bias, + lower_bound=GATE_LOWER_BOUND, + initial_state=recurrent_state, + cu_seqlens=sb.query_start_loc, + ssm_state_indices=sb.state_indices, + num_accepted_tokens=sb.num_accepted_tokens, + conv_state=inp.conv_state_view(conv_state), + conv_weight=inp.decode_conv1d_weight, + out_gate=inp.g2[:n_spec].unsqueeze(0), + norm_weight=inp.norm_weight_bf16, + norm_eps=NORM_EPS, + out=core_attn_out[:, :n_spec], + ) + return core_attn_out, conv_state, recurrent_state + + +@requires_aiter_gluon_kda +@pytest.mark.parametrize("num_heads", [12, 24]) +@pytest.mark.parametrize( + "num_spec, seq_lens, num_accepted", + [ + (3, [4, 4, 4], [1, 2, 4]), # Kimi-K3 recipe: 3 draft tokens, uniform + (3, [4, 2, 1, 4], [4, 1, 1, 3]), # ragged after partial acceptance + (6, [7, 3], [1, 4]), # longer window, two requests + ], +) +@torch.inference_mode() +def test_aiter_spec_decode_matches_triton_chain( + num_heads: int, num_spec: int, seq_lens: list[int], num_accepted: list[int] +) -> None: + sb = SpecBatch(seq_lens, num_accepted, num_heads, num_spec, seed=5) + + expected_out, expected_conv, expected_state = _run_spec_triton_chain(sb) + actual_out, actual_conv, actual_state = _run_aiter_spec(sb) + + torch.testing.assert_close(actual_out, expected_out, atol=3e-2, rtol=3e-2) + torch.testing.assert_close(actual_conv, expected_conv, atol=0, rtol=0) + torch.testing.assert_close(actual_state, expected_state, atol=2e-3, rtol=3e-2) + # Nothing outside the slot table may change. + touched = set(sb.state_indices.flatten().tolist()) + untouched = [s for s in range(sb.inp.conv_state.shape[0]) if s not in touched] + assert untouched + torch.testing.assert_close( + actual_conv[untouched], sb.inp.conv_state[untouched], atol=0, rtol=0 + ) + torch.testing.assert_close( + actual_state[untouched], sb.inp.recurrent_state[untouched], atol=0, rtol=0 ) - assert not layer.decode_norm_weight.requires_grad diff --git a/tests/models/language/generation/test_common.py b/tests/models/language/generation/test_common.py index 450e8badce..c65f8327aa 100644 --- a/tests/models/language/generation/test_common.py +++ b/tests/models/language/generation/test_common.py @@ -75,6 +75,9 @@ AITER_MODEL_LIST = [ "facebook/opt-125m", # opt marks=[pytest.mark.core_model, pytest.mark.cpu_model], ), + pytest.param( + "facebook/opt-350m", # opt with projected token embeddings + ), pytest.param( "microsoft/phi-2", # phi marks=[pytest.mark.core_model, pytest.mark.slow_test], diff --git a/tests/models/multimodal/generation/test_vit_cudagraph.py b/tests/models/multimodal/generation/test_vit_cudagraph.py index 8c7cf1a68e..e80dd9dc85 100644 --- a/tests/models/multimodal/generation/test_vit_cudagraph.py +++ b/tests/models/multimodal/generation/test_vit_cudagraph.py @@ -70,6 +70,10 @@ def gemma3_chat_template(content: str) -> str: return f"user\n{content}\nmodel\n" +def deepseek_vl2_chat_template(content: str) -> str: + return f"<|User|>: {content}\n\n<|Assistant|>:" + + def ernie45_vl_chat_template(content: str) -> str: return ( f"<|begin_of_sentence|>User: {content}" @@ -255,7 +259,7 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = { vllm_runner_kwargs={ "load_format": "dummy", "attention_backend": ( - "ROCM_FLASHMLA_SPARSE_DSV4" + "ROCM_FLASHMLA_SPARSE_DSV41" if current_platform.is_rocm() else "FLASHMLA_SPARSE_DSV41" ), @@ -356,6 +360,15 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = { }, skip=True, # TODO: Re-enable this once OOM issues are resolved on CI. ), + "deepseek_vl2": VitCudagraphTestConfig( + model="deepseek-ai/deepseek-vl2-tiny", + modalities=["image"], + image_prompt=deepseek_vl2_chat_template("\nWhat is in this image?"), + vllm_runner_kwargs={ + "hf_overrides": {"architectures": ["DeepseekVLV2ForCausalLM"]}, + }, + marks=[pytest.mark.core_model], + ), "gemma4": VitCudagraphTestConfig( model="google/gemma-4-E2B-it", image_prompt=( diff --git a/tests/models/multimodal/pooling/test_d3.py b/tests/models/multimodal/pooling/test_d3.py index 4eb6ce4c88..cae4ec9f8f 100644 --- a/tests/models/multimodal/pooling/test_d3.py +++ b/tests/models/multimodal/pooling/test_d3.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""d3-edge against the checkpoint's own runtime (`trust_remote_code`). +"""Decision 3.0 against the checkpoint's own runtime (`trust_remote_code`). The runtime reads a question's options from the 255 code logits of its readout at the last prompt token. vLLM's classify head must give the same code logits for the same prompt tokens, and for an image prompt that vLLM processes itself. +d3-edge has an irregular layer layout; d3-nano ships its backbone as a single +weight file without an index, next to the readout file. """ import pytest @@ -14,7 +16,7 @@ from transformers import AutoModel from vllm.assets.image import ImageAsset from vllm.distributed import cleanup_dist_env_and_memory -MODEL = "vllm-sr/d3-edge" +MODELS = ["vllm-sr/d3-edge", "vllm-sr/d3-nano"] STATE = ( "The order arrived damaged yesterday. The customer has a receipt and asks " "for a replacement today." @@ -47,7 +49,7 @@ def option_probs(logits: torch.Tensor, count: int) -> torch.Tensor: return torch.softmax(logits[:count].float(), dim=-1) -@pytest.mark.parametrize("model", [MODEL]) +@pytest.mark.parametrize("model", MODELS) def test_d3_matches_the_checkpoint_runtime(vllm_runner, model: str) -> None: reference = AutoModel.from_pretrained(model, trust_remote_code=True) runtime = reference.runtime diff --git a/tests/models/qwen4_exp/test_config.py b/tests/models/qwen4_exp/test_config.py index 8f7c187b12..61d619d282 100644 --- a/tests/models/qwen4_exp/test_config.py +++ b/tests/models/qwen4_exp/test_config.py @@ -70,20 +70,21 @@ def test_moe_sp_rejects_pipeline_parallel() -> None: ), ) with ( - patch("vllm.platforms.current_platform.is_cuda", return_value=True), + patch("vllm.platforms.current_platform.is_cuda_alike", return_value=True), pytest.raises(ValueError, match="MoE SP requires PP=1"), ): Qwen4ExpForConditionalGenerationConfig.verify_and_update_config(config) +@spawn_new_process_for_each_test +@pytest.mark.parametrize("backend", ["amd", "nvidia"]) @pytest.mark.parametrize("use_sp", [False, True]) -def test_qwen4_exp_mtp_returns_sample_and_multi_streams(use_sp: bool) -> None: +def test_qwen4_exp_mtp_returns_sample_and_multi_streams( + backend: str, use_sp: bool +) -> None: """MTP returns complete, contiguous outputs after gathering padded SP shards.""" - from vllm.models.qwen4_exp.nvidia.mtp import ( - Qwen4ExpMultiTokenPredictor, - ) - - model = object.__new__(Qwen4ExpMultiTokenPredictor) + mtp_module = f"vllm.models.qwen4_exp.{backend}.mtp" + model = object.__new__(import_module(mtp_module).Qwen4ExpMultiTokenPredictor) model.use_hc_sequence_parallel = use_sp model.use_sequence_parallel = False torch.nn.Module.__init__(model) @@ -98,7 +99,9 @@ def test_qwen4_exp_mtp_returns_sample_and_multi_streams(use_sp: bool) -> None: def decoder(**kwargs): """Check that the embedding residual shares the hidden state's token slice.""" hidden = kwargs["hidden_states"] - assert kwargs["prev_block_output"].shape == (hidden.shape[0], model.hidden_size) + if backend == "nvidia": + prev_block_output = kwargs["prev_block_output"] + assert prev_block_output.shape == (hidden.shape[0], model.hidden_size) return hidden, hidden, torch.zeros(hidden.shape[0], model.hc_count) model.layers = [decoder] @@ -120,9 +123,9 @@ def test_qwen4_exp_mtp_returns_sample_and_multi_streams(use_sp: bool) -> None: return torch.cat([local, remote, torch.full_like(remote, float("nan"))]) with ( - patch("vllm.models.qwen4_exp.nvidia.mtp.get_pp_group", return_value=pp_group), - patch("vllm.models.qwen4_exp.nvidia.mtp.sp_shard", side_effect=lambda x: x[:2]), - patch("vllm.models.qwen4_exp.nvidia.mtp.sp_all_gather", side_effect=gather), + patch(f"{mtp_module}.get_pp_group", return_value=pp_group), + patch(f"{mtp_module}.sp_shard", side_effect=lambda x: x[:2]), + patch(f"{mtp_module}.sp_all_gather", side_effect=gather), ): sample_hidden, returned_multi_hidden = model.forward( input_ids=None, diff --git a/tests/models/qwen4_exp/test_tpsp.py b/tests/models/qwen4_exp/test_tpsp.py index 5182ab09c3..33e01ab24f 100644 --- a/tests/models/qwen4_exp/test_tpsp.py +++ b/tests/models/qwen4_exp/test_tpsp.py @@ -14,8 +14,8 @@ from transformers import Qwen4ExpTextConfig from vllm.platforms import current_platform -if not current_platform.is_cuda(): - pytest.skip("Qwen4Exp SP requires CUDA", allow_module_level=True) +if not current_platform.is_cuda_alike(): + pytest.skip("Qwen4Exp SP requires CUDA or ROCm", allow_module_level=True) from tests.utils import multi_gpu_test from vllm.config import ( @@ -38,14 +38,19 @@ from vllm.models.common.ops.sequence_parallel import ( sp_padding_mask, sp_shard, ) -from vllm.models.qwen4_exp.nvidia import model as qwen4_model -from vllm.models.qwen4_exp.nvidia.model import Qwen4ExpDecoderLayer from vllm.utils.network_utils import get_open_port from vllm.v1.attention.backends.short_conv_attn import PleShortConvAttentionMetadata from vllm.v1.worker.workspace import init_workspace_manager, reset_workspace_manager from .test_ple import _ConvBatchCase, _make_conv_metadata +if current_platform.is_rocm(): + from vllm.models.qwen4_exp.amd import model as qwen4_model + from vllm.models.qwen4_exp.amd.model import Qwen4ExpDecoderLayer +else: + from vllm.models.qwen4_exp.nvidia import model as qwen4_model + from vllm.models.qwen4_exp.nvidia.model import Qwen4ExpDecoderLayer + class _TPAttentionProjection(nn.Module): """Keep real TP communication without setting up an attention KV cache.""" @@ -168,7 +173,8 @@ def _check_decoder_ple(rank: int) -> None: set_current_vllm_config(config), set_forward_context({layer.ple.prefix: metadata}, config), ): - layer.ple.start_prefetch(ids, qsl, context) + if hasattr(layer.ple, "start_prefetch"): + layer.ple.start_prefetch(ids, qsl, context) output = layer( *inputs, positions, @@ -177,7 +183,10 @@ def _check_decoder_ple(rank: int) -> None: ngram_context=context, ) if config.parallel_config.enable_hc_sp: - output = tuple(sp_all_gather(x)[:num_tokens] for x in output) + # ROCm all-gather requires contiguous inputs. + output = tuple( + sp_all_gather(x.contiguous())[:num_tokens] for x in output + ) # A later MoE call may reuse workspace-backed output storage. outputs.append(tuple(x.clone() for x in output)) torch.testing.assert_close(outputs[1], outputs[0], atol=0.002, rtol=0.02) @@ -223,11 +232,38 @@ def _check_moe_sp(rank: int, config: VllmConfig) -> None: query_start_loc=None, ngram_context=None, ) - actual = tuple(sp_all_gather(x)[:num_tokens] for x in output) + actual = tuple(sp_all_gather(x.contiguous())[:num_tokens] for x in output) torch.testing.assert_close(actual, expected, atol=0.002, rtol=0.02) -def _run_sp(rank: int, port: int, moe_sp: bool) -> None: +def _check_compiled_sp_ops() -> None: + """Compiled SP collectives stay correct for token counts unseen at trace time.""" + from vllm.models.qwen4_exp.amd.ops.sequence_parallel import ( + sp_reduce_scatter, + sp_shard, + ) + + def collectives(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + gathered = sp_all_gather(sp_shard(x) * 2)[: x.shape[0]] + reduced = sp_all_gather(sp_reduce_scatter(x))[: x.shape[0]] + return gathered, reduced + + # vLLM drops Dynamo guards, so one trace must serve every token count. + compiled = torch.compile( + collectives, + fullgraph=True, + options={"guard_filter_fn": torch.compiler.skip_all_guards_unsafe}, + ) + traced = torch.randn(8, 16) + torch._dynamo.mark_dynamic(traced, 0) + compiled(traced) + for num_tokens in (7, 3, 1): + torch.manual_seed(num_tokens) + x = torch.randn(num_tokens, 16) + torch.testing.assert_close(compiled(x), collectives(x), atol=0, rtol=0) + + +def _run_sp(rank: int, port: int, moe_sp: bool, compiled_ops: bool = False) -> None: """Initialize rank-local runtime state and run one equivalence check.""" torch.accelerator.set_device_index(rank) torch.set_num_threads(1) @@ -247,6 +283,8 @@ def _run_sp(rank: int, port: int, moe_sp: bool) -> None: torch.set_default_dtype(torch.bfloat16) if moe_sp: _check_moe_sp(rank, config) + elif compiled_ops: + _check_compiled_sp_ops() else: _check_decoder_ple(rank) finally: @@ -265,3 +303,10 @@ def test_moe_sp_equivalence(monkeypatch) -> None: """Native MoE SP preserves outputs with unequal DP batches and padding.""" monkeypatch.setenv("VLLM_MOE_SKIP_PADDING", "1") mp.spawn(_run_sp, args=(get_open_port(), True), nprocs=4) + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-only SP ops") +@multi_gpu_test(num_gpus=2) +def test_compiled_sp_ops_handle_untraced_token_counts() -> None: + """Padding in compiled SP collectives must not specialize on the traced size.""" + mp.spawn(_run_sp, args=(get_open_port(), False, True), nprocs=2) diff --git a/tests/models/registry.py b/tests/models/registry.py index febd9d1a83..f9dfafb404 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -462,7 +462,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "OlmoHybridForCausalLM": _HfExamplesInfo("allenai/Olmo-Hybrid-7B"), "OlmoeForCausalLM": _HfExamplesInfo("allenai/OLMoE-1B-7B-0924-Instruct"), "OPTForCausalLM": _HfExamplesInfo( - "facebook/opt-125m", {"1b": "facebook/opt-iml-max-1.3b"} + "facebook/opt-125m", + {"350m": "facebook/opt-350m", "1b": "facebook/opt-iml-max-1.3b"}, ), "OrionForCausalLM": _HfExamplesInfo( "OrionStarAI/Orion-14B-Chat", trust_remote_code=True @@ -731,7 +732,13 @@ _SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS = { ), "D3ForSequenceClassification": _HfExamplesInfo( "vllm-sr/d3-edge", - extras={"d3": "vllm-sr/d3", "d3-mini": "vllm-sr/d3-mini"}, + extras={ + "d3": "vllm-sr/d3", + "d3-flash": "vllm-sr/d3-flash", + "d3-mini": "vllm-sr/d3-mini", + "d3-nano": "vllm-sr/d3-nano", + "d3-lite": "vllm-sr/d3-lite", + }, max_model_len=4096, ), "GPT2ForSequenceClassification": _HfExamplesInfo( diff --git a/tests/models/test_deepseek_v41_pcp.py b/tests/models/test_deepseek_v41_pcp.py new file mode 100644 index 0000000000..59b185ca26 --- /dev/null +++ b/tests/models/test_deepseek_v41_pcp.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from vllm.models.deepseek_v41.nvidia import model as model_module + + +@pytest.mark.parametrize( + "pcp_size,use_ep", [(1, False), (2, True), (2, False), (4, False)] +) +@pytest.mark.parametrize("use_sp", [False, True]) +def test_pcp_gathers_routing_ids_once_after_local_input_preparation( + pcp_size, use_ep, use_sp, monkeypatch +): + """Embedding keeps local IDs; all decoder layers receive aligned routing IDs.""" + gather = Mock(side_effect=lambda ids, dim: torch.cat([ids] * pcp_size, dim=dim)) + monkeypatch.setattr( + model_module, "get_pcp_group", lambda: SimpleNamespace(all_gather=gather) + ) + monkeypatch.setattr( + model_module, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), + ) + monkeypatch.setattr(model_module, "is_forward_context_available", lambda: False) + monkeypatch.setattr(model_module, "sp_shard", lambda tensor: tensor[::2]) + monkeypatch.setattr( + model_module, "sp_all_gather", lambda tensor: tensor.repeat_interleave(2, dim=0) + ) + monkeypatch.setattr(model_module, "hc_collapse_triton", lambda hidden, mix: hidden) + embed = Mock(side_effect=lambda ids: ids[:, None].float()) + run_layers = Mock( + side_effect=lambda layer_ids, hidden, *args: (hidden, None, None, None, hidden) + ) + model = SimpleNamespace( + parallel_config=SimpleNamespace( + prefill_context_parallel_size=pcp_size, enable_expert_parallel=use_ep + ), + use_native_mega_moe=False, + use_sequence_parallel=use_sp, + engram_hash=None, + embed_input_ids=embed, + start_layer=0, + decoder_replay_start=2, + decoder_replay_layers=None, + aux_hidden_state_layers=(), + _mtp_hidden_buffer=None, + _run_layers=run_layers, + _collapse=lambda hidden, *args: hidden, + norm=lambda hidden: hidden, + ) + for offset in (0, 10): + local_ids = torch.tensor([3, 7, 11, 15]) + offset + model_module.DeepseekV4Model.forward(model, local_ids, torch.arange(4), None) + assert embed.call_args.args[0] is local_ids + routing_ids = local_ids[::2] if use_sp else local_ids + needs_gather = pcp_size > 1 and not use_ep + if needs_gather: + torch.testing.assert_close(gather.call_args.args[0], routing_ids) + routing_ids = routing_ids.repeat(pcp_size) + torch.testing.assert_close(run_layers.call_args.args[3], routing_ids) + assert gather.call_count == 2 * int(needs_gather) diff --git a/tests/models/test_deepseek_v41_replay_batch.py b/tests/models/test_deepseek_v41_replay_batch.py index 5a3625dcbf..d34f9a2002 100644 --- a/tests/models/test_deepseek_v41_replay_batch.py +++ b/tests/models/test_deepseek_v41_replay_batch.py @@ -4,7 +4,7 @@ from dataclasses import replace from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import numpy as np import pytest @@ -16,6 +16,9 @@ from vllm.forward_context import ( ForwardContext, override_forward_context, ) +from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( + RoutedExpertsCapturer, +) from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers from vllm.models.deepseek_v41.nvidia.model_state import DeepseekV41ModelState from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor @@ -63,6 +66,7 @@ def state(monkeypatch): cfg.scheduler_config.max_num_seqs = 16 cfg.scheduler_config.max_num_batched_tokens = 1024 cfg.parallel_config.data_parallel_size = 1 + cfg.aux_output_config.enable_return_routed_experts = False cfg.compilation_config.fast_moe_cold_start = False cfg.compilation_config.cudagraph_mode = CUDAGraphMode.NONE layers = DecoderReplayLayers(WINDOW, MagicMock(), MagicMock(), set(), "swa_first") @@ -172,6 +176,31 @@ def test_replay_batch_keeps_each_request_window(state): assert context.dp_metadata is None and context.is_padding is None +def test_replay_context_restores_tp_routing_rows(state): + """Capture uses the compact replay batch's size, not the outer batch's.""" + batch = _input_batch(QUERY_LENS, SEQ_LENS, PREFILLING) + replay, _ = _prepare(state, batch, CUDAGraphMode.NONE) + context = replay.forward_context + assert context.is_padding is None and context.dp_metadata is None + count = len(REPLAY_ROWS) + assert context.batch_descriptor.num_tokens == count + capturer = RoutedExpertsCapturer.__new__(RoutedExpertsCapturer) + capturer.dp_rank, capturer.tp_size = 0, 2 + capturer.device_buffer = torch.full((401, 1, 2), -1, device=DEVICE) + expected = torch.tensor(REPLAY_ROWS, device=DEVICE)[:, None].expand(-1, 2) + gathered = torch.cat((expected, torch.full((1, 2), -1, device=DEVICE))) + with ( + override_forward_context(context), + patch( + "vllm.model_executor.layers.fused_moe.routed_experts_capturer.get_tp_group" + ) as group, + ): + group.return_value.all_gather.return_value = gathered + capturer.capture(0, gathered[: (count + 1) // 2]) + torch.testing.assert_close(capturer.device_buffer[:count, 0], expected) + assert (capturer.device_buffer[count:] == -1).all() + + def test_replay_batch_keeps_device_decode_boundaries(state): """Adaptive verification resizes decodes on the GPU (CPU [2, 2], device [1, 3]): their rows stay and the boundaries shift only by trimmed rows.""" @@ -195,6 +224,27 @@ def test_prompt_logprobs_requests_keep_their_rows(state): assert build.batch.max_query_len == 300 +@pytest.mark.parametrize("cg_mode", [CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE]) +def test_aux_output_requests_keep_every_routing_row(state, cg_mode): + """R3 needs every prefill row even when final logits only need the window.""" + state.vllm_config.aux_output_config.enable_return_routed_experts = True + for index in range(3): + state.add_request( + index, + SimpleNamespace( + replay_start=0, + sampling_params=SimpleNamespace(prompt_logprobs=None), + ), + ) + batch = _input_batch(QUERY_LENS, SEQ_LENS, PREFILLING) + state.decoder_replay_layers.trim_threshold = WINDOW + 1 + replay, _ = _prepare(state, batch, cg_mode) + if cg_mode == CUDAGraphMode.NONE: + assert replay is None + else: + torch.testing.assert_close(replay.rows, torch.arange(401, device=DEVICE)) + + def test_replay_batch_keeps_adaptive_verification_query_bound(state): """Adaptive verification bounds the decodes' device-side query lengths above their CPU lengths; the replay batch keeps that bound.""" diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index 97dd1b376b..aa6e9c976d 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -8,10 +8,8 @@ import torch from transformers import DeepseekV4Config from vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe import ( - DeepGemmSm100MegaMoEBackend, -) -from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( - bind_routed_experts_capturer, + DeepGemmMegaMoEExperts, + MegaMoEActivation, ) from vllm.model_executor.models.interfaces import ( get_mixture_of_experts_model, @@ -20,17 +18,14 @@ from vllm.model_executor.models.interfaces import ( from vllm.models.deepseek_v4.nvidia.dspark import DSparkDeepseekV4ForCausalLM from vllm.models.deepseek_v4.nvidia.model import ( DeepseekV4ForCausalLM, - DeepseekV4MegaMoEExperts, DeepseekV4MoE, make_deepseek_v4_expert_params_mapping, prepare_mega_gate_routing_metadata, ) -from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs from vllm.models.deepseek_v41.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID from vllm.models.deepseek_v41.nvidia.model import DeepseekV4MoE as DeepseekV41MoE from vllm.platforms import current_platform from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config -from vllm.utils.deep_gemm import is_deep_gemm_supported from vllm.utils.torch_utils import set_default_torch_dtype pytestmark = pytest.mark.skipif( @@ -40,11 +35,12 @@ pytestmark = pytest.mark.skipif( @pytest.fixture(autouse=True) -def sm100_mega_moe_backend(monkeypatch): +def sm100_mega_moe_layout(monkeypatch): """Build the SM100 kernel layout on any GPU; kernels are gated per test.""" monkeypatch.setattr( - "vllm.models.deepseek_v4.nvidia.model.get_deep_gemm_mega_moe_backend", - lambda *args: DeepGemmSm100MegaMoEBackend(), + "vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe." + "check_deep_gemm_mega_moe_supported", + lambda *args: None, ) @@ -125,7 +121,8 @@ def test_deepseek_v41_moe_routes_without_hash_table( ) assert moe.gate.tid2eid is None assert moe.gate.weight.shape == (num_experts, config.hidden_size) - assert isinstance(moe.experts, DeepseekV4MegaMoEExperts) + assert isinstance(moe.experts, DeepGemmMegaMoEExperts) + assert moe.experts.activation == MegaMoEActivation("swiglu", clamp=10.0) assert moe.experts.w13_weight.shape == (num_experts, 256, 128) assert moe.experts.top_k == top_k assert (config.n_routed_experts, config.num_experts_per_tok) == (384, 6) @@ -151,10 +148,9 @@ def test_deepseek_v41_moe_routes_without_hash_table( routed = {} - def check_routing(x, weights, ids, *, activation_clamp): + def check_routing(x, weights, ids, shared_experts): routed["ids"], routed["weights"] = ids, weights - assert activation_clamp == config.swiglu_limit - return x.clone() + return x + shared_experts(x) monkeypatch.setattr(moe.experts, "forward", check_routing) monkeypatch.setattr(moe.shared_experts, "forward", lambda x: 2 * x) @@ -280,10 +276,10 @@ def test_deepseek_v4_mega_gate_hash_routing_correctness( dim=-1, keepdim=True ) - def check_routing(x, weights, ids, *, activation_clamp): + def check_routing(x, weights, ids, shared_experts): torch.testing.assert_close(ids, expected_ids) torch.testing.assert_close(weights, expected_weights, rtol=1e-3, atol=1e-4) - return x.clone() + return x + shared_experts(x) monkeypatch.setattr(moe.experts, "forward", check_routing) monkeypatch.setattr(moe.shared_experts, "forward", torch.zeros_like) @@ -308,634 +304,6 @@ def test_deepseek_v4_mega_moe_expert_mapping(): ] -@pytest.mark.parametrize("use_kimi", [False, True]) -def test_deep_gemm_mega_moe_capture_precedes_eplb(monkeypatch, use_kimi): - experts_cls = DeepseekV4MegaMoEExperts - if use_kimi: - from vllm.models.kimi_k3.nvidia.model import KimiK3MegaMoEExperts - - experts_cls = KimiK3MegaMoEExperts - - experts = experts_cls.__new__(experts_cls) - torch.nn.Module.__init__(experts) - if use_kimi: - experts.synchronize_first_launch = lambda: None - experts.prefix = "model.layers.3.ffn.experts" - experts.max_num_tokens = 4 - experts.capture_fn = None - experts.get_symm_buffer = lambda: object() - experts.eplb_state = SimpleNamespace( - logical_to_physical_map=torch.empty(1), - expert_load_view=torch.empty(1), - logical_replica_count=torch.empty(1), - should_record_tensor=torch.empty(1), - num_unpadded_tokens_tensors=None, - ) - - topk_ids = torch.tensor([[1, 2], [3, 4]]) - captured: list[tuple[int, torch.Tensor]] = [] - bind_routed_experts_capturer( - SimpleNamespace(modules=lambda: [experts]), - SimpleNamespace(capture=lambda layer_id, ids: captured.append((layer_id, ids))), - ) - - class MappingReached(Exception): - pass - - def map_ids(**kwargs): - assert captured == [(3, topk_ids)] - raise MappingReached - - monkeypatch.setattr( - f"{experts_cls.__module__}.eplb_map_to_physical_and_record", - map_ids, - ) - monkeypatch.setattr( - "vllm.utils.deep_gemm._import_deep_gemm", lambda: SimpleNamespace() - ) - - with pytest.raises(MappingReached): - experts( - torch.empty(2, 8), - torch.empty(2, 2), - topk_ids, - activation_clamp=None, - ) - - -def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership(): - vllm_config = SimpleNamespace( - scheduler_config=SimpleNamespace(max_num_batched_tokens=4), - compilation_config=SimpleNamespace(static_forward_context={}), - ) - experts = DeepseekV4MegaMoEExperts( - vllm_config, - num_experts=4, - num_local_experts=2, - experts_start_idx=2, - top_k=2, - hidden_size=128, - intermediate_size=128, - ) - - nonlocal_weight = torch.ones(128, 64, dtype=torch.uint8) - assert ( - experts.weight_loader( - experts.w13_weight, - nonlocal_weight, - "experts.w13_weight", - shard_id="w1", - expert_id=1, - return_success=True, - ) - is False - ) - - w1 = torch.full((128, 64), 3, dtype=torch.uint8) - w3 = torch.full((128, 64), 7, dtype=torch.uint8) - w2 = torch.full((128, 64), 11, dtype=torch.uint8) - - assert experts.weight_loader( - experts.w13_weight, - w1, - "experts.w13_weight", - shard_id="w1", - expert_id=2, - return_success=True, - ) - assert experts.weight_loader( - experts.w13_weight, - w3, - "experts.w13_weight", - shard_id="w3", - expert_id=2, - return_success=True, - ) - assert experts.weight_loader( - experts.w2_weight, - w2, - "experts.w2_weight", - shard_id="w2", - expert_id=2, - return_success=True, - ) - - interleaved = experts.w13_weight[0].view(torch.uint8).view(16, 2, 8, 64) - assert torch.equal(interleaved[:, 0], w1.view(16, 8, 64)) - assert torch.equal(interleaved[:, 1], w3.view(16, 8, 64)) - assert torch.equal(experts.w2_weight[0].view(torch.uint8), w2) - assert torch.count_nonzero(experts.w13_weight[1]) == 0 - - -def _reference_mega_moe_layout(w13, s13, w2, s2): - """CPU copy of DeepGEMM's transform_weights_for_mega_moe on packed scales.""" - - def interleave(t): - g, n, k = t.shape - halves = t.view(g, 2, n // 16, 8, k) - return torch.stack([halves[:, 0], halves[:, 1]], dim=2).reshape(g, n, k) - - def utccp(sf): - g, mn, k = sf.shape - return sf.reshape(g, -1, 4, 32, k).transpose(2, 3).reshape(g, mn, k) - - return ( - interleave(w13), - utccp(interleave(s13.view(torch.int32))), - w2, - utccp(s2.view(torch.int32)), - ) - - -def _make_mega_moe_experts(num_local_experts, num_logical_experts): - return DeepseekV4MegaMoEExperts( - SimpleNamespace( - scheduler_config=SimpleNamespace(max_num_batched_tokens=4), - compilation_config=SimpleNamespace(static_forward_context={}), - ), - num_experts=num_local_experts, - num_local_experts=num_local_experts, - experts_start_idx=0, - top_k=1, - hidden_size=256, - intermediate_size=384, - num_logical_experts=num_logical_experts, - ) - - -def _load_mega_moe_checkpoint_shuffled(experts, num_logical_experts, seed): - hidden, inter = experts.hidden_size, experts.intermediate_size - shapes = { - "w13_weight": (num_logical_experts, 2 * inter, hidden // 2), - "w13_weight_scale": (num_logical_experts, 2 * inter, hidden // 32), - "w2_weight": (num_logical_experts, hidden, inter // 2), - "w2_weight_scale": (num_logical_experts, hidden, inter // 32), - } - generator = torch.Generator().manual_seed(seed) - ckpt = { - name: torch.randint(0, 256, shape, dtype=torch.uint8, generator=generator) - for name, shape in shapes.items() - } - shards = [] - for expert_id in range(num_logical_experts): - for name, tensor in ckpt.items(): - if name.startswith("w2_"): - shards.append((name, tensor[expert_id], "w2", expert_id)) - else: - shards.append((name, tensor[expert_id, :inter], "w1", expert_id)) - shards.append((name, tensor[expert_id, inter:], "w3", expert_id)) - for i in torch.randperm(len(shards), generator=generator).tolist(): - name, loaded_weight, shard_id, expert_id = shards[i] - if name.endswith("scale"): - loaded_weight = loaded_weight.view(torch.float8_e8m0fnu) - assert experts.weight_loader( - getattr(experts, name), - loaded_weight, - f"experts.{name}", - shard_id=shard_id, - expert_id=expert_id, - return_success=True, - ) - return ckpt - - -def _expected_mega_moe_layout(ckpt, num_local_experts, num_logical_experts): - physical = torch.arange(num_local_experts) % num_logical_experts - ckpt = {name: tensor[physical] for name, tensor in ckpt.items()} - expected = _reference_mega_moe_layout(*ckpt.values()) - if is_deep_gemm_supported(): - from vllm.utils.deep_gemm import _import_deep_gemm - - deep_gemm = _import_deep_gemm() - w13, s13, w2, s2 = (t.cuda() for t in ckpt.values()) - hidden, inter = w2.shape[1], w2.shape[2] * 2 - scales = [ - deep_gemm.transform_sf_into_required_layout( - (sf.to(torch.int32) << 23).view(torch.float32), - mn, - k, - (1, 32), - num_local_experts, - ) - for sf, mn, k in ((s13, 2 * inter, hidden), (s2, hidden, inter)) - ] - (l1, l1_sf), (l2, l2_sf) = deep_gemm.transform_weights_for_mega_moe( - (w13.view(torch.int8), scales[0]), (w2.view(torch.int8), scales[1]) - ) - for actual, reference in zip((l1, l1_sf, l2, l2_sf), expected): - assert torch.equal(actual.cpu().view(reference.dtype), reference) - return expected - - -def _assert_mega_moe_layout(experts, expected): - actual = (*experts._transformed_l1_weights, *experts._transformed_l2_weights) - for tensor, reference in zip(actual, expected): - assert torch.equal(tensor.view(reference.dtype), reference) - - -def test_deepseek_v4_mega_moe_loads_shards_into_kernel_layout(): - """Shards in any order, including EPLB replicas, land in DeepGEMM's layout.""" - experts = _make_mega_moe_experts(num_local_experts=4, num_logical_experts=3) - ckpt = _load_mega_moe_checkpoint_shuffled(experts, 3, seed=0) - _assert_mega_moe_layout(experts, _expected_mega_moe_layout(ckpt, 4, 3)) - # DeepGEMM's MN-major scale strides. - assert experts.w13_weight_scale.stride() == (768 * 2, 1, 768) - assert experts.w2_weight_scale.stride() == (256 * 3, 1, 256) - - -def test_deepseek_v4_mega_moe_reload_overwrites_kernel_layout_in_place(): - """A second load must update the live kernel tensors without reallocating.""" - experts = _make_mega_moe_experts(num_local_experts=2, num_logical_experts=2) - _load_mega_moe_checkpoint_shuffled(experts, 2, seed=0) - data_ptrs = [param.data_ptr() for param in experts.parameters()] - ckpt = _load_mega_moe_checkpoint_shuffled(experts, 2, seed=1) - assert [param.data_ptr() for param in experts.parameters()] == data_ptrs - _assert_mega_moe_layout(experts, _expected_mega_moe_layout(ckpt, 2, 2)) - - -@pytest.mark.parametrize("intermediate_size", [512, 2304]) -def test_deepseek_v4_mega_moe_preserves_checkpoint_dimensions(intermediate_size): - """Keep native widths so V4.1 routed and shared experts can fuse.""" - experts = DeepseekV4MegaMoEExperts( - SimpleNamespace( - scheduler_config=SimpleNamespace(max_num_batched_tokens=4), - compilation_config=SimpleNamespace(static_forward_context={}), - ), - num_experts=2, - num_local_experts=2, - experts_start_idx=0, - top_k=1, - hidden_size=128, - intermediate_size=intermediate_size, - ) - - assert experts.intermediate_size == intermediate_size - assert experts.w13_weight.shape == (2, 2 * intermediate_size, 64) - assert experts.w13_weight_scale.shape == (2, 2 * intermediate_size, 1) - assert experts.w2_weight.shape == (2, 128, intermediate_size // 2) - assert experts.w2_weight_scale.shape == (2, 128, intermediate_size // 128) - transformed = (*experts._transformed_l1_weights, *experts._transformed_l2_weights) - for actual, param in zip(transformed, experts.parameters()): - assert actual.data_ptr() == param.data_ptr() - - -def test_deepseek_v4_mega_moe_loads_shared_experts_into_kernel_layout(monkeypatch): - """Shared expert shards land directly in the MegaMoE kernel layout.""" - monkeypatch.setattr( - DeepGemmSm100MegaMoEBackend, "supports_shared_experts", lambda *args: True - ) - experts = DeepseekV4MegaMoEExperts( - SimpleNamespace( - scheduler_config=SimpleNamespace(max_num_batched_tokens=4), - compilation_config=SimpleNamespace(static_forward_context={}), - ), - num_experts=2, - num_local_experts=2, - experts_start_idx=0, - top_k=1, - hidden_size=256, - intermediate_size=384, - num_shared_experts=1, - ) - assert experts.has_fused_shared_experts - hidden, inter = 256, 384 - generator = torch.Generator().manual_seed(42) - - def rand(*shape): - return torch.randint(0, 256, shape, dtype=torch.uint8, generator=generator) - - w1, w3, w2 = rand(inter, hidden), rand(inter, hidden), rand(hidden, inter) - s1, s3, s2 = ( - rand(inter, hidden // 32), - rand(inter, hidden // 32), - rand(hidden, inter // 32), - ) - fp8, e8m0 = torch.float8_e4m3fn, torch.float8_e8m0fnu - # The down shard arrives as `.w2.` (MTP, DSpark) or `.down_proj.` (mapped). - for name, data in [ - ("shared_experts.w1.weight", w1.view(fp8)), - ("shared_experts.w3.weight", w3.view(fp8)), - ("shared_experts.down_proj.weight", w2.view(fp8)), - ("shared_experts.w1.weight_scale", s1.view(e8m0)), - ("shared_experts.w3.weight_scale", s3.view(e8m0)), - ("shared_experts.w2.weight_scale", s2.view(e8m0)), - ]: - loaded = experts.load_shared_expert_weight(f"ffn.{name}", data) - # Weight tracking checks this name against the registered parameters. - assert loaded in {f"ffn.experts.{n}" for n, _ in experts.named_parameters()} - - def l1_halves(): - return experts.shared_l1_weight.view(torch.uint8).view(-1, 2, 8, hidden) - - # Every tensor matches DeepGEMM's batched transform of the checkpoint. - expected = _reference_mega_moe_layout( - *(t.unsqueeze(0) for t in (torch.cat([w1, w3]), torch.cat([s1, s3]), w2, s2)) - ) - actual = ( - experts.shared_l1_weight, - experts.shared_l1_weight_scale, - experts.shared_l2_weight, - experts.shared_l2_weight_scale, - ) - for tensor, reference in zip(actual, expected): - assert torch.equal(tensor.view(reference.dtype), reference[0]) - # A reload of one shard overwrites in place and leaves the other half alone. - ptr = experts.shared_l1_weight.data_ptr() - w1b = rand(inter, hidden) - experts.load_shared_expert_weight("shared_experts.w1.weight", w1b.view(fp8)) - assert experts.shared_l1_weight.data_ptr() == ptr - assert torch.equal(l1_halves()[:, 0], w1b.view(-1, 8, hidden)) - assert torch.equal(l1_halves()[:, 1], w3.view(-1, 8, hidden)) - with pytest.raises(ValueError): - experts.load_shared_expert_weight("shared_experts.gate_up_proj.weight", w1) - - -def _fake_deep_gemm(with_shared: bool): - def old_buffer(group): ... - - def old_kernel(y): ... - - def new_buffer(group, num_shared_experts=0): ... - - def new_kernel(y, shared_l1_weights=None, shared_l2_weights=None): ... - - return SimpleNamespace( - get_symm_buffer_for_mega_moe=new_buffer if with_shared else old_buffer, - fp8_fp4_mega_moe=new_kernel if with_shared else old_kernel, - get_block_m_for_mega_moe=lambda *args: 0, - ) - - -@pytest.mark.parametrize( - "weight_dtype,scale_dtype,scale_shape,with_shared,fused", - [ - (torch.float8_e4m3fn, torch.float8_e8m0fnu, (2, 2), True, True), - (torch.float8_e4m3fn, torch.uint8, (256, 8), True, True), - (torch.float8_e4m3fn, torch.float8_e8m0fnu, (2, 2), False, False), - (torch.bfloat16, None, None, True, False), - (torch.float8_e4m3fn, torch.float32, (2, 2), True, False), - (torch.float8_e4m3fn, torch.uint8, (256, 16), True, False), - ], -) -def test_deepseek_v4_mega_moe_decides_shared_fusion_at_construction( - monkeypatch, weight_dtype, scale_dtype, scale_shape, with_shared, fused -): - """Fuse only shared experts the installed kernel and the 1x32 layout fit.""" - monkeypatch.setattr( - "vllm.utils.deep_gemm._import_deep_gemm", lambda: _fake_deep_gemm(with_shared) - ) - - def linear(): - scale = ( - None if scale_dtype is None else torch.empty(scale_shape).to(scale_dtype) - ) - return SimpleNamespace( - weight=torch.empty(256, 256, dtype=weight_dtype), weight_scale_inv=scale - ) - - shared = SimpleNamespace(gate_up_proj=linear(), down_proj=linear()) - experts = DeepseekV4MegaMoEExperts( - SimpleNamespace( - scheduler_config=SimpleNamespace(max_num_batched_tokens=4), - compilation_config=SimpleNamespace(static_forward_context={}), - ), - num_experts=2, - num_local_experts=2, - experts_start_idx=0, - top_k=1, - hidden_size=256, - intermediate_size=128, - num_shared_experts=1, - shared_experts=shared, - ) - - assert experts.has_fused_shared_experts == fused - assert hasattr(experts, "shared_l1_weight") == fused - - -@pytest.mark.parametrize("fused", [False, True]) -def test_deepseek_v4_mega_moe_does_not_double_add_fused_shared_expert( - monkeypatch, fused -): - class FakeGate(torch.nn.Module): - weight = torch.empty(2, 128) - tid2eid = None - e_score_correction_bias = None - - def forward(self, hidden_states): - return torch.empty(hidden_states.shape[0], 2), None - - class FakeExperts(torch.nn.Module): - has_fused_shared_experts = fused - backend = DeepGemmSm100MegaMoEBackend() - - def forward(self, hidden_states, *args, **kwargs): - return torch.ones_like(hidden_states) - - class FakeSharedExperts(torch.nn.Module): - calls = 0 - - def forward(self, hidden_states): - self.calls += 1 - return torch.full_like(hidden_states, 2) - - moe = DeepseekV4MoE.__new__(DeepseekV4MoE) - torch.nn.Module.__init__(moe) - moe.use_native_mega_moe = True - moe.gate = FakeGate() - moe.experts = FakeExperts() - moe.shared_experts = None if fused else FakeSharedExperts() - moe.scoring_func = "sqrtsoftplus" - moe.n_activated_experts = 1 - moe.renormalize = True - moe.hash_indices_dtype = torch.int64 - moe.routed_scaling_factor = 1.0 - moe.swiglu_limit = 10.0 - monkeypatch.setattr( - "vllm.models.deepseek_v4.nvidia.model.fused_topk_bias", - lambda **kwargs: ( - torch.ones(kwargs["hidden_states"].shape[0], 1), - torch.zeros(kwargs["hidden_states"].shape[0], 1, dtype=torch.int64), - ), - ) - - output = moe(torch.zeros(2, 128)) - - expected = 1 if fused else 3 - assert torch.all(output == expected) - if not fused: - assert moe.shared_experts.calls == 1 - - -@pytest.mark.skipif( - not torch.cuda.is_available(), - reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.", -) -@pytest.mark.parametrize("num_tokens", [7, 63, 64, 65, 127, 128, 135, 434, 16384]) -@pytest.mark.parametrize("hidden_size", [256, 5120]) -def test_deepseek_v4_mega_moe_fused_input_staging_is_bitwise_exact( - num_tokens, hidden_size -): - from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8 - - device = torch.device("cuda") - top_k = 8 - - generator = torch.Generator(device=device) - generator.manual_seed(0) - hidden_states = ( - torch.randn( - num_tokens, - hidden_size, - device=device, - dtype=torch.float32, - generator=generator, - ) - * 17.0 - ).to(torch.bfloat16) - hidden_states[0, :32] = 0 - hidden_states[1, 32:64] = 1.0e-6 - hidden_states[2, 64:96] = -1.0e-6 - - topk_ids = torch.randint( - 0, - 256, - (num_tokens, top_k), - device=device, - dtype=torch.int32, - generator=generator, - ) - topk_weights = torch.randn( - num_tokens, - top_k, - device=device, - dtype=torch.float32, - generator=generator, - ) - - ref_x, ref_x_sf = per_token_cast_to_fp8( - hidden_states, - use_ue8m0=True, - gran_k=32, - use_packed_ue8m0=True, - ) - ref_topk_idx = topk_ids.to(torch.int64) - ref_topk_weights = topk_weights.clone() - - fused_x = torch.empty_like(ref_x) - fused_x_sf = torch.empty_like(ref_x_sf) - fused_topk_idx = torch.empty_like(ref_topk_idx) - fused_topk_weights = torch.empty_like(ref_topk_weights) - - prepare_megamoe_inputs( - hidden_states, - topk_weights, - topk_ids, - fused_x, - fused_x_sf, - fused_topk_idx, - fused_topk_weights, - ) - torch.accelerator.synchronize() - - assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8)) - assert torch.equal(fused_x_sf, ref_x_sf) - assert torch.equal(fused_topk_idx, ref_topk_idx) - assert torch.equal( - fused_topk_weights.view(torch.uint8), - ref_topk_weights.view(torch.uint8), - ) - - -@pytest.mark.parametrize("shared_block_m", [8, 32, 96, 128, 192]) -def test_deepseek_v4_mega_moe_stages_shared_scale_tma_layout(shared_block_m): - from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8 - - device = torch.device("cuda") - num_tokens = shared_block_m + 7 - hidden_size = 256 - top_k = 8 - generator = torch.Generator(device=device) - generator.manual_seed(shared_block_m) - hidden_states = torch.randn( - num_tokens, - hidden_size, - device=device, - dtype=torch.bfloat16, - generator=generator, - ) - topk_ids = torch.randint( - 0, - 256, - (num_tokens, top_k), - device=device, - dtype=torch.int32, - generator=generator, - ) - topk_weights = torch.randn( - num_tokens, - top_k, - device=device, - dtype=torch.float32, - generator=generator, - ) - - ref_x, ref_x_sf = per_token_cast_to_fp8( - hidden_states, - use_ue8m0=True, - gran_k=32, - use_packed_ue8m0=True, - ) - aligned_block_m = ((shared_block_m + 127) // 128) * 128 - num_shared_rows = ((num_tokens + shared_block_m - 1) // shared_block_m) * ( - aligned_block_m - ) - ref_shared_x_sf = torch.zeros( - num_shared_rows, - hidden_size // 128, - dtype=torch.int32, - device=device, - ) - for token_id in range(num_tokens): - m_in_block = token_id % shared_block_m - transposed_m = ( - (m_in_block // 128) * 128 + (m_in_block % 32) * 4 + (m_in_block % 128) // 32 - ) - shared_row = token_id // shared_block_m * aligned_block_m + transposed_m - ref_shared_x_sf[shared_row].copy_(ref_x_sf[token_id]) - - fused_x = torch.empty_like(ref_x) - fused_x_sf = torch.empty_like(ref_x_sf) - fused_shared_storage = torch.full( - (hidden_size // 128, num_shared_rows), - -1, - dtype=torch.int32, - device=device, - ) - fused_shared_x_sf = fused_shared_storage.t() - fused_topk_idx = torch.empty_like(topk_ids, dtype=torch.int64) - fused_topk_weights = torch.empty_like(topk_weights) - - prepare_megamoe_inputs( - hidden_states, - topk_weights, - topk_ids, - fused_x, - fused_x_sf, - fused_topk_idx, - fused_topk_weights, - shared_x_sf=fused_shared_x_sf, - shared_block_m=shared_block_m, - ) - torch.accelerator.synchronize() - - populated = ref_shared_x_sf != 0 - assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8)) - assert torch.equal(fused_x_sf, ref_x_sf) - assert torch.equal(fused_shared_x_sf[populated], ref_shared_x_sf[populated]) - assert torch.all(fused_shared_x_sf[~populated] == -1) - - def test_deepseek_v4_pwal_hook_finalizes_mhc_broadcast(): """The loader invokes the model-level PWAL hook for every load format, so it must finalize mhc broadcast weights to cover dummy load, which @@ -1040,85 +408,3 @@ def test_dspark_draft_registers_mixture_of_experts( assert len(draft.moe_layers) == draft.model.num_dspark_layers == 2 assert draft.num_routed_experts == hf_config.n_routed_experts assert draft.num_redundant_experts == 4 - - -@pytest.mark.skipif( - not torch.cuda.is_available(), - reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.", -) -@pytest.mark.parametrize("nonfinite_padding", [False, True]) -@pytest.mark.parametrize("num_tokens", [7, 65, 135]) -def test_deepseek_v4_mega_moe_fused_input_staging_masks_padding( - nonfinite_padding, num_tokens -): - from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8 - - device = torch.device("cuda") - hidden_size = 256 - top_k = 8 - - generator = torch.Generator(device=device) - generator.manual_seed(1) - hidden_states = torch.randn( - num_tokens, - hidden_size, - device=device, - dtype=torch.bfloat16, - generator=generator, - ) - topk_ids = torch.randint( - 0, - 256, - (num_tokens, top_k), - device=device, - dtype=torch.int32, - generator=generator, - ) - topk_weights = torch.randn( - num_tokens, - top_k, - device=device, - dtype=torch.float32, - generator=generator, - ) - is_padding = torch.zeros(num_tokens, device=device, dtype=torch.bool) - is_padding[1::3] = True - is_padding[-1] = True - if nonfinite_padding: - topk_weights[is_padding] = float("nan") - - ref_x, ref_x_sf = per_token_cast_to_fp8( - hidden_states, - use_ue8m0=True, - gran_k=32, - use_packed_ue8m0=True, - ) - ref_topk_idx = topk_ids.to(torch.int64) - ref_topk_idx[is_padding] = -1 - ref_topk_weights = topk_weights.clone() - ref_topk_weights[is_padding] = 0.0 - - fused_x = torch.empty_like(ref_x) - fused_x_sf = torch.empty_like(ref_x_sf) - fused_topk_idx = torch.empty_like(ref_topk_idx) - fused_topk_weights = torch.empty_like(ref_topk_weights) - - prepare_megamoe_inputs( - hidden_states, - topk_weights, - topk_ids, - fused_x, - fused_x_sf, - fused_topk_idx, - fused_topk_weights, - is_padding=is_padding, - ) - torch.accelerator.synchronize() - - assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8)) - assert torch.equal(fused_x_sf, ref_x_sf) - assert torch.equal(fused_topk_idx, ref_topk_idx) - assert torch.equal( - fused_topk_weights.view(torch.uint8), - ref_topk_weights.view(torch.uint8), - ) diff --git a/tests/models/test_registry.py b/tests/models/test_registry.py index 37a95d9f7f..354e7fd689 100644 --- a/tests/models/test_registry.py +++ b/tests/models/test_registry.py @@ -127,9 +127,9 @@ def test_registry_is_pp(model_arch, is_pp): @pytest.mark.parametrize( "model_arch,supported", [ - ("Qwen4ExpForCausalLM", not current_platform.is_rocm()), - ("Qwen4ExpForConditionalGeneration", not current_platform.is_rocm()), - ("Qwen4ExpMTP", not current_platform.is_rocm()), + ("Qwen4ExpForCausalLM", True), + ("Qwen4ExpForConditionalGeneration", True), + ("Qwen4ExpMTP", True), ("LlamaForCausalLM", False), ], ) diff --git a/tests/parser/engine/test_parser_engine.py b/tests/parser/engine/test_parser_engine.py index 811f9cb7b8..ac0203bc0b 100644 --- a/tests/parser/engine/test_parser_engine.py +++ b/tests/parser/engine/test_parser_engine.py @@ -161,12 +161,15 @@ class TestReasoningEndTokenIds: ) engine = _make_engine(cfg) assert engine.reasoning_end_token_ids == {201, 202} + assert engine.implicit_reasoning_end_token_ids == {202} assert engine.find_reasoning_end([5, 201]) == ReasoningEnd(1, False) assert engine.find_reasoning_end([5, 202]) == ReasoningEnd(1, True) def test_transition_staying_in_reasoning_is_ignored(self): cfg = _with_reasoning_exits(("THINK_START", ParserState.REASONING, ())) - assert _make_engine(cfg).reasoning_end_token_ids == {201} + engine = _make_engine(cfg) + assert engine.reasoning_end_token_ids == {201} + assert engine.implicit_reasoning_end_token_ids == frozenset() def test_unreported_exit_fails_closed(self): cfg = _with_reasoning_exits( diff --git a/tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/__init__.py b/tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/__init__.py index 457187e4b4..74ffac1f24 100644 --- a/tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/__init__.py +++ b/tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/__init__.py @@ -6,11 +6,16 @@ from vllm import ModelRegistry def register(): # Test directly passing the model - from .my_opt import MyOPTForCausalLM + from .my_opt import MyOPTForCausalLM, MyOPTRawInputTokensForCausalLM if "MyOPTForCausalLM" not in ModelRegistry.get_supported_archs(): ModelRegistry.register_model("MyOPTForCausalLM", MyOPTForCausalLM) + if "MyOPTRawInputTokensForCausalLM" not in ModelRegistry.get_supported_archs(): + ModelRegistry.register_model( + "MyOPTRawInputTokensForCausalLM", MyOPTRawInputTokensForCausalLM + ) + # Test passing lazy model if "MyGemma2Embedding" not in ModelRegistry.get_supported_archs(): ModelRegistry.register_model( diff --git a/tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/my_opt.py b/tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/my_opt.py index f1e6e7b10f..da3ad64253 100644 --- a/tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/my_opt.py +++ b/tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/my_opt.py @@ -5,6 +5,7 @@ import torch from vllm.model_executor.models.opt import OPTForCausalLM +from vllm.sequence import IntermediateTensors class MyOPTForCausalLM(OPTForCausalLM): @@ -15,3 +16,21 @@ class MyOPTForCausalLM(OPTForCausalLM): logits.zero_() logits[:, 0] += 1.0 return logits + + +class MyOPTRawInputTokensForCausalLM(MyOPTForCausalLM): + # Like models whose decoder layers read token ids, this model needs + # input_ids on every pipeline-parallel stage, not only the first. + requires_raw_input_tokens = True + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + assert input_ids is not None, "requires_raw_input_tokens model got no input_ids" + return super().forward( + input_ids, positions, intermediate_tensors, inputs_embeds + ) diff --git a/tests/plugins_tests/test_oot_registration_online.py b/tests/plugins_tests/test_oot_registration_online.py index 151373d82f..1c0880bf1d 100644 --- a/tests/plugins_tests/test_oot_registration_online.py +++ b/tests/plugins_tests/test_oot_registration_online.py @@ -7,7 +7,7 @@ chatml_jinja_path = VLLM_PATH / "examples/template_chatml.jinja" assert chatml_jinja_path.exists() -def run_and_test_dummy_opt_api_server(model, tp=1): +def run_and_test_dummy_opt_api_server(model, tp=1, extra_args=(), env_dict=None): # the model is registered through the plugin server_args = [ "--gpu-memory-utilization", @@ -20,8 +20,9 @@ def run_and_test_dummy_opt_api_server(model, tp=1): "dummy", "-tp", f"{tp}", + *extra_args, ] - with RemoteOpenAIServer(model, server_args) as server: + with RemoteOpenAIServer(model, server_args, env_dict=env_dict) as server: client = server.get_client() completion = client.chat.completions.create( model=model, diff --git a/tests/quantization/test_humming_compile.py b/tests/quantization/test_humming_compile.py new file mode 100644 index 0000000000..d69973b68a --- /dev/null +++ b/tests/quantization/test_humming_compile.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.utils.import_utils import has_humming + +pytestmark = pytest.mark.skipif( + not (has_humming() and current_platform.is_cuda()), + reason="requires humming and CUDA", +) + + +def test_humming_input_quant_compiles_fullgraph(): + """Humming's forward helpers look up ``DeviceInfo`` on every call; vLLM + traces them with ``fullgraph=True`` (vllm-project/humming#147).""" + import vllm.utils.humming as hm + + weight_schema = hm.HummingWeightSchema(b_dtype="int4", weight_scale_group_size=128) + input_schema = hm.HummingInputSchema(a_dtype="int8") + from humming.transform import prepare_layer_config + + config = prepare_layer_config( + shape_n=256, + shape_k=256, + weight_schema=weight_schema, + input_schema=input_schema, + torch_dtype=torch.bfloat16, + device=torch.device("cuda"), + ) + x = torch.randn(16, 256, dtype=torch.bfloat16, device="cuda") + + def quant(t: torch.Tensor): + return hm.may_process_input(config, t) + + expected = quant(x) + actual = torch.compile(quant, fullgraph=True)(x) + for a, e in zip(actual, expected): + if e is None: + assert a is None + else: + torch.testing.assert_close(a, e) diff --git a/tests/renderers/chat_utils/test_prompt_embeds.py b/tests/renderers/chat_utils/test_prompt_embeds.py index d6f1673ca7..c036a950b5 100644 --- a/tests/renderers/chat_utils/test_prompt_embeds.py +++ b/tests/renderers/chat_utils/test_prompt_embeds.py @@ -105,6 +105,7 @@ def _make_mock_model_config(*, enable_prompt_embeds: bool = True) -> mock.MagicM # `safe_load_prompt_embeds` pins each tensor to the model's hidden_size # and dtype, so the mock must return concrete values. mc.get_hidden_size.return_value = _MOCK_HIDDEN_SIZE + mc.hf_text_config = object() mc.dtype = _MOCK_DTYPE return mc diff --git a/tests/renderers/test_completions.py b/tests/renderers/test_completions.py index 4fac063807..9d37ff78d1 100644 --- a/tests/renderers/test_completions.py +++ b/tests/renderers/test_completions.py @@ -34,6 +34,7 @@ class MockModelConfig: tokenizer_revision = None tokenizer_mode = "auto" hf_config = MockHFConfig() + hf_text_config = hf_config encoder_config: dict[str, Any] | None = None enable_prompt_embeds: bool = True skip_tokenizer_init: bool = False diff --git a/tests/renderers/test_prompt_embeds_payload_validation.py b/tests/renderers/test_prompt_embeds_payload_validation.py index bec28c4c49..d8342f52ee 100644 --- a/tests/renderers/test_prompt_embeds_payload_validation.py +++ b/tests/renderers/test_prompt_embeds_payload_validation.py @@ -65,6 +65,40 @@ def _saved_tensor(tensor: torch.Tensor) -> bytes: return buffer.getvalue() +@pytest.mark.parametrize( + "model, hidden_size, input_size, buffer_size", + [ + ("facebook/opt-350m", 1024, 512, 512), + ("facebook/opt-125m", 768, 768, 768), + ("openai/clip-vit-base-patch32", 512, 512, 1024), + ("google/siglip-base-patch16-224", 768, 768, 1024), + ], +) +def test_prompt_embeds_input_width(model, hidden_size, input_size, buffer_size): + """OPT uses its token width; dual encoders retain their text input width.""" + from vllm.config import ModelConfig + + config = ModelConfig( + model=model, + dtype="float32", + enable_prompt_embeds=True, + ) + if "opt-" not in model: + config.hf_text_config.projection_dim = buffer_size + assert config.get_hidden_size() == hidden_size + assert config.get_inputs_embeds_size() == buffer_size + + tensor = torch.randn(4, input_size) + loaded = safe_load_prompt_embeds(config, _encode(_saved_tensor(tensor))) + torch.testing.assert_close(loaded, tensor) + + wrong_size = hidden_size if hidden_size != input_size else input_size + 1 + with pytest.raises(VLLMValidationError, match="does not match"): + safe_load_prompt_embeds( + config, _encode(_saved_tensor(torch.randn(4, wrong_size))) + ) + + UNPARSABLE_PAYLOADS = [ pytest.param(b"", id="empty"), pytest.param(b"\x80", id="single-byte"), diff --git a/tests/test_config.py b/tests/test_config.py index a44aa03e82..6151124610 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -2132,8 +2132,8 @@ def test_draft_parallel_config_preserves_ep_without_model( [ pytest.param("eagle3", False, 0, id="eagle3"), pytest.param("eagle3", True, 7, id="p-eagle"), - pytest.param("dflash", True, 8, id="dflash"), - pytest.param("dspark", True, 7, id="dspark"), + pytest.param("dflash", True, 0, id="dflash"), + pytest.param("dspark", True, 0, id="dspark"), pytest.param("mtp", False, 0, id="mtp"), pytest.param("ngram", False, 0, id="ngram"), pytest.param("draft_model", False, 1, id="draft-model"), @@ -3769,34 +3769,48 @@ def test_dual_key_gumbel_requires_probabilistic_drafting(): config._check_supports_watermarking() -@pytest.mark.parametrize("method", ["eagle", "eagle3", "mtp"]) -def test_dual_key_gumbel_supports_probabilistic_speculative_decoding(method): +@pytest.mark.parametrize( + ("method", "parallel_drafting"), + [ + ("eagle", False), + ("eagle3", False), + ("mtp", False), + ("dspark", True), + ("dflash", True), + ], +) +def test_dual_key_gumbel_supports_probabilistic_speculative_decoding( + method, parallel_drafting +): config = _watermarked_vllm_config() config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42) config.speculative_config = SimpleNamespace( method=method, draft_sample_method="probabilistic", rejection_sample_method="standard", - parallel_drafting=False, + parallel_drafting=parallel_drafting, ) config._check_supports_watermarking() -def test_dual_key_gumbel_supports_dspark(): +@pytest.mark.parametrize("architecture", ["DFlash2DraftModel", "LiLiCorrDraftModel"]) +def test_dual_key_gumbel_rejects_dflash_candidate_drafts(architecture): config = _watermarked_vllm_config() config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42) config.speculative_config = SimpleNamespace( - method="dspark", + method="dflash", draft_sample_method="probabilistic", rejection_sample_method="standard", parallel_drafting=True, + draft_model_config=SimpleNamespace(architectures=[architecture]), ) - config._check_supports_watermarking() + with pytest.raises(ValueError, match="candidate-head drafters"): + config._check_supports_watermarking() -def test_dual_key_gumbel_rejects_non_autoregressive_speculation(): +def test_dual_key_gumbel_rejects_unsupported_speculative_method(): config = _watermarked_vllm_config() config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42) config.speculative_config = SimpleNamespace( @@ -3806,7 +3820,7 @@ def test_dual_key_gumbel_rejects_non_autoregressive_speculation(): parallel_drafting=False, ) - with pytest.raises(ValueError, match="autoregressive model-based"): + with pytest.raises(ValueError, match="Watermarking supports only"): config._check_supports_watermarking() @@ -4167,7 +4181,7 @@ def test_revision_resolved_for_model(mock_resolve): @pytest.mark.parametrize( - "enabled,supported,is_cuda,tp_size,pp_size,error", + "enabled,supported,cuda_alike,tp_size,pp_size,error", [ (True, True, True, 2, 1, None), (True, False, True, 2, 1, "not supported for this model"), @@ -4178,7 +4192,7 @@ def test_revision_resolved_for_model(mock_resolve): (False, False, False, 2, 1, None), ], ) -def test_hc_sp_model_support(enabled, supported, is_cuda, tp_size, pp_size, error): +def test_hc_sp_model_support(enabled, supported, cuda_alike, tp_size, pp_size, error): """Validate active HC SP configurations and ignore the flag with TP=1.""" from vllm.model_executor.models.registry import ModelRegistry @@ -4196,7 +4210,10 @@ def test_hc_sp_model_support(enabled, supported, is_cuda, tp_size, pp_size, erro distributed_executor_backend="mp", ) with ( - patch("vllm.config.model.current_platform.is_cuda", return_value=is_cuda), + patch( + "vllm.config.model.current_platform.is_cuda_alike", + return_value=cuda_alike, + ), patch.object(ModelRegistry, "is_hc_sp_supported_model", return_value=supported), patch.object(ModelRegistry, "is_pp_supported_model", return_value=True), ): diff --git a/tests/tokenizers_/test_basic.py b/tests/tokenizers_/test_basic.py index a911245494..922820b89c 100644 --- a/tests/tokenizers_/test_basic.py +++ b/tests/tokenizers_/test_basic.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from pathlib import Path +from types import SimpleNamespace from typing import _get_protocol_attrs # type: ignore from unittest.mock import patch @@ -15,7 +16,11 @@ from transformers import ( from vllm.tokenizers import TokenizerLike, get_tokenizer from vllm.tokenizers.hf import HfTokenizer from vllm.tokenizers.mistral import MistralTokenizer -from vllm.tokenizers.registry import resolve_tokenizer_args +from vllm.tokenizers.registry import ( + cached_tokenizer_from_config, + resolve_tokenizer_args, + tokenizer_args_from_config, +) def _get_missing_attrs(obj: object, target: type): @@ -65,7 +70,11 @@ def test_tokenizer_revision(tokenizer_name: str): @pytest.mark.parametrize( "code_revision, error_match", [ - pytest.param("never", "not a valid git identifier", id="unknown-revision"), + pytest.param( + "never", + "not a valid git identifier|couldn't find them in the cached files", + id="unknown-revision", + ), pytest.param( "73293ac504e43f4a895688b57c65cebe24f01a38", r"tokenization_fast\.py", @@ -143,6 +152,141 @@ def test_tokenizer_code_revision_default_or_matching(code_revision: str | None): _assert_tokenizer_like(tokenizer) +@pytest.mark.skip_global_cleanup +@pytest.mark.parametrize( + "code_revision", + [ + None, + "c8ae8a2e5845fc3ad4e5237e4346e5a9ba1e87d2", + "2efddce40dddaccd37bae208c3c7ca66dbedf68a", + ], +) +def test_tokenizer_config_code_revision(code_revision): + """Load real remote config code at its pin, not the data revision.""" + revision = "c8ae8a2e5845fc3ad4e5237e4346e5a9ba1e87d2" + + with patch( + "vllm.tokenizers.registry.CachedHfTokenizer.from_pretrained", + return_value=SimpleNamespace(is_fast=True), + ) as load_tokenizer: + get_tokenizer( + "hf-internal-testing/test_dynamic_model", + tokenizer_mode="hf", + trust_remote_code=True, + revision=revision, + code_revision=code_revision, + ) + + config = load_tokenizer.call_args.kwargs["config"] + assert (code_revision or revision) in type(config).__module__.split(".") + assert config.new_hidden_size == 12 + + +@pytest.mark.skip_global_cleanup +@pytest.mark.parametrize( + "loader", (cached_tokenizer_from_config, tokenizer_args_from_config) +) +@pytest.mark.parametrize( + "model, code_revision, kwargs, error_match", + [ + pytest.param( + None, + "never", + {}, + "not a valid git identifier|couldn't find them in the cached files", + id="same-repo", + ), + pytest.param( + None, + "73293ac504e43f4a895688b57c65cebe24f01a38", + {}, + r"tokenization_fast\.py", + id="same-repo-missing-code", + ), + pytest.param( + "facebook/opt-125m", + "27dcfa74d334bc871f3234de431e71c6eeba5dd6", + {}, + None, + id="separate-repo", + ), + pytest.param( + None, + "never", + {"code_revision": "08a35780b7a78794436d8654aafa26ee0165c149"}, + None, + id="explicit-revision", + ), + pytest.param(None, "never", {"code_revision": None}, None, id="explicit-none"), + ], +) +def test_config_tokenizer_code_revision( + loader, model, code_revision, kwargs, error_match +): + """Only same-repo tokenizers inherit model code revisions; overrides win.""" + revision = "08a35780b7a78794436d8654aafa26ee0165c149" + repo = "hf-internal-testing/test_dynamic_tokenizer" + config = SimpleNamespace( + model=model or repo, + tokenizer=repo, + runner_type="generate", + tokenizer_mode="hf", + tokenizer_revision=revision, + code_revision=code_revision, + trust_remote_code=True, + skip_tokenizer_init=False, + ) + + def load(): + result = loader(config, **kwargs) + if loader is tokenizer_args_from_config: + mode, name, args, resolved_kwargs = result + return get_tokenizer(name, *args, tokenizer_mode=mode, **resolved_kwargs) + return result + + if error_match is not None: + with pytest.raises(OSError, match=error_match): + load() + else: + tokenizer = load() + reference = AutoTokenizer.from_pretrained( + repo, + revision=revision, + code_revision=kwargs.get("code_revision"), + trust_remote_code=True, + ) + assert tokenizer.encode("hello world") == reference.encode("hello world") + _assert_tokenizer_like(tokenizer) + + +@pytest.mark.parametrize( + "loader", (cached_tokenizer_from_config, tokenizer_args_from_config) +) +@pytest.mark.parametrize("tokenizer_mode", ["mistral", "auto"]) +@pytest.mark.parametrize("code_revision", [None, "never"]) +def test_config_mistral_code_revision(loader, tokenizer_mode, code_revision): + """Config loading must not pass HF-only revisions to the Mistral backend.""" + repo = "mistralai/Mistral-7B-Instruct-v0.3" + config = SimpleNamespace( + model=repo, + tokenizer=repo, + runner_type="generate", + tokenizer_mode=tokenizer_mode, + tokenizer_revision=None, + code_revision=code_revision, + trust_remote_code=False, + skip_tokenizer_init=False, + ) + tokenizer = loader(config) + if loader is tokenizer_args_from_config: + mode, name, args, kwargs = tokenizer + tokenizer = get_tokenizer(name, *args, tokenizer_mode=mode, **kwargs) + + reference = get_tokenizer(repo, tokenizer_mode="mistral") + assert isinstance(tokenizer, MistralTokenizer) + assert tokenizer.encode("hello world") == reference.encode("hello world") + + def test_hub_tokenizer_loads_from_local_snapshot(): """Transformers makes a Hub request per file when given a repo ID.""" with patch.object( diff --git a/tests/tokenizers_/test_detokenize.py b/tests/tokenizers_/test_detokenize.py index 1dd2a000a6..37e7252d2d 100644 --- a/tests/tokenizers_/test_detokenize.py +++ b/tests/tokenizers_/test_detokenize.py @@ -15,6 +15,7 @@ from vllm.tokenizers.detokenizer_utils import ( ) from vllm.tokenizers.mistral import MistralTokenizer from vllm.tokenizers.protocol import TokenizerLike +from vllm.transformers_utils.repo_utils import with_retry from vllm.v1.engine import EngineCoreRequest from vllm.v1.engine.detokenizer import ( FastIncrementalDetokenizer, @@ -95,10 +96,13 @@ def _run_incremental_decode( @pytest.fixture def tokenizer(tokenizer_name): - return ( - MistralTokenizer.from_pretrained(tokenizer_name) - if "mistral" in tokenizer_name - else AutoTokenizer.from_pretrained(tokenizer_name) + return with_retry( + lambda: ( + MistralTokenizer.from_pretrained(tokenizer_name) + if "mistral" in tokenizer_name + else AutoTokenizer.from_pretrained(tokenizer_name) + ), + log_msg=f"Error loading tokenizer {tokenizer_name}", ) diff --git a/tests/tokenizers_/test_registry.py b/tests/tokenizers_/test_registry.py index b9d9c732f3..73ff4bc5df 100644 --- a/tests/tokenizers_/test_registry.py +++ b/tests/tokenizers_/test_registry.py @@ -9,7 +9,7 @@ import pytest from transformers import AutoConfig, Qwen3_5MoeConfig from transformers.models.auto.configuration_auto import CONFIG_MAPPING -from vllm.tokenizers import TokenizerLike +from vllm.tokenizers import TokenizerLike, registry from vllm.tokenizers.registry import ( TokenizerRegistry, cached_get_tokenizer, @@ -17,6 +17,7 @@ from vllm.tokenizers.registry import ( cached_tokenizer_from_config, get_tokenizer, resolve_tokenizer_args, + tokenizer_args_from_config, ) @@ -116,11 +117,13 @@ def test_cached_tokenizer_from_config_registers_local_config(tmp_path: Path): ) model_config = SimpleNamespace( + model=str(tmp_path), skip_tokenizer_init=False, tokenizer=str(tmp_path), runner_type="generate", tokenizer_mode="hf", tokenizer_revision=None, + code_revision=None, trust_remote_code=True, hf_config=Qwen3_5MoeConfig(), ) @@ -164,3 +167,117 @@ def test_cached_tokenizer_from_config_registers_local_config(tmp_path: Path): CONFIG_MAPPING._extra_content.pop("qwen3_5_moe", None) if registered_config is not None: CONFIG_MAPPING._extra_content["qwen3_5_moe"] = registered_config + + +@pytest.fixture +def config_tokenizer_backend(monkeypatch): + """Keep config helpers and dispatch real; replace Hub/backend loading.""" + calls = [] + + class Backend: + @classmethod + def from_pretrained(cls, name, *args, **kwargs): + calls.append(kwargs) + return SimpleNamespace(is_fast=True) + + monkeypatch.setattr(registry, "get_config", lambda *a, **kw: None) + monkeypatch.setattr(registry, "_maybe_register_hf_config", lambda config: None) + monkeypatch.setattr(registry, "is_mistral_model_repo", lambda **kw: True) + monkeypatch.setattr(registry, "any_pattern_in_repo_files", lambda **kw: True) + monkeypatch.setattr( + registry, "get_non_weight_snapshot_path", lambda name, **kw: name + ) + monkeypatch.setattr(TokenizerRegistry, "load_tokenizer_cls", lambda mode: Backend) + cached_get_tokenizer.cache_clear() + cached_resolve_tokenizer_args.cache_clear() + yield calls + cached_get_tokenizer.cache_clear() + cached_resolve_tokenizer_args.cache_clear() + + +def _load_config_tokenizer(loader, config, **kwargs): + result = loader(config, **kwargs) + if loader is tokenizer_args_from_config: + mode, name, args, resolved_kwargs = result + return get_tokenizer(name, *args, tokenizer_mode=mode, **resolved_kwargs) + return result + + +@pytest.mark.parametrize( + "loader", (tokenizer_args_from_config, cached_tokenizer_from_config) +) +@pytest.mark.parametrize("mode", ["mistral", "auto"]) +@pytest.mark.parametrize("code_revision", [None, "model-code"]) +@pytest.mark.parametrize("overrides", [{}, {"code_revision": None}]) +def test_config_mistral_does_not_forward_code_revision( + loader, mode, code_revision, overrides, config_tokenizer_backend +): + """HF code revisions must not reach MistralCommonBackend's strict kwargs.""" + config = SimpleNamespace( + model="test-model", + tokenizer="test-model", + runner_type="generate", + tokenizer_mode=mode, + tokenizer_revision=None, + code_revision=code_revision, + trust_remote_code=False, + skip_tokenizer_init=False, + ) + _load_config_tokenizer(loader, config, **overrides) + assert len(config_tokenizer_backend) == 1 + assert "code_revision" not in config_tokenizer_backend[0] + + +@pytest.mark.parametrize( + "loader", (tokenizer_args_from_config, cached_tokenizer_from_config) +) +@pytest.mark.parametrize("missing", [False, True]) +def test_config_tokenizer_without_code_revision( + loader, missing, config_tokenizer_backend +): + """Default or minimal configs must not introduce an HF-only plugin kwarg.""" + config = SimpleNamespace( + model="test-model", + tokenizer="test-model", + runner_type="generate", + tokenizer_mode="test_tokenizer", + tokenizer_revision=None, + trust_remote_code=False, + skip_tokenizer_init=False, + ) + if not missing: + config.code_revision = None + _load_config_tokenizer(loader, config) + assert len(config_tokenizer_backend) == 1 + assert "code_revision" not in config_tokenizer_backend[0] + + +@pytest.mark.parametrize( + "loader", (tokenizer_args_from_config, cached_tokenizer_from_config) +) +@pytest.mark.parametrize("mode", ["hf", "cohere", "kimi_k3", "inkling", "deepseek_v4"]) +@pytest.mark.parametrize( + "overrides, expected", + [ + ({}, "model-code"), + ({"code_revision": None}, None), + ({"code_revision": "override"}, "override"), + ], +) +def test_config_hf_backed_tokenizer_preserves_code_revision( + loader, mode, overrides, expected, config_tokenizer_backend +): + """Filtering Mistral kwargs must preserve HF wrappers and caller overrides.""" + config = SimpleNamespace( + model="test-model", + tokenizer="test-model", + runner_type="generate", + tokenizer_mode=mode, + tokenizer_revision=None, + code_revision="model-code", + trust_remote_code=True, + skip_tokenizer_init=False, + ) + _load_config_tokenizer(loader, config, **overrides) + assert len(config_tokenizer_backend) == 1 + assert config_tokenizer_backend[0]["code_revision"] == expected diff --git a/tests/tool_parsers/test_apertus_tool_parser.py b/tests/tool_parsers/test_apertus_tool_parser.py index 5ab0a2bcc7..4c8e75e40b 100644 --- a/tests/tool_parsers/test_apertus_tool_parser.py +++ b/tests/tool_parsers/test_apertus_tool_parser.py @@ -244,6 +244,22 @@ class TestStreamingExtraction: assert tcs[0]["name"] == "get_weather" assert json.loads(tcs[0]["arguments"]) == {"location": "Paris, France"} + @pytest.mark.parametrize( + "arguments", + ['{"value": -2e-3}', '{"value": 1.5e3}', '{"values": [1.5e3, -2e-3]}'], + ) + def test_streaming_scientific_notation(self, parser, mock_request, arguments): + model_output = ( + f'{TOOL_CALLS_PREFIX}[{{"calculate": {arguments}}}]{TOOL_CALLS_SUFFIX}' + ) + + results = self._simulate_streaming(parser, mock_request, list(model_output)) + tool_calls = self._collect_tool_calls(results) + + assert len(tool_calls) == 1 + assert tool_calls[0]["name"] == "calculate" + assert json.loads(tool_calls[0]["arguments"]) == json.loads(arguments) + def test_streaming_missing_tool_suffix(self, parser, mock_request): chunks = [ "<|tools_prefix|>", diff --git a/tests/v1/attention/test_attention_backends.py b/tests/v1/attention/test_attention_backends.py index b537a17823..abe13137f7 100644 --- a/tests/v1/attention/test_attention_backends.py +++ b/tests/v1/attention/test_attention_backends.py @@ -763,7 +763,7 @@ def test_flashinfer_cross_layer_layout( _test_backend_correctness( batch_spec=BATCH_SPECS[batch_spec_name], - model="meta-llama/Meta-Llama-3-8B", + model="Qwen/Qwen3-0.6B", backend_to_test=[AttentionBackendEnum.FLASHINFER], mask_mod=causal_mask_mod, kv_cache_dtype=kv_cache_dtype, @@ -786,7 +786,7 @@ def test_flashinfer_cross_layer_layout( "single_prefill", ], ) -@pytest.mark.parametrize("model", ["meta-llama/Meta-Llama-3-8B"]) +@pytest.mark.parametrize("model", ["Qwen/Qwen3-0.6B"]) @pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4]) @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e4m3"]) def test_causal_backend_correctness( @@ -1582,7 +1582,7 @@ def test_flashinfer_native_prefill_with_sinks(default_vllm_config): _test_backend_correctness( BATCH_SPECS["small_prefill"], - "meta-llama/Meta-Llama-3-8B", + "Qwen/Qwen3-0.6B", [AttentionBackendEnum.FLASHINFER], causal_mask_mod, use_sinks=True, @@ -1629,7 +1629,7 @@ def test_flashinfer_xqa_decode_correctness(default_vllm_config): batch_spec = BATCH_SPECS["small_decode"] vllm_config = create_vllm_config( - model_name="meta-llama/Meta-Llama-3-8B", + model_name="Qwen/Qwen3-0.6B", max_model_len=max(batch_spec.seq_lens), block_size=16, ) @@ -1691,7 +1691,7 @@ def test_flashinfer_xqa_decode_correctness(default_vllm_config): _test_backend_correctness( batch_spec, - "meta-llama/Meta-Llama-3-8B", + "Qwen/Qwen3-0.6B", [AttentionBackendEnum.FLASHINFER], causal_mask_mod, ) @@ -1994,7 +1994,7 @@ if current_platform.is_rocm(): "mixed_small", ], ) -@pytest.mark.parametrize("model", ["meta-llama/Meta-Llama-3-8B"]) +@pytest.mark.parametrize("model", ["Qwen/Qwen3-0.6B"]) def test_non_causal_backend_correctness( default_vllm_config, batch_spec_name: str, diff --git a/tests/v1/attention/test_deepseek_v4_rocm_adaptive.py b/tests/v1/attention/test_deepseek_v4_rocm_adaptive.py index 7c18949881..8157bd6397 100644 --- a/tests/v1/attention/test_deepseek_v4_rocm_adaptive.py +++ b/tests/v1/attention/test_deepseek_v4_rocm_adaptive.py @@ -67,7 +67,7 @@ def test_deepseek_rocm_adaptive_builders_support_varlen_full_graphs(): DeepseekV4ROCMAiterMLASparseMetadataBuilder, DeepseekV4ROCMAiterSparseSWAMetadataBuilder, DeepseekV41ROCMAiterMLASparseMetadataBuilder, - v41_rocm.DeepseekV4ROCMAiterSparseSWAMetadataBuilder, + v41_rocm.DeepseekV41ROCMAiterSparseSWAMetadataBuilder, ): assert ( builder_cls.get_cudagraph_support(adaptive_config, SimpleNamespace()) @@ -126,7 +126,7 @@ def test_deepseek_v41_rocm_fixed_k_path_is_unchanged( for builder_cls in ( DeepseekV41ROCMAiterMLASparseMetadataBuilder, - v41_rocm.DeepseekV4ROCMAiterSparseSWAMetadataBuilder, + v41_rocm.DeepseekV41ROCMAiterSparseSWAMetadataBuilder, ): assert ( builder_cls.get_cudagraph_support(fixed_config, SimpleNamespace()) diff --git a/tests/v1/attention/test_deepseek_v4_swa_visible.py b/tests/v1/attention/test_deepseek_v4_swa_visible.py index 5dcea926f0..a46495f790 100644 --- a/tests/v1/attention/test_deepseek_v4_swa_visible.py +++ b/tests/v1/attention/test_deepseek_v4_swa_visible.py @@ -529,7 +529,7 @@ def make_builder( window: int = WINDOW, max_image_tokens: int = MAX_IMG, max_num_batched_tokens: int = 64, - model_name: str = "meta-llama/Meta-Llama-3-8B", + model_name: str = "Qwen/Qwen3-0.6B", builder_cls: type[ DeepseekSparseSWAMetadataBuilder ] = DeepseekSparseSWAMetadataBuilder, diff --git a/tests/v1/attention/test_indexer_tp_row_shard.py b/tests/v1/attention/test_indexer_tp_row_shard.py index e65e4789f1..4bbbf93eb9 100644 --- a/tests/v1/attention/test_indexer_tp_row_shard.py +++ b/tests/v1/attention/test_indexer_tp_row_shard.py @@ -10,9 +10,11 @@ import torch import vllm.model_executor.layers.sparse_attn_indexer as sparse_indexer import vllm.v1.attention.backends.mla.indexer as indexer -from vllm.config import CUDAGraphMode +from vllm.config import AttentionConfig, CUDAGraphMode from vllm.model_executor.layers.sparse_attn_indexer import PrefillRowShard +MIN_ROWS = AttentionConfig().indexer_tp_shard_min_rows + def _chunk(start: int, end: int) -> SimpleNamespace: """Chunk over rows [start, end); bounds derive from the global row.""" @@ -80,11 +82,11 @@ def test_shard_exchange_reassembles_rows_only( assert torch.all(buffer[decode + sum(sizes) :] == -7), "padding untouched" -@pytest.mark.parametrize("tp_size", [2, 4]) +@pytest.mark.parametrize("tp_size", [2, 4, 8]) def test_balanced_row_shard_partitions_and_declines(tp_size: int) -> None: """The split covers every row once with >= 1 row per rank and declines below the activation floor (or on a single rank).""" - n = indexer.MIN_TP_SHARD_ROWS_PER_RANK * tp_size + n = MIN_ROWS shapes = [ ([4 * n], [4 * n]), # fresh prompt: the full causal cost ramp # uneven requests at different context depths @@ -97,12 +99,38 @@ def test_balanced_row_shard_partitions_and_declines(tp_size: int) -> None: torch.tensor(query_lens, dtype=torch.int32), ratio, tp_size, + n, ) assert sizes is not None and len(sizes) == tp_size assert sum(sizes) == sum(query_lens) and min(sizes) >= 1 lens = torch.tensor([n - 1], dtype=torch.int32) - assert indexer.balanced_prefill_row_shard(lens, lens, ratio, tp_size) is None - assert indexer.balanced_prefill_row_shard(lens, lens, ratio, 1) is None + assert indexer.balanced_prefill_row_shard(lens, lens, ratio, tp_size, n) is None + assert indexer.balanced_prefill_row_shard(lens, lens, ratio, 1, n) is None + + +def test_row_shard_threshold_is_total_not_per_rank() -> None: + """The activation floor counts TOTAL rows regardless of tp_size: a + per-rank floor would need 128K rows at TP8, i.e. an unschedulable + --max-num-batched-tokens of 128K on 80 GB Hopper nodes.""" + lens = torch.tensor([MIN_ROWS], dtype=torch.int32) + for tp_size in (2, 4, 8): + sizes = indexer.balanced_prefill_row_shard(lens, lens, 4, tp_size, MIN_ROWS) + assert sizes is not None + + +def test_tp2_old_per_rank_window_stays_replicated() -> None: + """32K-64K total rows engaged on TP2 under the old per-rank floor + (16K rows/rank); the total floor now keeps them replicated — the + all_gatherv latency term is largest at small tp.""" + window = torch.tensor([3 * MIN_ROWS // 4], dtype=torch.int32) + assert indexer.balanced_prefill_row_shard(window, window, 4, 2, MIN_ROWS) is None + + +def test_row_shard_floor_is_at_least_tp_size() -> None: + """A floor below tp_size still leaves every rank at least one row.""" + lens = torch.tensor([8], dtype=torch.int32) + assert indexer.balanced_prefill_row_shard(lens, lens, 4, 8, 1) is not None + assert indexer.balanced_prefill_row_shard(lens - 1, lens - 1, 4, 8, 1) is None def test_row_sharding_gate_rejects_unsupported_configurations( diff --git a/tests/v1/attention/test_mamba_align_block_table.py b/tests/v1/attention/test_mamba_align_block_table.py new file mode 100644 index 0000000000..0380771388 --- /dev/null +++ b/tests/v1/attention/test_mamba_align_block_table.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Equivalence tests for the fused mamba "align" block-table gather. + +`mamba_get_block_table_tensor`'s align branch used to be five PyTorch +statements that expanded to seven kernel launches over a `(num_reqs,)` tensor. +It is now one Triton kernel. These tests pin the fused kernel to the exact +semantics of the expression it replaced, including cases a single-request +decode never exercises but other callers can hit. + +`block_table` is deliberately tested non-contiguous: the original +`torch.gather(block_table, 1, ...)` honours any layout, and this is a shared +utility (`gdn_attn`, `mamba_attn`, `short_conv_attn` and `linear_attn` all call +it), so a column stride of 1 must not be assumed. +""" + +import pytest +import torch + +from vllm.v1.attention.backends.utils import mamba_get_block_table_tensor +from vllm.v1.kv_cache_interface import MambaSpec + +WIDTH = 64 + + +def _spec(block_size: int, num_speculative_blocks: int) -> MambaSpec: + # Only block_size and num_speculative_blocks are read by the align branch; + # shapes/dtypes just have to be well formed for the frozen dataclass. + return MambaSpec( + block_size=block_size, + shapes=((1,),), + dtypes=(torch.float32,), + num_speculative_blocks=num_speculative_blocks, + ) + + +def _reference(block_table, seq_lens, block_size, num_spec_blocks): + """The align branch exactly as it was before the fusion.""" + start_indices = (seq_lens - 1) // block_size + start_indices = start_indices.clamp(min=0) + offsets = torch.arange( + 1 + num_spec_blocks, device=block_table.device, dtype=torch.int32 + ) + indices_to_gather = (start_indices.unsqueeze(1) + offsets).to(torch.int64) + return torch.gather(block_table, 1, indices_to_gather) + + +def _make(nreq, block_size, nspec, device, noncontig=False): + torch.manual_seed(0) + width = WIDTH * 2 if noncontig else WIDTH + bt = torch.randint(0, 5000, (nreq, width), dtype=torch.int32, device=device) + if noncontig: + # Every other column: stride(1) == 2, which a kernel assuming a unit + # column stride would read straight through, silently returning the + # wrong blocks. + bt = bt[:, ::2] + assert not bt.is_contiguous() + # Keep start + nspec inside the row, which the gather already requires. + hi = max(1, (WIDTH - 1 - nspec) * block_size) + sl = torch.randint(0, hi, (nreq,), dtype=torch.int32, device=device) + return bt, sl + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +@pytest.mark.parametrize("nreq", [1, 2, 8, 64]) +@pytest.mark.parametrize("block_size", [16, 64, 256]) +@pytest.mark.parametrize("nspec", [0, 1, 2, 3, 4, 7]) +def test_matches_reference(nreq, block_size, nspec): + bt, sl = _make(nreq, block_size, nspec, "cuda") + got = mamba_get_block_table_tensor(bt, sl, _spec(block_size, nspec), "align") + want = _reference(bt, sl, block_size, nspec) + assert got.shape == want.shape + assert got.dtype == want.dtype + torch.testing.assert_close(got, want, atol=0, rtol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +@pytest.mark.parametrize("nspec", [0, 3, 7]) +def test_non_contiguous_block_table(nspec): + """A column stride != 1 must be honoured, not assumed away.""" + bt, sl = _make(8, 64, nspec, "cuda", noncontig=True) + got = mamba_get_block_table_tensor(bt, sl, _spec(64, nspec), "align") + want = _reference(bt, sl, 64, nspec) + torch.testing.assert_close(got, want, atol=0, rtol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_zero_seq_len_clamps_to_first_row(): + """CUDA-graph padding rows carry seq_len 0; the old code clamped to row 0.""" + bt = torch.randint(0, 5000, (4, WIDTH), dtype=torch.int32, device="cuda") + sl = torch.zeros(4, dtype=torch.int32, device="cuda") + got = mamba_get_block_table_tensor(bt, sl, _spec(64, 3), "align") + torch.testing.assert_close(got, _reference(bt, sl, 64, 3), atol=0, rtol=0) + # start clamps to 0, so this is just the first 1 + nspec columns. + torch.testing.assert_close(got, bt[:, :4], atol=0, rtol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_zero_requests_returns_empty(): + """An empty batch must not launch a (0,) grid.""" + bt = torch.zeros((0, WIDTH), dtype=torch.int32, device="cuda") + sl = torch.zeros(0, dtype=torch.int32, device="cuda") + got = mamba_get_block_table_tensor(bt, sl, _spec(64, 3), "align") + assert got.shape == (0, 4) + assert got.dtype == bt.dtype + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_none_mode_returns_the_input_unchanged(): + """`MambaCacheMode` is Literal["align", "none"]; only "none" passes through. + + An earlier version of this test also parametrised a "mode" of "all", which + is not a member of that Literal. Anything that is not "none" takes the + align branch and returns a NEW tensor, so asserting identity for it failed + -- correctly. The passthrough contract applies to "none" alone. + """ + bt = torch.randint(0, 5000, (4, WIDTH), dtype=torch.int32, device="cuda") + sl = torch.randint(0, 1000, (4,), dtype=torch.int32, device="cuda") + assert mamba_get_block_table_tensor(bt, sl, _spec(64, 3), "none") is bt + + +# --------------------------------------------------------------------------- +# CPU tensors. Every test above runs on "cuda", which is exactly how the +# regression reached review: `mamba_get_block_table_tensor` is a SHARED +# utility and its callers legitimately pass CPU tensors -- the GDN and +# linear-attention metadata builders are unit-tested with +# `DEVICE = torch.device("cpu")`. A Triton kernel cannot run on those at all, +# while the torch expression it replaced could. These run without a GPU. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("nspec", [0, 1, 3]) +@pytest.mark.parametrize("block_size", [1, 16, 64]) +def test_cpu_tensors_match_the_reference(nspec, block_size): + """The align branch must work on CPU, not just on the accelerator.""" + width = 2 + nspec + 8 + bt = torch.arange(5 * width, dtype=torch.int32).reshape(5, width) + sl = torch.tensor([0, 1, block_size, block_size * 3 + 1, 7], dtype=torch.int32) + got = mamba_get_block_table_tensor(bt, sl, _spec(block_size, nspec), "align") + want = _reference(bt, sl, block_size, nspec) + assert torch.equal(got, want), f"cpu mismatch: {got.tolist()} != {want.tolist()}" + assert got.device == bt.device + assert got.dtype == bt.dtype + + +def test_cpu_none_mode_passes_through(): + bt = torch.arange(12, dtype=torch.int32).reshape(3, 4) + sl = torch.tensor([1, 2, 3], dtype=torch.int32) + assert mamba_get_block_table_tensor(bt, sl, _spec(16, 3), "none") is bt diff --git a/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py b/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py index 6849c4e3eb..e132a45c36 100644 --- a/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py +++ b/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py @@ -171,6 +171,11 @@ def test_paged_prefill_indices_match_gathered_prefill( from vllm.v1.kv_cache_interface import SlidingWindowMLASpec rocm = importlib.import_module(f"vllm.models.{model}.amd.rocm") + builder_cls = ( + rocm.DeepseekV41ROCMAiterSparseSWAMetadataBuilder + if model == "deepseek_v41" + else rocm.DeepseekV4ROCMAiterSparseSWAMetadataBuilder + ) device = torch.device("cuda") gen = torch.Generator().manual_seed(compress_ratio) window, width, block_size, max_model_len = 16, 32, 16, 128 @@ -233,9 +238,7 @@ def test_paged_prefill_indices_match_gathered_prefill( monkeypatch.setattr( rocm_aiter_ops, "is_triton_sparse_mla_enabled", lambda: enabled ) - return rocm.DeepseekV4ROCMAiterSparseSWAMetadataBuilder( - swa_spec, ["swa"], vllm_config, device - ).build(0, common) + return builder_cls(swa_spec, ["swa"], vllm_config, device).build(0, common) assert build(enabled=False).prefill_swa_ragged_indices is None md = build(enabled=True) diff --git a/tests/v1/attention/test_rocm_attention_backends_selection.py b/tests/v1/attention/test_rocm_attention_backends_selection.py index f103b43650..011720c042 100644 --- a/tests/v1/attention/test_rocm_attention_backends_selection.py +++ b/tests/v1/attention/test_rocm_attention_backends_selection.py @@ -8,6 +8,8 @@ from unittest.mock import MagicMock, patch import pytest import torch +from vllm.config.attention import AttentionConfig +from vllm.config.parallel import ParallelConfig from vllm.platforms import current_platform from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.attention.selector import AttentionSelectorConfig @@ -44,6 +46,48 @@ def cleared_attention_selector_cache(): _cached_get_attn_backend.cache_clear() +def test_deepseek_rocm_backend_registry_paths_are_model_specific(): + assert AttentionBackendEnum.ROCM_FLASHMLA_SPARSE_DSV4.get_path() == ( + "vllm.models.deepseek_v4.amd.rocm.DeepseekV4ROCMAiterMLASparseBackend" + ) + assert AttentionBackendEnum.ROCM_FLASHMLA_SPARSE_DSV41.get_path() == ( + "vllm.models.deepseek_v41.amd.rocm.DeepseekV41ROCMAiterMLASparseBackend" + ) + + +@pytest.mark.parametrize( + "backend_name, expect_warning", + [ + ("ROCM_FLASHMLA_SPARSE_DSV4", True), + ("ROCM_FLASHMLA_SPARSE_DSV41", False), + ], + ids=["legacy", "canonical"], +) +def test_deepseek_v41_rocm_accepts_legacy_backend_alias(backend_name, expect_warning): + from vllm.models.deepseek_v41.amd import model as dsv41_model + from vllm.models.deepseek_v41.amd.rocm import ( + DeepseekV41ROCMAiterMLAAttention, + DeepseekV41ROCMAiterMLASparseBackend, + ) + + config = SimpleNamespace( + attention_config=AttentionConfig(backend=backend_name), + parallel_config=ParallelConfig(), + ) + + with patch.object(dsv41_model.logger, "warning_once") as warning_once: + selected = dsv41_model._select_dsv4_attn_cls(config) + + assert selected is DeepseekV41ROCMAiterMLAAttention + if expect_warning: + warning_once.assert_called_once() + else: + warning_once.assert_not_called() + assert ( + DeepseekV41ROCMAiterMLASparseBackend.get_name() == "ROCM_FLASHMLA_SPARSE_DSV41" + ) + + def test_aiter_unified_attention_uses_dedicated_metadata_builder(): from vllm.v1.attention.backends.rocm_aiter_unified_attn import ( RocmAiterUnifiedAttentionBackend, diff --git a/tests/v1/attention/utils.py b/tests/v1/attention/utils.py index abb29a03c4..a0a24fd176 100644 --- a/tests/v1/attention/utils.py +++ b/tests/v1/attention/utils.py @@ -180,7 +180,7 @@ def create_standard_kv_cache_spec( def create_vllm_config( - model_name: str = "meta-llama/Meta-Llama-3-8B", + model_name: str = "Qwen/Qwen3-0.6B", tensor_parallel_size: int = 1, max_model_len: int = 1024, dtype: ModelDType | torch.dtype = "auto", diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 1489315ec2..6a0b6393a2 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -52,6 +52,7 @@ from vllm.v1.core.kv_cache_utils import ( hash_block_tokens, init_none_hash, make_block_hash_with_group_id, + null_head_len, tensor_data, to_event_extra_keys, ) @@ -2156,3 +2157,20 @@ def test_wrapped_mamba_group_requires_block_zeroing(): ) assert config.needs_kv_cache_zeroing + + +def test_null_head_len(): + null = KVCacheBlock(0, is_null=True) + real = [KVCacheBlock(i) for i in range(1, 4)] + tables = [ + [], + real, + [null] * 5, + [null] * 5 + real, + [null] * 2 + real[:1] + [null] * 3 + real[1:], # not one contiguous head + ] + for table in tables: + expected = next((i for i, b in enumerate(table) if not b.is_null), len(table)) + assert null_head_len(table) == expected + blocks = KVCacheBlocks(([null] * 4 + real, [null] + real[:1] + [null])) + assert blocks.get_unhashed_block_ids_all_groups() == [[1, 2, 3], [1]] diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 99eb02ff33..7b2b5b13b6 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -558,6 +558,249 @@ def test_schedule_prefills_gating(has_running: bool): assert any(r.req_id == "new0" for r in output.scheduled_new_reqs) +def _decode_one(scheduler, output, req_ids: list[str]) -> None: + """Feed back one sampled token per request, so they keep decoding.""" + scheduler.update_from_output( + output, + ModelRunnerOutput( + req_ids=req_ids, + req_id_to_index={req_id: i for i, req_id in enumerate(req_ids)}, + sampled_token_ids=[[0] for _ in req_ids], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ), + ) + + +def _start_decoding(scheduler, req_id: str = "run0"): + """Prefill one request and take it into the decode state.""" + (request,) = create_requests( + num_requests=1, num_tokens=8, max_tokens=100, req_ids=[req_id] + ) + scheduler.add_request(request) + output = scheduler.schedule() + assert req_id in output.num_scheduled_tokens + _decode_one(scheduler, output, [req_id]) + assert len(scheduler.running) == 1 + return output + + +@pytest.mark.parametrize("interval", [2, 4]) +def test_local_prefill_interval_holds_prefill_between_bursts(interval: int): + """A single engine applies `prefill_schedule_interval` in the scheduler, + timed from the last prefill-carrying step: a waiting prefill is held for the + remainder of the interval while running requests keep decoding, then + admitted once the interval has elapsed. + """ + scheduler = create_scheduler(prefill_schedule_interval=interval) + # This prefill starts the interval. + _start_decoding(scheduler) + assert scheduler.last_prefill_step == scheduler.num_executed_steps + + (new_req,) = create_requests(num_requests=1, num_tokens=8, req_ids=["new0"]) + scheduler.add_request(new_req) + + # Steps inside the interval decode alone; the waiting prefill is held. + for _ in range(interval - 1): + output = scheduler.schedule() + assert "new0" not in output.num_scheduled_tokens + assert "run0" in output.num_scheduled_tokens + assert new_req.status == RequestStatus.WAITING + _decode_one(scheduler, output, ["run0"]) + + # The interval has now elapsed, so prefill is admitted again. + output = scheduler.schedule() + assert "new0" in output.num_scheduled_tokens + assert scheduler.last_prefill_step == scheduler.num_executed_steps + + +def test_empty_schedule_does_not_advance_the_prefill_cadence(): + """An empty schedule() must not consume interval budget. + + `current_step` counts every schedule() call. With async scheduling, + step_with_batch_queue() calls schedule() again whenever the batch queue has + room and has_requests() is true -- so a call lands while the previous batch + is still in flight and nothing new can be scheduled. An async KV connector + makes this common, since requests sit in lookup or WAITING_FOR_REMOTE_KVS. + Timing the cadence on current_step lets those empty calls expire the hold + early. Reported on #54627 against the LMCache MP connector. + """ + scheduler = create_scheduler(prefill_schedule_interval=4) + _start_decoding(scheduler) + + # One real decode step, left in flight (no update_from_output yet). + output = scheduler.schedule() + assert output.total_num_scheduled_tokens > 0 + executed_before = scheduler.num_executed_steps + step_before = scheduler.current_step + + # Re-entering schedule() while that batch is still unaccounted for + # schedules nothing -- this is the async-scheduling shape. + for _ in range(3): + empty = scheduler.schedule() + assert empty.total_num_scheduled_tokens == 0 + + assert scheduler.current_step == step_before + 3, "current_step should advance" + assert scheduler.num_executed_steps == executed_before, ( + "an empty schedule() must not count as an executed step" + ) + + +def test_prefill_cadence_counts_executed_steps_not_schedule_calls(): + """The hold spans `interval` EXECUTED steps, however many empty schedule() + calls are interleaved. + + Counting schedule() calls instead lets each empty call eat interval budget, + so with an async KV connector prefill lands on nearly every step and the + feature stops doing anything. + """ + interval = 4 + scheduler = create_scheduler(prefill_schedule_interval=interval) + _start_decoding(scheduler) + prefill_at = scheduler.num_executed_steps + + (new_req,) = create_requests(num_requests=1, num_tokens=8, req_ids=["new0"]) + scheduler.add_request(new_req) + + empty_calls = 0 + admitted_at = None + for _ in range(50): + output = scheduler.schedule() + if "new0" in output.num_scheduled_tokens: + admitted_at = scheduler.num_executed_steps + break + if output.total_num_scheduled_tokens == 0: + empty_calls += 1 + continue + # A real decode step. Leave it in flight and re-enter schedule() twice, + # as async scheduling does while the batch queue has room. + for _ in range(2): + extra = scheduler.schedule() + if "new0" in extra.num_scheduled_tokens: + admitted_at = scheduler.num_executed_steps + break + empty_calls += 1 + if admitted_at is not None: + break + _decode_one(scheduler, output, ["run0"]) + + assert admitted_at is not None, "prefill was never admitted" + assert empty_calls > 0, "test never exercised an empty schedule() call" + # The hold is measured in executed steps. Counting schedule() calls would + # admit after far fewer, because the empty re-entries above would count. + assert admitted_at - prefill_at == interval, ( + f"hold lasted {admitted_at - prefill_at} executed steps, expected " + f"{interval}; empty schedule() calls were counted against the cadence" + ) + assert scheduler.current_step > scheduler.num_executed_steps + + +def test_local_prefill_interval_defaults_to_no_gating(): + """The default interval of 1 leaves admission exactly as it was: a waiting + prefill is admitted on the very next step. + """ + scheduler = create_scheduler() + assert scheduler.local_prefill_interval == 1 + _start_decoding(scheduler) + + (new_req,) = create_requests(num_requests=1, num_tokens=8, req_ids=["new0"]) + scheduler.add_request(new_req) + output = scheduler.schedule() + assert "new0" in output.num_scheduled_tokens + + +def test_local_prefill_interval_left_to_engine_under_dp(): + """Under data parallelism the cadence must stay aligned across ranks, so the + DP engine core drives it and the scheduler must not gate on its own. + """ + scheduler = create_scheduler(prefill_schedule_interval=4, data_parallel_size=2) + assert scheduler.local_prefill_interval == 1 + _start_decoding(scheduler) + + (new_req,) = create_requests(num_requests=1, num_tokens=8, req_ids=["new0"]) + scheduler.add_request(new_req) + output = scheduler.schedule() + assert "new0" in output.num_scheduled_tokens + + +def test_local_prefill_interval_never_idles_the_model(): + """With nothing running to decode in prefill's place, the interval must not + hold prefill back, or the step would be wasted entirely. + """ + scheduler = create_scheduler(prefill_schedule_interval=8) + _start_decoding(scheduler) + + # Finish the only running request, leaving nothing to decode. + scheduler.finish_requests("run0", RequestStatus.FINISHED_ABORTED) + assert not scheduler.running + + (new_req,) = create_requests(num_requests=1, num_tokens=8, req_ids=["new0"]) + scheduler.add_request(new_req) + output = scheduler.schedule() + assert "new0" in output.num_scheduled_tokens + + +@pytest.mark.parametrize("interval", [2, 3]) +def test_local_prefill_interval_defers_inflight_prefill_chunk(interval: int): + """The single-engine interval gates prefill WORK, not just new admissions. + + A request that is already partially prefilled is held for the remainder of + the interval too, exactly as `throttle_prefills` does for DP. That is the + point of the cadence -- a held step must carry no prefill at all, or decode + does not get the uninterrupted stretch the interval exists to give it. + """ + scheduler = create_scheduler( + max_num_seqs=16, + max_num_batched_tokens=50, + enable_chunked_prefill=True, + prefill_schedule_interval=interval, + ) + + # A short request that finishes prefill in one step -> a running decode. + (decode_req,) = create_requests(num_requests=1, num_tokens=4, req_ids=["dec0"]) + scheduler.add_request(decode_req) + output = scheduler.schedule() + _decode_one(scheduler, output, ["dec0"]) + assert decode_req in scheduler.running and not decode_req.is_prefill_chunk + + # A long request (80 tokens, budget 50) -> prefilled in chunks. That first + # prefill closed the interval, so walk to the next step that opens it. + (chunk_req,) = create_requests(num_requests=1, num_tokens=80, req_ids=["chk0"]) + scheduler.add_request(chunk_req) + for _ in range(interval - 1): + output = scheduler.schedule() + assert "chk0" not in output.num_scheduled_tokens + _decode_one(scheduler, output, ["dec0"]) + + output = scheduler.schedule() + assert output.num_scheduled_tokens["chk0"] > 0 + scheduler.update_from_output( + output, + ModelRunnerOutput( + req_ids=["dec0", "chk0"], + req_id_to_index={"dec0": 0, "chk0": 1}, + sampled_token_ids=[[0], []], # no token sampled for partial prefill + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ), + ) + assert chunk_req.is_prefill_chunk # still mid-prefill, in running + + # Held steps: the in-flight chunk is deferred along with new admissions, + # while the decode keeps running. + for _ in range(interval - 1): + output = scheduler.schedule() + assert "chk0" not in output.num_scheduled_tokens + assert "dec0" in output.num_scheduled_tokens + _decode_one(scheduler, output, ["dec0"]) + + # Once the interval has elapsed the chunk resumes. + output = scheduler.schedule() + assert "chk0" in output.num_scheduled_tokens + + def _setup_remote_kv_resume(num_prompt_tokens: int, matched_tokens: int): """Drive a remote-KV request `r2` to the resume point (async load complete) while another request `r1` is already decoding, so the step is throttle- @@ -1916,6 +2159,46 @@ def test_draft_slots_budgeted_per_scheduled_request(tmp_path, monkeypatch): assert scheduler.schedule().num_scheduled_tokens == {"0": 10, "1": 4} +@pytest.mark.parametrize(("method", "max_num_reqs"), [("dflash", 4), ("dspark", 5)]) +def test_block_drafters_cap_requests_instead_of_draft_slots( + tmp_path, method, max_num_reqs +): + (tmp_path / "config.json").write_text( + '{"architectures": ["OPTForCausalLM"], "model_type": "opt"}' + ) + base = create_scheduler( + model=str(tmp_path), + max_num_batched_tokens=20, + num_speculative_tokens=4, + parallel_drafting=True, + skip_tokenizer_init=True, + ) + # DFlash/DSpark need a resolvable draft model, so switch the ngram config in + # place and redo the config step. DSpark then has 4 lookahead tokens, not 5. + vllm_config = base.vllm_config + object.__setattr__(vllm_config.speculative_config, "method", method) + object.__setattr__( + vllm_config.speculative_config, + "draft_model_config", + SimpleNamespace(hf_config=SimpleNamespace()), + ) + vllm_config._set_max_num_scheduled_tokens() + scheduler = Scheduler( + vllm_config=vllm_config, + kv_cache_config=base.kv_cache_config, + structured_output_manager=base.structured_output_manager, + block_size=base.block_size, + ) + for request in create_requests(num_requests=8, num_tokens=2): + scheduler.add_request(request) + + # Admission stops at 20 // num_lookahead_tokens, and no draft slots are + # charged: each admitted prompt is scheduled in full. + assert scheduler.schedule().num_scheduled_tokens == { + str(i): 2 for i in range(max_num_reqs) + } + + # Note - these test cases mirror some of those in test_rejection_sampler.py @pytest.mark.parametrize( "spec_tokens,output_tokens,expected,expected_per_req", diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index 499258f20e..5c234f20f4 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/tests/v1/core/test_single_type_kv_cache_manager.py @@ -1013,6 +1013,7 @@ def test_predictor_matches_allocator_blocks_calculation_with_admission_cap(): ) total_computed = num_tokens + def test_mamba_honors_drop_eagle_block(): """MambaManager must drop the EAGLE lookahead block, like FullAttentionManager. @@ -1067,7 +1068,7 @@ def test_mamba_honors_drop_eagle_block(): kv_cache_spec=spec, drop_eagle_block=drop_eagle_block, alignment_tokens=block_size, - )[0] + )[0][0] ) # Control: with no eagle drop, both groups see the whole cached prefix. @@ -1082,3 +1083,35 @@ def test_mamba_honors_drop_eagle_block(): f"MambaManager ignored drop_eagle_block: full-attention dropped the " f"lookahead block (hit={full_hit}), mamba kept it (hit={mamba_hit})" ) + + +def test_pop_blocks_for_free_trims_null_head(): + block_size = 2 + spec = SlidingWindowSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=4, + ) + block_pool = BlockPool( + num_gpu_blocks=10, enable_caching=True, hash_block_size=block_size + ) + manager = get_sliding_window_manager(spec, block_pool) + null = block_pool.null_block + real = block_pool.get_new_blocks(2) + manager.req_to_blocks["swa"] = [null] * 3 + real + manager.req_to_blocks["dense"] = dense = block_pool.get_new_blocks(2) + + assert manager.pop_blocks_for_free("swa") == real + assert manager.pop_blocks_for_free("dense") is dense # no null head, no copy + assert "swa" not in manager.req_to_blocks + assert all(b.ref_cnt == 1 for b in real + dense) + + # The null block is never pooled, so touch / free_blocks leave it alone. + null_ref_cnt = null.ref_cnt + num_free_blocks = block_pool.get_num_free_blocks() + block_pool.touch([null]) + block_pool.free_blocks([null]) + assert null.ref_cnt == null_ref_cnt + assert block_pool.get_num_free_blocks() == num_free_blocks diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index 541e55f5b6..25e9373888 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -85,6 +85,7 @@ def create_scheduler( kv_cache_spec: KVCacheSpec | None = None, runner: RunnerOption = "auto", per_request_spec_decode_metrics: str = "none", + prefill_schedule_interval: int = 1, scheduling_policy: SchedulerPolicy = "fcfs", diffusion_canvas_length: int | None = None, scheduler_cls: type[Scheduler] | None = None, @@ -134,6 +135,7 @@ def create_scheduler( enable_chunked_prefill=enable_chunked_prefill, async_scheduling=async_scheduling, is_encoder_decoder=model_config.is_encoder_decoder, + prefill_schedule_interval=prefill_schedule_interval, # Ensure admission/preemption mechanics are deterministic watermark=0.0, policy=scheduling_policy, diff --git a/tests/v1/cudagraph/test_encoder_cudagraph.py b/tests/v1/cudagraph/test_encoder_cudagraph.py index 2799a3ef4d..8015e9af33 100644 --- a/tests/v1/cudagraph/test_encoder_cudagraph.py +++ b/tests/v1/cudagraph/test_encoder_cudagraph.py @@ -160,6 +160,8 @@ def _make_manager_with_budgets(budgets: list[int]) -> EncoderCudaGraphManager: Skips the parts of __init__ that require a real VllmConfig / model by patching the attributes directly after construction. """ + from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphConfig + mgr = object.__new__(EncoderCudaGraphManager) mgr.token_budgets = sorted(budgets) mgr.path_token_budgets = {"default": mgr.token_budgets} @@ -1044,6 +1046,278 @@ class TestGetInputModality: assert model.get_input_modality(mm_kwargs) == "video" +# --------------------------------------------------------------------------- +# GPU tests — DeepSeek-VL2 batched-tile encoder CUDA graph +# --------------------------------------------------------------------------- + +# DeepSeek-VL2 mock constants (kept small for fast capture) +_DSV2_H = 2 # projector output side per tile (h = w) +_DSV2_DIM = 16 # hidden size +_DSV2_IMG_SIZE = 4 # each tile is 4×4 pixels +_DSV2_BUDGETS = [16, 64] +_DSV2_MAX_BATCH = 4 + + +def _dsv2_output_tokens(tw: int, th: int) -> int: + """Expected output tokens for an image with (tw, th) local tiles.""" + h = w = _DSV2_H + return h * (w + 1) + th * h * (tw * w + 1) + 1 + + +def _make_deepseek_vl2_mm_kwargs( + spatial_crops: list[list[int]], + device: torch.device, + dtype: torch.dtype, +) -> dict[str, Any]: + """spatial_crops: [[tw, th], ...] per image.""" + total_tiles = sum(1 + int(tw) * int(th) for tw, th in spatial_crops) + pixel_values = torch.randn( + total_tiles, 3, _DSV2_IMG_SIZE, _DSV2_IMG_SIZE, device=device, dtype=dtype + ) + return { + "pixel_values": pixel_values, + "images_spatial_crop": torch.tensor( + spatial_crops, dtype=torch.long, device=device + ), + } + + +class MockDeepseekVL2Model(torch.nn.Module, SupportsEncoderCudaGraph): + """Mock of DeepseekVLV2ForCausalLM for CUDA graph tests. + + Simulates the batched-tile ViT pattern: pixel_values is + [N_tiles, 3, H, W], output is assembled with newlines and + view-separator in postprocess_encoder_output. + """ + + def __init__(self): + super().__init__() + hw = _DSV2_H * _DSV2_H + in_features = 3 * _DSV2_IMG_SIZE * _DSV2_IMG_SIZE + self.proj = torch.nn.Linear(in_features, hw * _DSV2_DIM) + self.image_newline = torch.nn.Parameter(torch.randn(_DSV2_DIM)) + self.view_seperator = torch.nn.Parameter(torch.randn(_DSV2_DIM)) + self.global_view_pos = "head" + + def get_max_frames_per_video(self) -> int: + return 1 + + def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig: + return EncoderCudaGraphConfig( + modalities=["image"], + buffer_keys=["pixel_values"], + out_hidden_size=_DSV2_DIM, + ) + + def get_encoder_cudagraph_budget_range(self, vllm_config) -> tuple[int, int]: + return (_dsv2_output_tokens(1, 1), 128) + + def get_encoder_cudagraph_item_specs( + self, mm_kwargs: dict[str, Any] + ) -> list[EncoderItemSpec]: + h = w = _DSV2_H + specs = [] + for row in mm_kwargs["images_spatial_crop"].tolist(): + tw, th = int(row[0]), int(row[1]) + if tw == 0 or th == 0: + break + specs.append( + EncoderItemSpec( + input_size=1 + tw * th, + output_tokens=h * (w + 1) + th * h * (tw * w + 1) + 1, + ) + ) + return specs + + def select_encoder_cudagraph_items( + self, mm_kwargs: dict[str, Any], indices: list[int] + ) -> dict[str, Any]: + pixel_values = mm_kwargs["pixel_values"] + images_spatial_crop = mm_kwargs["images_spatial_crop"] + + if not indices: + return { + "pixel_values": pixel_values[:0], + "images_spatial_crop": images_spatial_crop[:0], + } + + tiles_per_image = [ + 1 + int(row[0]) * int(row[1]) for row in images_spatial_crop.tolist() + ] + cum = [0] + for t in tiles_per_image: + cum.append(cum[-1] + t) + + selected_pv = torch.cat([pixel_values[cum[i] : cum[i + 1]] for i in indices]) + return { + "pixel_values": selected_pv, + "images_spatial_crop": images_spatial_crop[list(indices)], + } + + def prepare_encoder_cudagraph_capture_inputs( + self, + token_budget: int, + max_batch_size: int, + max_frames_per_batch: int, + device: torch.device, + dtype: torch.dtype, + path: str = "default", + axis_keys: tuple[Hashable, ...] | None = None, + ) -> EncoderCudaGraphCaptureInputs: + hw = _DSV2_H * _DSV2_H + max_tiles = (token_budget + hw - 1) // hw + dummy = torch.randn( + max_tiles, 3, _DSV2_IMG_SIZE, _DSV2_IMG_SIZE, device=device, dtype=dtype + ) + return EncoderCudaGraphCaptureInputs(values={"pixel_values": dummy}) + + def prepare_encoder_cudagraph_replay_buffers( + self, + mm_kwargs: dict[str, Any], + max_batch_size: int, + max_frames_per_batch: int, + path: str = "default", + ) -> EncoderCudaGraphReplayBuffers: + return EncoderCudaGraphReplayBuffers( + values={"pixel_values": mm_kwargs["pixel_values"]} + ) + + def _tile_forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + """Project tiles: [N, 3, H, W] → [N, hw, DIM].""" + n = pixel_values.shape[0] + flat = pixel_values.view(n, -1) + hw = _DSV2_H * _DSV2_H + return self.proj(flat).view(n, hw, _DSV2_DIM) + + def encoder_cudagraph_forward( + self, values: dict[str, torch.Tensor], path: str = "default" + ) -> torch.Tensor: + return self._tile_forward(values["pixel_values"]) + + def encoder_eager_forward( + self, mm_kwargs: dict[str, Any], path: str = "default" + ) -> torch.Tensor: + # Manager always runs `postprocess_encoder_output` after this, whether + # the batch went through eager or CUDA graph replay, so this must + # return the same raw per-tile shape as `encoder_cudagraph_forward`. + return self._tile_forward(mm_kwargs["pixel_values"]) + + def postprocess_encoder_output( + self, + outputs: dict[str, torch.Tensor], + indices: list[int], + per_item_out_tokens: list[int], + dest: dict[int, torch.Tensor] | list[torch.Tensor | None], + clone: bool = False, + batch_mm_kwargs: dict[str, Any] | None = None, + ) -> None: + assert batch_mm_kwargs is not None + output = outputs["default"] + images_spatial_crop = batch_mm_kwargs["images_spatial_crop"] + _, hw, n_dim = output.shape + h = w = int(hw**0.5) + + tile_offset = 0 + for rank, img_idx in enumerate(indices): + tw = int(images_spatial_crop[rank][0]) + th = int(images_spatial_crop[rank][1]) + n_tiles = 1 + tw * th + tiles = output[tile_offset : tile_offset + n_tiles] + tile_offset += n_tiles + + global_f = tiles[0].view(h, w, n_dim) + newline_g = self.image_newline.view(1, 1, n_dim).expand(h, 1, n_dim) + global_f = torch.cat([global_f, newline_g], dim=1).view(-1, n_dim) + + local_f = tiles[1:].view(th, tw, h, w, n_dim) + local_f = local_f.permute(0, 2, 1, 3, 4).reshape(th * h, tw * w, n_dim) + newline_l = self.image_newline.view(1, 1, n_dim).expand(th * h, 1, n_dim) + local_f = torch.cat([local_f, newline_l], dim=1).view(-1, n_dim) + + if self.global_view_pos == "head": + emb = torch.cat([global_f, self.view_seperator[None], local_f]) + else: + emb = torch.cat([local_f, self.view_seperator[None], global_f]) + + if isinstance(dest, dict): + dest[img_idx] = emb.clone() if clone else emb + else: + dest[rank] = emb.clone() if clone else emb + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda") +class TestDeepseekVL2CudaGraph: + """CUDA graph tests for the batched-tile encoder pattern (DeepSeek-VL2).""" + + def setup_method(self): + self.device = torch.device("cuda:0") + self.dtype = torch.float16 + self.model = MockDeepseekVL2Model().to(self.device).half() + self.mgr = _make_manager_for_gpu( + self.model, _DSV2_BUDGETS, _DSV2_MAX_BATCH, self.device, self.dtype + ) + self.graph_pool = current_platform.graph_pool_handle() + self.mgr.capture(graph_pool=self.graph_pool) + + def test_capture_creates_one_graph_per_budget(self): + assert set(self.mgr.budget_graphs["default"].keys()) == set(_DSV2_BUDGETS) + + def test_execute_returns_one_tensor_per_image(self): + mm_kwargs = _make_deepseek_vl2_mm_kwargs( + [[1, 1], [1, 1]], self.device, self.dtype + ) + result = self.mgr.execute(mm_kwargs) + assert result is not None + assert len(result) == 2 + + def test_output_tokens_1x1_tile(self): + # global=h*(w+1)=2*3=6, local=1*2*(1*2+1)=6, sep=1 → 13 + expected = _dsv2_output_tokens(1, 1) + mm_kwargs = _make_deepseek_vl2_mm_kwargs([[1, 1]], self.device, self.dtype) + result = self.mgr.execute(mm_kwargs) + assert result is not None + assert result[0].shape == (expected, _DSV2_DIM) + + def test_output_tokens_2x1_tile(self): + # global=6, local=1*2*(2*2+1)=10, sep=1 → 17 + expected = _dsv2_output_tokens(2, 1) + mm_kwargs = _make_deepseek_vl2_mm_kwargs([[2, 1]], self.device, self.dtype) + result = self.mgr.execute(mm_kwargs) + assert result is not None + assert result[0].shape == (expected, _DSV2_DIM) + + def test_output_tokens_multi_image(self): + # Image 0: tw=1,th=1 → 13; Image 1: tw=2,th=1 → 17 + e0 = _dsv2_output_tokens(1, 1) + e1 = _dsv2_output_tokens(2, 1) + mm_kwargs = _make_deepseek_vl2_mm_kwargs( + [[1, 1], [2, 1]], self.device, self.dtype + ) + result = self.mgr.execute(mm_kwargs) + assert result is not None + assert result[0].shape == (e0, _DSV2_DIM) + assert result[1].shape == (e1, _DSV2_DIM) + + def test_eager_fallback_when_tokens_exceed_all_budgets(self): + # tw=4,th=4: global=6, local=4*2*(4*2+1)=72, sep=1 → 79 > max budget 64 + tw, th = 4, 4 + expected = _dsv2_output_tokens(tw, th) + assert expected > _DSV2_BUDGETS[-1], "test precondition: must exceed max budget" + mm_kwargs = _make_deepseek_vl2_mm_kwargs([[tw, th]], self.device, self.dtype) + result = self.mgr.execute(mm_kwargs) + assert result is not None + assert len(result) == 1 + assert result[0].shape == (expected, _DSV2_DIM) + assert self.mgr.graph_misses >= 1 + + def test_graph_hit_counter(self): + mm_kwargs = _make_deepseek_vl2_mm_kwargs( + [[1, 1], [1, 1]], self.device, self.dtype + ) + self.mgr.execute(mm_kwargs) + assert self.mgr.graph_hits == 2 + + # --------------------------------------------------------------------------- # GPU tests — video capture, replay, fallback, and mixed image+video # --------------------------------------------------------------------------- diff --git a/tests/v1/distributed/test_async_llm_dp.py b/tests/v1/distributed/test_async_llm_dp.py index f4b4ca2a79..7b350fc55e 100644 --- a/tests/v1/distributed/test_async_llm_dp.py +++ b/tests/v1/distributed/test_async_llm_dp.py @@ -368,9 +368,11 @@ async def test_dp_pause_late_request_does_not_block_drain(): # down first. The front-end sets its own copy optimistically when it # forwards the wake, so sample only after several coordinator # publishes (every 100ms while stats change) have overwritten it. + # Rank 0 reports the wave only after its first step, which can outlast + # the 2s on a cold device, so wait for the report rather than sample it. await asyncio.sleep(2) assert not long_request.done(), "the warmup request was too short to sample" - assert engine.engine_core.dp_engines_running(), ( + assert await _poll_flag(engine, True, timeout=30), ( "the coordinator does not report the engines as running while they are" ) diff --git a/tests/v1/distributed/test_external_lb_dp.py b/tests/v1/distributed/test_external_lb_dp.py index b53d0d8af0..1afb134ef5 100644 --- a/tests/v1/distributed/test_external_lb_dp.py +++ b/tests/v1/distributed/test_external_lb_dp.py @@ -4,17 +4,28 @@ import asyncio import os import threading import time +from concurrent.futures import Future, ThreadPoolExecutor from contextlib import AsyncExitStack +from types import SimpleNamespace +from typing import TYPE_CHECKING, cast +import msgspec import openai # use the official client for correctness check import pytest import pytest_asyncio import requests from tests.utils import RemoteOpenAIServer +from vllm.distributed.elastic_ep.external_elastic_ep import ( + ExternalElasticEPScaleCoordinator, +) from vllm.platforms import current_platform +if TYPE_CHECKING: + from vllm.v1.engine.core_client import DPAsyncMPClient + MODEL_NAME = os.getenv("MODEL_NAME", "ibm-research/PowerMoE-3b") +ELASTIC_EP_MODEL_NAME = "deepseek-ai/DeepSeek-V2-Lite-Chat" # Number of data parallel ranks for external LB testing DP_SIZE = int(os.getenv("DP_SIZE", "2")) @@ -22,6 +33,111 @@ DP_SIZE = int(os.getenv("DP_SIZE", "2")) TP_SIZE = int(os.getenv("TP_SIZE", "1")) +class DictStore: + def __init__(self, values: dict[str, bytes]): + self.values = values + + def check(self, keys: list[str]) -> bool: + return all(key in self.values for key in keys) + + def get(self, key: str) -> bytes: + return self.values[key] + + def set(self, key: str, value: bytes) -> None: + self.values[key] = value + + +def test_external_elastic_ep_visible_epoch_has_bootstrap_and_phase(monkeypatch): + client = SimpleNamespace( + vllm_config=SimpleNamespace( + parallel_config=SimpleNamespace( + data_parallel_master_port=12345, + _data_parallel_master_port_list=[12346, 12347], + ) + ) + ) + coordinator = ExternalElasticEPScaleCoordinator(cast("DPAsyncMPClient", client)) + control_store, reconfig_store = DictStore({}), DictStore({}) + coordinator.reconfig_store_ref = reconfig_store + monkeypatch.setattr( + coordinator, "_setup_reconfig_bootstrap", lambda: ("127.0.0.1", 12348) + ) + epoch, bootstrap = coordinator._prepare_reconfig_bootstrap(control_store, 2, 3) + + for store in (control_store, reconfig_store): + # DictStore preserves insertion order: metadata must precede the epoch. + assert list(store.values)[-1] == coordinator.key("current_epoch") + assert store.get(coordinator.key(epoch, "bootstrap")) == msgspec.msgpack.encode( + bootstrap + ) + assert store.get(coordinator.key(epoch, "phase")) == b"preparing" + assert store.get(coordinator.key("current_epoch")) == epoch.encode() + monkeypatch.setattr( + coordinator, "_get_reconfig_store", lambda store=store: store + ) + assert coordinator.get_status()["requested_data_parallel_size"] == 3 + + +def test_external_elastic_ep_calculates_target_expert_redundancy(): + client = SimpleNamespace( + vllm_config=SimpleNamespace( + model_config=SimpleNamespace(get_num_experts=lambda: 128), + parallel_config=SimpleNamespace( + eplb_config=SimpleNamespace(num_redundant_experts=0) + ), + ) + ) + coordinator = ExternalElasticEPScaleCoordinator(cast("DPAsyncMPClient", client)) + + assert coordinator._calculate_num_redundant_experts(2, 3) == 64 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("dp_rank", range(4)) +async def test_validate_expert_capacity_on_every_rank(dp_rank, monkeypatch): + """Validate expert capacity on every rank before waiting for an epoch.""" + client = SimpleNamespace( + vllm_config=SimpleNamespace( + model_config=SimpleNamespace(get_num_experts=lambda: 128), + parallel_config=SimpleNamespace( + data_parallel_rank=dp_rank, + elastic_ep_max_dp_size=4, + data_parallel_master_ip="127.0.0.1", + _coord_store_port=12345, + eplb_config=SimpleNamespace(num_redundant_experts=0), + ), + ) + ) + coordinator = ExternalElasticEPScaleCoordinator(cast("DPAsyncMPClient", client)) + store = DictStore({}) + monkeypatch.setattr( + "vllm.distributed.utils.get_cached_tcp_store_client", lambda *args: store + ) + + with pytest.raises(ValueError, match="cannot retain all model experts"): + await asyncio.wait_for(coordinator.prepare(4, 2), timeout=1) + + assert not store.values + assert coordinator.prepared_scale is None + + +@pytest.mark.asyncio +async def test_external_elastic_ep_late_rank_observes_epoch_error(): + coordinator = ExternalElasticEPScaleCoordinator( + cast("DPAsyncMPClient", SimpleNamespace()) + ) + epoch = "failed-epoch" + store = DictStore( + { + coordinator.key("current_epoch"): epoch.encode(), + coordinator.key(epoch, "error"): b"scale failed", + } + ) + + with pytest.raises(RuntimeError, match="scale failed"): + await coordinator._wait_for_bootstrap(store, requested_new_dp_size=3) + + class ExternalLBServerManager: """Manages data parallel vLLM server instances for external load balancer testing.""" @@ -42,57 +158,69 @@ class ExternalLBServerManager: self.servers: list[tuple[RemoteOpenAIServer, list[str]]] = [] self.server_threads: list[threading.Thread] = [] + def start_rank( + self, + rank: int, + dp_size: int | None = None, + elastic_ep_scale_up_launch: bool = False, + ) -> tuple[RemoteOpenAIServer, list[str]]: + """Start one externally managed DP rank.""" + dp_size = self.dp_size if dp_size is None else dp_size + server_args = self.base_server_args.copy() + server_args.extend( + [ + "--data-parallel-size", + str(dp_size), + "--data-parallel-rank", + str(rank), + "--data-parallel-size-local", + "1", + "--tensor-parallel-size", + str(self.tp_size), + "--port", + str(8000 + rank), + "--api-server-count", + str(self.api_server_count), + ] + ) + + env_dict = { + "VLLM_SERVER_DEV_MODE": "1", + current_platform.device_control_env_var: ",".join( + str(current_platform.device_id_to_physical_device_id(i)) + for i in range(rank * self.tp_size, (rank + 1) * self.tp_size) + ), + } + if elastic_ep_scale_up_launch: + env_dict["VLLM_ELASTIC_EP_SCALE_UP_LAUNCH"] = "1" + + server = RemoteOpenAIServer( + self.model_name, + server_args, + auto_port=False, + env_dict=env_dict, + ) + print( + f"Server rank {rank} started successfully with " + f"{self.api_server_count} API servers" + ) + server_and_args = (server, server_args) + self.servers.append(server_and_args) + return server_and_args + def __enter__(self) -> list[tuple[RemoteOpenAIServer, list[str]]]: - """Start all server instances for external LB mode.""" + """Start all initial server instances for external LB mode.""" + start_errors: list[Exception] = [] for rank in range(self.dp_size): - # Create server args for this specific rank - server_args = self.base_server_args.copy() - - # Add external LB specific arguments - server_args.extend( - [ - "--data-parallel-size", - str(self.dp_size), - "--data-parallel-rank", - str(rank), - "--data-parallel-size-local", - "1", - "--tensor-parallel-size", - str(self.tp_size), - "--port", - str(8000 + rank), # Different port for each rank - "--api-server-count", - str(self.api_server_count), - ] - ) - # Use a thread to start each server to allow parallel initialization - def start_server(r: int, sargs: list[str]): + def start_server(r: int): try: - # Start the server - server = RemoteOpenAIServer( - self.model_name, - sargs, - auto_port=False, - env_dict={ - "VLLM_SERVER_DEV_MODE": "1", - current_platform.device_control_env_var: ",".join( - str(current_platform.device_id_to_physical_device_id(i)) - for i in range(r * TP_SIZE, (r + 1) * TP_SIZE) - ), - }, - ) - server.__enter__() - print( - f"Server rank {r} started successfully with " - f"{self.api_server_count} API servers" - ) - self.servers.append((server, sargs)) + self.start_rank(r) except Exception as e: print(f"Failed to start server rank {r}: {e}") - raise + start_errors.append(e) - thread = threading.Thread(target=start_server, args=(rank, server_args)) + thread = threading.Thread(target=start_server, args=(rank,)) thread.start() self.server_threads.append(thread) @@ -101,6 +229,10 @@ class ExternalLBServerManager: for thread in self.server_threads: thread.join() + if start_errors: + start_error = start_errors[0] + raise RuntimeError("Failed to start external LB servers") from start_error + # Give servers additional time to fully initialize and coordinate time.sleep(2) @@ -184,12 +316,104 @@ def _assert_expert_parallel_config( assert parallel_config["enable_expert_parallel"] is expected +def _send_scale_command( + server: RemoteOpenAIServer, new_dp_size: int +) -> requests.Response: + return requests.post( + server.url_for("scale_elastic_ep"), + json={"new_data_parallel_size": new_dp_size}, + timeout=600, + ) + + +def _wait_for_scale_requests_to_start( + scale_requests: list[Future[requests.Response]], + timeout: float = 30, +) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + for request in scale_requests: + if request.done(): + response = request.result() + pytest.fail( + "Scale request finished before the new rank was launched: " + f"{response.status_code} {response.text}" + ) + + if all(request.running() for request in scale_requests): + return + time.sleep(0.1) + + pytest.fail("Timed out waiting for all scale requests to start") + + +@pytest.mark.distributed(num_gpus=3) +@pytest.mark.skipif( + TP_SIZE != 1 or current_platform.device_count() < 3, + reason="External Elastic EP scale-up requires three GPUs with TP_SIZE=1", +) +@pytest.mark.skipif( + os.getenv("VLLM_USE_RUST_FRONTEND") == "1", + reason="Elastic EP scaling is not supported by the Rust frontend", +) +def test_external_lb_elastic_ep_scale_up(default_server_args) -> None: + server_args = [ + *default_server_args, + "--enable-expert-parallel", + "--all2all-backend", + "allgather_reducescatter", + "--attention-backend", + "TRITON_MLA", + "--enable-elastic-ep", + "--elastic-ep-max-dp-size", + "3", + "--enable-eplb", + "--eplb-config", + '{"communicator":"pynccl","num_redundant_experts":0,"use_async":false}', + ] + server_manager = ExternalLBServerManager( + ELASTIC_EP_MODEL_NAME, + dp_size=2, + api_server_count=1, + base_server_args=server_args, + ) + + with server_manager: + existing_servers = [server for server, _ in server_manager.servers] + with ThreadPoolExecutor(max_workers=2) as executor: + scale_requests = [ + executor.submit(_send_scale_command, server, 3) + for server in existing_servers + ] + _wait_for_scale_requests_to_start(scale_requests) + server_manager.start_rank( + rank=2, + dp_size=3, + elastic_ep_scale_up_launch=True, + ) + + for request in scale_requests: + response = request.result(timeout=600) + assert response.status_code == 200, response.text + + for server, _ in server_manager.servers: + parallel_config = _get_parallel_config(server) + assert parallel_config["data_parallel_size"] == 3 + assert parallel_config["eplb_config"]["num_redundant_experts"] == 32 + with server.get_client() as client: + completion = client.completions.create( + model=ELASTIC_EP_MODEL_NAME, + prompt="Hello, my name is", + max_tokens=5, + temperature=0.0, + ) + assert completion.choices[0].finish_reason in ("length", "stop") + + def test_external_lb_server_info(server_manager): servers = server_manager.servers api_server_count = server_manager.api_server_count - _assert_expert_parallel_config(servers) - for i, (server, _) in enumerate(servers): print(f"Testing {i=}") diff --git a/tests/v1/e2e/spec_decode/mtp/qwen3_5/test_mtp.py b/tests/v1/e2e/spec_decode/mtp/qwen3_5/test_mtp.py index c61838a655..0feaa9efc6 100644 --- a/tests/v1/e2e/spec_decode/mtp/qwen3_5/test_mtp.py +++ b/tests/v1/e2e/spec_decode/mtp/qwen3_5/test_mtp.py @@ -5,7 +5,13 @@ import pytest from tests.utils import single_gpu_only from vllm import SamplingParams +from vllm.platforms import current_platform +from ...utils import ( + assert_request_outputs_match, + compute_acceptance_len, + get_test_prompts, +) from .._correctness import check_mtp_correctness @@ -37,3 +43,107 @@ def test_mtp_correctness( expected_accuracy_threshold, vllm_runner, ) + + +@single_gpu_only +def test_mtp_draft_lm_head_quantization( + monkeypatch: pytest.MonkeyPatch, + sampling_config: SamplingParams, + vllm_runner, +): + """A quantized draft copy of the shared lm_head only changes drafting: + greedy outputs match stock MTP and acceptance stays within 10%.""" + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + prompts = get_test_prompts(mm_enabled=False) + results = {} + # NVFP4 needs a CUDA W4A16 kernel. + quants = (None, "fp8", "nvfp4") if current_platform.is_cuda() else (None, "fp8") + for quant in quants: + with vllm_runner( + "Qwen/Qwen3.5-0.8B-Base", + block_size=None, + enable_chunked_prefill=None, + max_model_len=2048, + limit_mm_per_prompt={"image": 0, "video": 0}, + speculative_config={ + "method": "mtp", + "num_speculative_tokens": 3, + "draft_lm_head_quantization": quant, + }, + disable_log_stats=False, + ) as runner: + outputs = runner.llm.chat(prompts, sampling_config) + results[quant] = (outputs, compute_acceptance_len(runner.llm.get_metrics())) + + ref_outputs, ref_al = results.pop(None) + for quant, (outputs, al) in results.items(): + assert_request_outputs_match( + ref_outputs, + outputs, + # Not bit-exact: a different acceptance length changes the target's + # verify batch shapes, which can flip a near-tied greedy token. + required_matches=int(0.8 * len(ref_outputs)) + 1, + context=f"draft_lm_head_quantization={quant}", + ) + assert al >= 0.9 * ref_al, ( + f"{quant}: acceptance length {al:.3f} vs {ref_al:.3f}" + ) + + +@single_gpu_only +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="nvfp4 needs a CUDA W4A16 kernel" +) +def test_mtp_draft_lm_head_quantization_survives_sleep_and_reload( + monkeypatch: pytest.MonkeyPatch, + sampling_config: SamplingParams, + vllm_runner, +): + """The quantized draft head is derived from the target head: it must come + back intact after a level 1 sleep and be re-derived after a target weight + reload, so outputs and acceptance match the run before.""" + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + prompts = get_test_prompts(mm_enabled=False) + with vllm_runner( + "Qwen/Qwen3.5-0.8B-Base", + block_size=None, + enable_chunked_prefill=None, + max_model_len=2048, + limit_mm_per_prompt={"image": 0, "video": 0}, + speculative_config={ + "method": "mtp", + "num_speculative_tokens": 3, + "draft_lm_head_quantization": "nvfp4", + }, + enable_sleep_mode=True, + disable_log_stats=False, + ) as runner: + llm = runner.llm + + def run(): + before = llm.get_metrics() + outputs = llm.chat(prompts, sampling_config) + return outputs, compute_acceptance_len(llm.get_metrics(), before) + + ref_outputs, ref_al = run() + + llm.sleep(level=1) + llm.wake_up() + outputs, al = run() + assert_request_outputs_match( + ref_outputs, + outputs, + required_matches=int(0.8 * len(ref_outputs)) + 1, + context="after level 1 sleep", + ) + assert al == pytest.approx(ref_al, rel=0.03) + + llm.collective_rpc("reload_weights") + outputs, al = run() + assert_request_outputs_match( + ref_outputs, + outputs, + required_matches=int(0.8 * len(ref_outputs)) + 1, + context="after weight reload", + ) + assert al == pytest.approx(ref_al, rel=0.03) diff --git a/tests/v1/engine/test_core_engine_actor_manager.py b/tests/v1/engine/test_core_engine_actor_manager.py index 1b23356134..f8e2373d4c 100644 --- a/tests/v1/engine/test_core_engine_actor_manager.py +++ b/tests/v1/engine/test_core_engine_actor_manager.py @@ -141,7 +141,7 @@ def _make_addresses() -> EngineZmqAddresses: False, ), ({"data_parallel_backend": "mp", "data_parallel_rank_local": 0}, False), - ({"data_parallel_backend": "mp", "enable_elastic_ep": True}, False), + ({"data_parallel_backend": "mp", "enable_elastic_ep": True}, True), ({"data_parallel_backend": "mp", "data_parallel_size": 1}, False), ], ids=["mp", "ray", "non-master-node", "offline", "elastic-ep", "single-engine"], @@ -150,8 +150,9 @@ def test_coordination_store_held_only_by_the_online_dp_master( monkeypatch: pytest.MonkeyPatch, overrides: dict[str, Any], expect_store: bool ) -> None: """The DP master node holds a coordination store for both engine backends, - so engines there pick world-group ports at bind time. Every other - deployment keeps the pre-allocated ports and must not hold one.""" + so engines there pick world-group ports at bind time. Elastic EP also + creates its coordination store before startup. The remaining cases + below keep the pre-allocated ports and must not hold a store.""" class FakeManager: def __init__(self, **kwargs) -> None: diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index 9bcc9d91eb..b58aed8547 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -142,6 +142,9 @@ def test_mp_client_uses_env_timeout(monkeypatch: pytest.MonkeyPatch): def setsockopt(self, *_args, **_kwargs): pass + def getsockopt_string(self, _option): + return "inproc://dummy" + monkeypatch.setattr(core_client_mod.zmq.Socket, "shadow", lambda *_: ShadowSocket()) monkeypatch.setattr( core_client_mod, "make_zmq_socket", lambda *_, **__: DummySocket() @@ -210,6 +213,7 @@ def _make_dplb_client(num_engines: int = 3, client_count: int = 1) -> DPLBAsyncM client = object.__new__(DPLBAsyncMPClient) client.client_count = client_count client.reqs_in_flight = {} + client.streaming_reqs = set() client.engine_inflight = Counter() client.core_engines = [bytes([i, 0]) for i in range(num_engines)] client.lb_engines = [[0, 0, 0.0] for _ in range(num_engines)] @@ -407,11 +411,47 @@ def test_dplb_finished_requests_release_inflight(): assert req.request_id not in client.reqs_in_flight +@pytest.mark.parametrize("final_sent", [False, True]) +def test_dplb_streaming_session_abort_goes_to_all_engines(final_sent: bool): + """A streaming-input chunk in flight when its session finished may + re-create it on any engine, so an open session's abort reaches every + engine. Other aborts go only to the routed engine.""" + client = _make_dplb_client(num_engines=3) + client.resources = SimpleNamespace(engine_dead=False) + sent: dict[bytes, list[str]] = {} + + async def abort_requests(request_ids, engine): + sent[engine] = request_ids + + client._abort_requests = abort_requests + + chunk = make_request(SamplingParams(max_tokens=1)) + chunk.resumable = True + client.get_core_engine_for_request(chunk) + if final_sent: + chunk.resumable = False + client.get_core_engine_for_request(chunk) + # The engine finished the session, dropping its routing entry. + outputs = EngineCoreOutputs(finished_requests={chunk.request_id}) + asyncio.run(DPLBAsyncMPClient.process_engine_outputs(client, outputs)) + client.reqs_in_flight["live"] = client.core_engines[1] + + asyncio.run(client.abort_requests_async([chunk.request_id, "live"])) + + if final_sent: + expected = {client.core_engines[1]: ["live"]} + else: + expected = {engine: [chunk.request_id] for engine in client.core_engines} + expected[client.core_engines[1]].append("live") + assert sent == expected + assert not client.streaming_reqs + + @pytest.mark.parametrize( ("effective_size", "other_size"), [(None, None), (4224, 4224), (4224, 1056), (4224, None)], ) -def test_apply_ready_response_syncs_block_size(effective_size, other_size): +def test_apply_ready_response_syncs_engine_metadata(effective_size, other_size): import msgspec client = object.__new__(MPClient) @@ -419,8 +459,11 @@ def test_apply_ready_response_syncs_block_size(effective_size, other_size): client.vllm_config = SimpleNamespace( cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0), model_config=SimpleNamespace(max_model_len=8192), + parallel_config=SimpleNamespace(_coord_store_port=0), ) client.stats_update_address = None + client.coordinator_input_address = None + client.coordinator_output_address = None payload = msgspec.msgpack.encode( EngineCoreReadyResponse( @@ -441,6 +484,9 @@ def test_apply_ready_response_syncs_block_size(effective_size, other_size): instance_id="test-instance", supports_lora=False, max_loras=0, + coord_store_port=1234, + coordinator_input_address="tcp://127.0.0.1:1235", + coordinator_output_address="tcp://127.0.0.1:1236", ) ) fields = msgspec.msgpack.decode(payload) @@ -461,6 +507,9 @@ def test_apply_ready_response_syncs_block_size(effective_size, other_size): client._apply_ready_response(b"") assert cache_config.effective_attention_block_size is None + assert client.vllm_config.parallel_config._coord_store_port == 1234 + assert client.coordinator_input_address == "tcp://127.0.0.1:1235" + assert client.coordinator_output_address == "tcp://127.0.0.1:1236" def test_apply_ready_response_syncs_mamba_block_size(): @@ -471,6 +520,7 @@ def test_apply_ready_response_syncs_mamba_block_size(): client.vllm_config = SimpleNamespace( cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0), model_config=SimpleNamespace(max_model_len=8192), + parallel_config=SimpleNamespace(_coord_store_port=0), ) client.stats_update_address = None @@ -1541,6 +1591,7 @@ def test_apply_ready_response_retains_kv_event_sources(): client.vllm_config = SimpleNamespace( cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0), model_config=SimpleNamespace(max_model_len=8192), + parallel_config=SimpleNamespace(_coord_store_port=0), ) def ready(rank: int, config: KVEventsConfig | None) -> bytes: diff --git a/tests/v1/engine/test_output_processor.py b/tests/v1/engine/test_output_processor.py index 3c186a4e73..c07f747a27 100644 --- a/tests/v1/engine/test_output_processor.py +++ b/tests/v1/engine/test_output_processor.py @@ -42,6 +42,7 @@ from vllm.v1.engine.output_processor import ( OutputProcessor, RequestOutputCollector, RequestState, + StreamingInputOutputCollector, ) from vllm.v1.engine.parallel_sampling import ParentRequest from vllm.v1.metrics.loggers import PrometheusStatLogger @@ -1642,6 +1643,137 @@ async def test_cumulative_output_collector_n(): assert third[0].text == "c" +def _make_boundary_output( + token_id: int, finish_reason: str | None, finished: bool = False +) -> RequestOutput: + return RequestOutput( + request_id="my-request-id", + prompt=None, + prompt_token_ids=[1, 2, 3], + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=0, + text="a", + token_ids=[token_id], + cumulative_logprob=None, + logprobs=None, + finish_reason=finish_reason, + ) + ], + finished=finished, + ) + + +@pytest.mark.asyncio +async def test_streaming_input_collector_preserves_chunk_boundary(): + """A pending chunk-boundary marker (finish_reason on a non-finished + output) must not be erased by merging the next chunk's output into it.""" + collector = StreamingInputOutputCollector( + RequestOutputKind.DELTA, request_id="my-request-id-int" + ) + + collector.put(_make_boundary_output(0, "stop")) + collector.put(_make_boundary_output(1, None)) + collector.put(_make_boundary_output(2, None)) + + boundary = await collector.get() + assert boundary.outputs[0].finish_reason == "stop" + assert boundary.outputs[0].token_ids == [0] + + following = await collector.get() + assert following.outputs[0].finish_reason is None + assert following.outputs[0].token_ids == [1, 2] + + assert collector.output is None + assert not collector._backlog + assert not collector.ready.is_set() + + +def _make_streaming_input_request( + request_id: str, prompt_token_ids: list[int] +) -> EngineCoreRequest: + """One chunk of a streaming-input (resumable) request.""" + return EngineCoreRequest( + request_id=request_id, + prompt_token_ids=prompt_token_ids, + mm_features=None, + sampling_params=SamplingParams( + detokenize=False, output_kind=RequestOutputKind.DELTA + ), + pooling_params=None, + arrival_time=0, + lora_request=None, + cache_salt=None, + data_parallel_rank=None, + resumable=True, + external_req_id=f"{request_id}-ext", + ) + + +def _add_streaming_request_with_queued_chunk( + processor: OutputProcessor, request_id: str +) -> RequestOutputCollector: + queue = RequestOutputCollector(RequestOutputKind.DELTA, request_id=request_id) + processor.add_request( + _make_streaming_input_request(request_id, [1, 2, 3]), None, queue=queue + ) + processor.add_request(_make_streaming_input_request(request_id, [7, 8]), None) + return queue + + +def test_streaming_input_stop_is_chunk_boundary(): + """STOP is a per-chunk boundary: finished is forced False, the queued + chunk is applied, and the request state stays alive.""" + processor = OutputProcessor(None, log_stats=False) + queue = _add_streaming_request_with_queued_chunk(processor, "stream-req") + + processed = processor.process_outputs( + [ + EngineCoreOutput( + request_id="stream-req", + new_token_ids=[4], + finish_reason=FinishReason.STOP, + ) + ] + ) + + assert processed.reqs_to_abort == [] + assert processor.get_num_unfinished_requests() == 1 + output = queue.get_nowait() + assert output is not None + assert output.finished is False + assert output.outputs[0].finish_reason == "stop" + req_state = processor.request_states["stream-req"] + assert req_state.prompt_token_ids is not None + assert req_state.prompt_token_ids[-2:] == [7, 8] + + +@pytest.mark.parametrize("finish_reason", [FinishReason.ABORT, FinishReason.ERROR]) +def test_streaming_input_abort_error_is_terminal(finish_reason: FinishReason): + """ABORT/ERROR is terminal, not a chunk boundary: finished stays True, + the state is freed, and an abort is routed back to the engine core.""" + processor = OutputProcessor(None, log_stats=False) + queue = _add_streaming_request_with_queued_chunk(processor, "stream-req") + + processed = processor.process_outputs( + [ + EngineCoreOutput( + request_id="stream-req", + new_token_ids=[], + finish_reason=finish_reason, + ) + ] + ) + + assert processor.get_num_unfinished_requests() == 0 + assert processed.reqs_to_abort == ["stream-req"] + output = queue.get_nowait() + assert output is not None + assert output.finished is True + assert output.outputs[0].finish_reason == str(finish_reason) + + @pytest.mark.parametrize("runner", ["generate", "pooling"]) @pytest.mark.parametrize("abort_by", ["internal", "external"]) def test_abort_requests(runner: str, abort_by: str, dummy_test_vectors): diff --git a/tests/v1/engine/test_startup_watch_processes.py b/tests/v1/engine/test_startup_watch_processes.py index 19b2d7bca8..bf325a6108 100644 --- a/tests/v1/engine/test_startup_watch_processes.py +++ b/tests/v1/engine/test_startup_watch_processes.py @@ -281,6 +281,7 @@ class _FinishedProcess: def test_wait_for_engine_startup_reports_watched_process_exit(): + """Update the startup test argument so it reaches frontend failure handling.""" ctx = zmq.Context() handshake_socket = ctx.socket(zmq.ROUTER) recv, send = connection.Pipe(duplex=False) diff --git a/tests/v1/fault_tolerance/test_fault_tolerance_e2e.py b/tests/v1/fault_tolerance/test_fault_tolerance_e2e.py index f6d15343b5..27e45f6635 100644 --- a/tests/v1/fault_tolerance/test_fault_tolerance_e2e.py +++ b/tests/v1/fault_tolerance/test_fault_tolerance_e2e.py @@ -190,11 +190,13 @@ def _apply_ft(server, instruction: str, params: dict | None = None) -> dict: def _kill_worker_process(server) -> None: """SIGKILL only the worker proc, leaving EngineCore and API server alive.""" - workers = [ - p - for p in psutil.Process(server.proc.pid).children(recursive=True) - if "Worker" in " ".join(p.cmdline()) - ] + workers = [] + for p in psutil.Process(server.proc.pid).children(recursive=True): + # Skip descendants that already exited: cmdline() raises ZombieProcess + # if unreaped, NoSuchProcess if reaped. A dead worker still fails below. + with contextlib.suppress(psutil.ZombieProcess, psutil.NoSuchProcess): + if "Worker" in " ".join(p.cmdline()): + workers.append(p) assert len(workers) == 1, f"expected 1 worker proc, found: {workers}" workers[0].kill() diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index db7dbdab09..f90de6297a 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -73,6 +73,7 @@ from vllm.v1.kv_cache_interface import ( KVCacheConfig, KVCacheGroupSpec, KVCacheLayout, + KVQuantMode, MLAAttentionSpec, compute_layer_kv_cache_shape_bytes, ) @@ -1551,6 +1552,95 @@ class TestNixlHandshake: with pytest.raises(AssertionError): worker.add_remote_agent(bad_meta, remote_tp_size=2) + @pytest.mark.parametrize( + ("kv_quant_mode", "local_tp", "remote_tp", "num_kv_heads", "rejected"), + [ + (KVQuantMode.NVFP4, 2, 1, 8, True), + (KVQuantMode.NVFP4, 1, 2, 8, True), + (KVQuantMode.NVFP4, 2, 2, 8, False), + # Both sides hold one replicated head per rank: no head split. + (KVQuantMode.NVFP4, 16, 8, 8, False), + (KVQuantMode.NONE, 2, 1, 8, False), + ], + ) + @patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", + FakeNixlWrapper, + ) + def test_handshake_rejects_nvfp4_kv_head_split( + self, + default_vllm_config, + dist_init, + kv_quant_mode, + local_tp, + remote_tp, + num_kv_heads, + rejected, + ): + """NVFP4 pages keep all heads' data before all heads' scales, so + heterogeneous-TP head splitting would silently corrupt the KV. + """ + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) + kv_cache_config = KVCacheConfig( + num_blocks=1, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer0"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=max(1, num_kv_heads // local_tp), + head_size=16, + dtype=torch.uint8, + kv_quant_mode=kv_quant_mode, + ), + ) + ], + ) + with patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.get_tensor_model_parallel_world_size", # noqa: E501 + return_value=local_tp, + ): + worker = FakeNixlConnectorWorker( + vllm_config, + "engine", + hand_shake_latency=0, + kv_cache_config=kv_cache_config, + ) + worker.transfer_topo.total_num_kv_heads = num_kv_heads + worker.transfer_topo.local_physical_heads = max(1, num_kv_heads // local_tp) + + local_len = 4096 * worker.block_size + remote_len = ( + local_len + * max(1, num_kv_heads // remote_tp) + // worker.transfer_topo.local_physical_heads + ) + worker.block_len_per_layer = [local_len] + worker.num_blocks = 1 + worker.dst_num_blocks[worker.engine_id] = worker.num_blocks + meta = NixlAgentMetadata( + engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID, + agent_metadata=FakeNixlWrapper.AGENT_METADATA, + kv_caches_base_addr=[0], + device_id=0, + num_blocks=1, + block_lens=[remote_len], + block_strides=[remote_len], + kv_cache_layout=worker.kv_cache_layout, + block_size=worker.block_size, + ssm_sizes=(0, 0), + attn_backend_name=worker.backend_name, + physical_blocks_per_logical_kv_block=1, + ) + + if rejected: + with pytest.raises(RuntimeError, match="NVFP4"): + worker.add_remote_agent(meta, remote_tp_size=remote_tp) + else: + worker.add_remote_agent(meta, remote_tp_size=remote_tp) + # NOTE: resource cleanup in mp backend is a bit finicky, so the order in which # we put here is important. First run ray, it will clean up the resources, then diff --git a/tests/v1/kv_offload/cpu/test_shared_offload_region.py b/tests/v1/kv_offload/cpu/test_shared_offload_region.py index 21ef6d2761..6fc557c3a8 100644 --- a/tests/v1/kv_offload/cpu/test_shared_offload_region.py +++ b/tests/v1/kv_offload/cpu/test_shared_offload_region.py @@ -653,6 +653,41 @@ def test_madvise_unexpected_oserror_propagates(iid, monkeypatch): assert exc_info.value.errno == errno.EIO +def test_segmented_region_maps_files_back_to_back(iid, monkeypatch): + """Rows split over three shm files mapped back to back: each worker sees + the other's writes in every file, and the creator removes every file.""" + monkeypatch.setattr(region_module, "_MIN_SEGMENT_BYTES", PAGE_SIZE) + regions = [ + SharedOffloadRegion( + engine_id=iid, + num_chunks=5, + rank=rank, + kv_bytes_per_chunk=2 * PAGE_SIZE, + cpu_page_size=PAGE_SIZE, + num_segments=3, + ) + for rank in range(2) + ] + segments = regions[0]._segments + try: + assert [length // PAGE_SIZE for _, _, length in segments] == [4, 4, 2] + for rank, region in enumerate(regions): + ranges = list(region.iter_worker_slot_ranges()) + first_row = rank * 5 // 2 + assert ranges[0] == ((first_row * 2 + rank) * PAGE_SIZE, PAGE_SIZE) + assert sorted(ranges) == [ + ((row * 2 + rank) * PAGE_SIZE, PAGE_SIZE) for row in range(5) + ] + region.create_next_worker_view(PAGE_SIZE).fill_(rank + 1) + rows = regions[0].base_tensor.view(5, 2, PAGE_SIZE) + assert (rows[:, 0] == 1).all() and (rows[:, 1] == 2).all() + del rows + finally: + for region in regions: + region.cleanup() + assert not any(os.path.exists(path) for path, _, _ in segments) + + # --------------------------------------------------------------------------- # Multi-worker race — concurrent construction # --------------------------------------------------------------------------- @@ -899,6 +934,7 @@ def _unmapped_registration_region(num_chunks, slot_size, num_workers): region.cpu_page_size = slot_size region._row_stride = num_workers * slot_size region.total_size_bytes = region.num_chunks * region._row_stride + region._segments = [("unused", 0, region.total_size_bytes)] region._base = SimpleNamespace(data_ptr=lambda: 0x100000) region.pinned_addresses = [] region.is_pinned = False diff --git a/tests/v1/spec_decode/test_dspark_topk.py b/tests/v1/spec_decode/test_dspark_topk.py index 2de1ff052a..164ef14ff4 100644 --- a/tests/v1/spec_decode/test_dspark_topk.py +++ b/tests/v1/spec_decode/test_dspark_topk.py @@ -15,6 +15,8 @@ from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ) from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead from vllm.platforms import current_platform +from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample +from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator def _markov_head(weight: torch.Tensor) -> DSparkMarkovHead: @@ -70,6 +72,46 @@ def test_gathered_markov_bias_matches_dense_at_full_vocab(): torch.testing.assert_close(result, expected) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("use_fp64", [False, True]) +def test_dspark_revisited_position_uses_step_seed(use_fp64): + """A cached draft at a revisited position must use its current step.""" + device = torch.device("cuda") + batch = 4096 + spec = DSparkSpeculator.__new__(DSparkSpeculator) + spec.acceptance_estimator = None + spec._d2t_scatter_index = None + spec.draft_watermarker = None + spec.use_fp64_gumbel = use_fp64 + spec.temperature = torch.ones(batch, device=device) + spec.seeds = torch.arange(batch, dtype=torch.int64, device=device) + spec._step_cols = torch.arange(2, dtype=torch.int32, device=device) + spec.draft_logits = torch.empty(batch, 2, 2, device=device) + logits = torch.tensor([0.7, 0.3], device=device).log().expand(batch, -1) + mapping = torch.arange(batch, device=device) + positions = torch.full((batch,), 22, dtype=torch.int64, device=device) + + first = spec._sample_logits(logits, mapping, positions, 1) + repeated = spec._sample_logits(logits, mapping, positions, 1) + next_round = spec._sample_logits(logits, mapping, positions, 0) + torch.testing.assert_close(first, repeated) + expected = gumbel_sample( + logits, + mapping, + spec.temperature, + spec.seeds, + positions - 1, + apply_temperature=True, + is_drafting=True, + logits_cache=spec.draft_logits, + logits_cache_col=spec._step_cols[0], + use_fp64=use_fp64, + ) + torch.testing.assert_close(next_round, expected) + mismatch_rate = (first != next_round).float().mean().item() + assert 0.35 < mismatch_rate < 0.5 + + @pytest.mark.parametrize("cuda_alike_platform", [False, True]) def test_gathered_markov_bias_dequantizes_selected_w4a16_rows( monkeypatch, cuda_alike_platform diff --git a/tests/v1/spec_decode/test_lilicorr.py b/tests/v1/spec_decode/test_lilicorr.py index 46dbaf41af..95fac9f834 100644 --- a/tests/v1/spec_decode/test_lilicorr.py +++ b/tests/v1/spec_decode/test_lilicorr.py @@ -22,6 +22,7 @@ from vllm.model_executor.layers import logits_processor from vllm.model_executor.layers.linear import LinearBase from vllm.model_executor.models.lilicorr import LiLiCorrHead from vllm.transformers_utils.configs.eagle import EAGLEConfig +from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import CandidateSampler from vllm.v1.worker.gpu.spec_decode.lilicorr.speculator import LiLiCorrSpeculator @@ -500,6 +501,60 @@ def test_candidate_walk_preserves_conditional_scores_and_padding(probabilistic): assert cache[1].isneginf().all() +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("use_fp64", [False, True]) +def test_candidate_walk_revisited_position_matches_step_seed(use_fp64): + """The candidate walk must use the same step key as ordinary draft sampling.""" + device = torch.device("cuda") + batch, steps, top_k = 4096, 2, 2 + sampler = CandidateSampler(batch, steps, top_k, device) + candidates = torch.arange(top_k, device=device).expand(batch, steps, top_k) + logits = torch.tensor([0.7, 0.3], device=device).log() + scores = logits.expand(batch, steps, top_k, top_k) + mapping = torch.arange(batch, device=device) + temperature = torch.ones(batch, device=device) + seeds = mapping.to(torch.int64) + positions = torch.arange(21, 21 + steps, device=device).repeat(batch) + tokens = torch.empty(batch, steps, dtype=torch.int64, device=device) + cache = torch.empty(batch, steps, top_k, device=device) + first_round = None + for start in (21, 22): + positions.view(batch, steps).copy_( + torch.arange(start, start + steps, device=device) + ) + sampler.sample( + candidates, + scores, + batch, + positions, + mapping.repeat_interleave(steps), + temperature, + seeds, + tokens, + cache, + use_fp64, + ) + for step in range(steps): + expected = gumbel_sample( + logits.expand(batch, -1), + mapping, + temperature, + seeds, + positions.view(batch, steps)[:, step] - 1, + apply_temperature=True, + is_drafting=True, + logits_cache=cache, + logits_cache_col=torch.tensor(step, device=device), + use_fp64=use_fp64, + ) + torch.testing.assert_close(tokens[:, step], expected) + if first_round is None: + first_round = tokens.clone() + else: + mismatch_rate = (first_round[:, 1] != tokens[:, 0]).float().mean().item() + assert 0.35 < mismatch_rate < 0.5 + + @pytest.mark.parametrize( ("convolution", "mismatch", "quantized"), [ diff --git a/tests/v1/spec_decode/test_mtp_structured_output.py b/tests/v1/spec_decode/test_mtp_structured_output.py index 8a72765db1..72fa3e7ad7 100644 --- a/tests/v1/spec_decode/test_mtp_structured_output.py +++ b/tests/v1/spec_decode/test_mtp_structured_output.py @@ -114,6 +114,10 @@ class _EngineReasonerStub(ParserEngineReasoningAdapter): def reasoning_end_token_ids(self): return self._end_token_ids + @property + def implicit_reasoning_end_token_ids(self): + return frozenset() + def find_reasoning_end(self, token_ids): self.windows.append(list(token_ids)) for offset, token in enumerate(token_ids): diff --git a/tests/v1/spec_decode/test_rejection_sampler_utils.py b/tests/v1/spec_decode/test_rejection_sampler_utils.py index f635a66a46..ab94ec54d2 100644 --- a/tests/v1/spec_decode/test_rejection_sampler_utils.py +++ b/tests/v1/spec_decode/test_rejection_sampler_utils.py @@ -557,6 +557,79 @@ def test_gumbel_drafted_rejection_sample_is_unbiased(num_speculative_steps: int) ) +@pytest.mark.parametrize("use_block_verification", [False, True]) +@pytest.mark.parametrize("probabilistic_draft", [False, True]) +def test_multiround_rejection_sample_is_unbiased( + use_block_verification: bool, probabilistic_draft: bool +): + """Uncommitted positions must receive fresh noise on later verification rounds.""" + torch.manual_seed(42) + num_trials, num_tokens, k = 32768, 6, 2 + device = "cuda" + target_probs = torch.tensor( + [0.3, 0.7] if probabilistic_draft else [0.5, 0.5], device=device + ) + draft_logits = torch.tensor( + [0.7, 0.3] if probabilistic_draft else [1.0, 0.0], device=device + ).log() + inputs = _build_rejection_sample_inputs( + target_probs.log(), draft_logits, k, 1.0, num_trials + ) + seeds = inputs["seed"] + output = torch.full((num_trials, num_tokens + k), -1, device=device) + output[:, 0] = torch.multinomial(target_probs, num_trials, replacement=True) + lengths = torch.ones(num_trials, dtype=torch.int64, device=device) + accepted = 0 + while (lengths < num_tokens).any(): + active = (lengths < num_tokens).nonzero().flatten() + n = active.numel() + starts = lengths[active] + local_positions = torch.arange(k + 1, device=device).repeat(n) + positions = (starts[:, None] + torch.arange(k + 1, device=device)).flatten() + draft_sampled = torch.zeros((n, k + 1), dtype=torch.int64, device=device) + draft_sampled[:, 0] = output[active, starts - 1] + logits_cache = inputs["draft_logits"] if probabilistic_draft else None + if probabilistic_draft: + for step in range(k): + draft_sampled[:, step + 1] = gumbel_sample( + draft_logits.expand(n, -1), + active, + inputs["temperature"], + seeds, + starts + step, + apply_temperature=True, + is_drafting=True, + logits_cache=logits_cache, + logits_cache_col=torch.tensor(step, device=device), + ) + sampled, counts = rejection_sample( + target_logits=target_probs.log().expand(n * (k + 1), -1).contiguous(), + draft_logits=logits_cache, + draft_sampled=draft_sampled.flatten(), + cu_num_logits=torch.arange(n + 1, dtype=torch.int32, device=device) + * (k + 1), + pos=positions, + idx_mapping=active, + expanded_idx_mapping=active.repeat_interleave(k + 1), + expanded_local_pos=local_positions, + temperature=inputs["temperature"], + seed=seeds, + num_speculative_steps=k, + use_block_verification=use_block_verification, + ) + for column in range(k + 1): + valid = counts > column + output[active[valid], starts[valid] + column] = sampled[valid, column] + lengths[active] += counts + accepted += (counts - 1).sum().item() + + assert accepted > num_trials, "The test must exercise accepted draft tokens." + for position in range(num_tokens): + _assert_distribution_match( + output[:, position], target_probs, device, label=f"position {position}" + ) + + @pytest.mark.parametrize("num_speculative_steps", [1, 3]) def test_greedy_rejection_sample(num_speculative_steps: int): """Verify that greedy (temperature=0) always outputs the target argmax diff --git a/tests/v1/streaming_input/test_async_llm_streaming.py b/tests/v1/streaming_input/test_async_llm_streaming.py index 21a49024ce..9dac5ee652 100644 --- a/tests/v1/streaming_input/test_async_llm_streaming.py +++ b/tests/v1/streaming_input/test_async_llm_streaming.py @@ -3,6 +3,7 @@ import asyncio from collections.abc import AsyncGenerator +from unittest import mock from unittest.mock import AsyncMock, MagicMock import pytest @@ -198,3 +199,52 @@ async def test_empty_input_stream_finishes_without_engine_request(): assert queue.get_nowait() is STREAM_FINISHED llm._add_request.assert_not_awaited() + + +def _make_streaming_llm(num_chunks: int, has_request: bool = True): + llm = AsyncLLM.__new__(AsyncLLM) + llm._validate_streaming_input_sampling_params = MagicMock() + llm.get_supported_tasks = AsyncMock(return_value=("generate",)) + llm._run_output_handler = MagicMock() + llm._add_request = AsyncMock() + llm.model_config = MagicMock() + llm.output_processor = MagicMock() + llm.output_processor.has_request.return_value = has_request + + # The first processed request is the final (end-of-stream) request. + reqs = [MagicMock(prompt_embeds=None) for _ in range(num_chunks + 1)] + reqs[0].request_id = "r1-int" + llm.input_processor = MagicMock() + llm.input_processor.process_inputs.side_effect = reqs + return llm, reqs + + +async def _stream_chunks(llm: AsyncLLM, num_chunks: int) -> None: + async def chunks(): + for i in range(num_chunks): + yield StreamingInput(prompt=f"frame-{i}") + + with mock.patch( + "vllm.v1.engine.async_llm.extract_prompt_components", + return_value=("", None, None), + ): + queue = await llm._add_streaming_input_request( + "r1", chunks(), SamplingParams(max_tokens=4) + ) + task = queue._input_stream_task + assert task is not None + await task + + +@pytest.mark.asyncio +@pytest.mark.parametrize("session_live", [True, False]) +async def test_handle_inputs_stops_after_session_finished(session_live: bool): + """Once the session's request state is gone (aborted or errored), no + further chunks or final request are sent: they would recreate a request + state that is never cleaned up.""" + llm, reqs = _make_streaming_llm(num_chunks=2, has_request=session_live) + await _stream_chunks(llm, num_chunks=2) + + sent = [call.args[0] for call in llm._add_request.await_args_list] + # Both chunks then the final request, or only the first chunk. + assert sent == ([reqs[1], reqs[2], reqs[0]] if session_live else [reqs[1]]) diff --git a/tests/v1/streaming_input/test_scheduler_streaming.py b/tests/v1/streaming_input/test_scheduler_streaming.py index 1f97d8a1c0..c6df585a58 100644 --- a/tests/v1/streaming_input/test_scheduler_streaming.py +++ b/tests/v1/streaming_input/test_scheduler_streaming.py @@ -5,9 +5,11 @@ import unittest from unittest.mock import MagicMock import numpy as np +import pytest import torch import vllm.v1.core.kv_cache_utils as kv_cache_utils +from tests.v1.core.utils import create_scheduler as create_spec_scheduler from vllm.config import DeviceConfig, VllmConfig from vllm.multimodal.inputs import ( MultiModalFeatureSpec, @@ -28,7 +30,7 @@ from vllm.v1.kv_cache_interface import ( KVCacheConfig, KVCacheGroupSpec, ) -from vllm.v1.outputs import LogprobsLists, ModelRunnerOutput +from vllm.v1.outputs import DraftTokenIds, LogprobsLists, ModelRunnerOutput from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm.v1.structured_output import StructuredOutputManager @@ -676,3 +678,41 @@ class TestStreamingScheduler(unittest.TestCase): cached_state_cycle1["prompt_token_ids"] is not cached_state_cycle3["prompt_token_ids"] ), "Cached states from different cycles should be independent objects." + + +def _model_output(request_id: str, token_ids: list[int]) -> ModelRunnerOutput: + return ModelRunnerOutput( + req_ids=[request_id], + req_id_to_index={request_id: 0}, + sampled_token_ids=[token_ids], + logprobs=None, + prompt_logprobs_dict={request_id: None}, + pooler_output=[], + ) + + +@pytest.mark.parametrize("async_scheduling", [False, True]) +def test_next_chunk_fully_computed_after_stop_on_accepted_draft(async_scheduling): + """A stop on an accepted draft trims the step's later tokens. The next + chunk must still be computed in full rather than landing on positions + whose KV belongs to the trimmed drafts.""" + scheduler = create_spec_scheduler( + num_speculative_tokens=2, + async_scheduling=async_scheduling, + speculative_method="ngram_gpu" if async_scheduling else None, + ) + session = DummyRequest("session", prompt_token_ids=[1, 2, 3], max_tokens=2) + scheduler.add_request(session) + scheduler.update_from_output(scheduler.schedule(), _model_output("session", [10])) + scheduler.update_draft_token_ids(DraftTokenIds(["session"], [[11, 12]])) + # Both drafts are accepted, but max_tokens stops at 11. + scheduler.update_from_output( + scheduler.schedule(), _model_output("session", [11, 12, 13]) + ) + assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ + + scheduler.add_request(DummyRequest("session", prompt_token_ids=[7, 8])) + output = scheduler.schedule() + + assert list(session.prompt_token_ids) == [1, 2, 3, 10, 7, 8] + assert output.num_scheduled_tokens["session"] == 2 diff --git a/tests/v1/structured_output/test_structured_output_manager.py b/tests/v1/structured_output/test_structured_output_manager.py index 1d1e7b78df..0cb051f0fd 100644 --- a/tests/v1/structured_output/test_structured_output_manager.py +++ b/tests/v1/structured_output/test_structured_output_manager.py @@ -6,6 +6,7 @@ from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass from threading import Event +import numpy as np import pytest from transformers import AutoTokenizer @@ -72,20 +73,26 @@ class MockReasoner: class MockEngineReasoner(ParserEngineReasoningAdapter): - """Ends reasoning implicitly on ``marker``, which is also content.""" + """Ends reasoning implicitly on ``marker``, which is also content, and + explicitly on ``end_marker`` when one is given.""" - def __init__(self, tokenizer, marker: int): + def __init__(self, tokenizer, marker: int, end_marker: int | None = None): self.marker = marker + self.end_marker = end_marker @property def reasoning_end_token_ids(self): + return frozenset({self.marker, self.end_marker} - {None}) + + @property + def implicit_reasoning_end_token_ids(self): return frozenset({self.marker}) def find_reasoning_end(self, token_ids): - ids = list(token_ids) - if self.marker in ids: - return ReasoningEnd(ids.index(self.marker), True) - return ReasoningEnd(len(ids), False) + for offset, token in enumerate(token_ids): + if token in self.reasoning_end_token_ids: + return ReasoningEnd(offset, token == self.marker) + return ReasoningEnd(len(token_ids), False) def _single_token(tokenizer, text: str) -> int: @@ -122,6 +129,7 @@ def _build_harness( enable_in_reasoning: bool = False, reasoning_parser_kwargs: dict | None = None, structured_outputs: StructuredOutputsParams | None = None, + num_spec_tokens: int = NUM_SPEC_TOKENS, ) -> tuple[StructuredOutputManager, Request]: vllm_config = VllmConfig( model_config=ModelConfig(tokenizer=TOKENIZER), @@ -132,8 +140,10 @@ def _build_harness( ), speculative_config=SpeculativeConfig( model="[ngram]", - num_speculative_tokens=NUM_SPEC_TOKENS, - ), + num_speculative_tokens=num_spec_tokens, + ) + if num_spec_tokens + else None, ) manager = StructuredOutputManager(vllm_config) if use_reasoner: @@ -171,8 +181,20 @@ def _build_harness( return manager, request -def _row_pattern(bitmask) -> str: - return "".join("U" if (row == -1).all() else "C" for row in bitmask) +def _row_pattern(bitmask, masked_ids=()) -> str: + """U: unconstrained; R: all but `masked_ids` allowed; C: grammar rows.""" + reasoning_row = np.full(bitmask.shape[1], -1, dtype=np.int32) + for token_id in masked_ids: + reasoning_row.view(np.uint32)[token_id >> 5] &= ~np.uint32(1 << (token_id & 31)) + + def kind(row) -> str: + if (row == -1).all(): + return "U" + if masked_ids and (row == reasoning_row).all(): + return "R" + return "C" + + return "".join(kind(row) for row in bitmask) def _run_real_flow( @@ -183,6 +205,7 @@ def _run_real_flow( expected_row_pattern: str, expected_reasoning: bool | None, expect_terminated: bool, + masked_ids: list[int] | None = None, ) -> None: assert len(expected_validated) <= len(raw_drafts) @@ -197,7 +220,7 @@ def _run_real_flow( ) assert bitmask is not None assert bitmask.shape[0] == len(raw_drafts) + 1 - assert _row_pattern(bitmask) == expected_row_pattern + assert _row_pattern(bitmask, masked_ids or ()) == expected_row_pattern # grammar_bitmask() must rollback any speculative state it advanced. assert manager.validate_tokens(request, list(raw_drafts)) == expected_validated @@ -375,7 +398,8 @@ FLOW_CASES = [ FlowCase( raw_drafts=(" ", IMPLICIT_THINK_END, "{", "}"), expected_validated=(" ",), - expected_row_pattern="UUUUU", + # The grammar cannot start with the marker, so reasoning rows mask it. + expected_row_pattern="RRUUU", expected_reasoning=False, expect_terminated=False, engine_reasoner=True, @@ -421,9 +445,93 @@ def test_real_flow( expected_row_pattern=case.expected_row_pattern, expected_reasoning=case.expected_reasoning, expect_terminated=case.expect_terminated, + masked_ids=[reasoner_kwargs["marker"]] if case.engine_reasoner else None, ) +@pytest.mark.parametrize("backend", BACKENDS) +@pytest.mark.parametrize( + ("structured_outputs", "raw_drafts", "expected_row_pattern"), + [ + # JSON cannot start with the marker: reasoning rows mask it and keep + # every other token, including the explicit end, allowed. + pytest.param( + StructuredOutputsParams(json=JSON_SCHEMA), (" ", " "), "RRR", id="json" + ), + # The grammar can start with the marker: it stays allowed and, once + # sampled, is the first token the grammar consumes. + pytest.param( + StructuredOutputsParams(regex="z[0-9]+"), + (" ", IMPLICIT_THINK_END, "1"), + "UUCC", + id="regex_starts_with_marker", + ), + ], +) +def test_reasoning_rows_mask_grammar_invalid_implicit_end( + tokenizer, + backend: StructuredOutputsBackend, + structured_outputs: StructuredOutputsParams, + raw_drafts: tuple[str, ...], + expected_row_pattern: str, +): + marker = _single_token(tokenizer, IMPLICIT_THINK_END) + manager, request = _build_harness( + tokenizer, + backend, + reasoning_ended=False, + reasoning_parser_kwargs={ + "marker": marker, + "end_marker": _single_token(tokenizer, THINK_END), + }, + structured_outputs=structured_outputs, + ) + manager.reasoner_cls = MockEngineReasoner + drafts = _to_token_ids(tokenizer, raw_drafts) + assert manager.validate_tokens(request, drafts) == drafts + + bitmask = manager.grammar_bitmask( + requests={request.request_id: request}, + structured_output_request_ids=[request.request_id], + scheduled_spec_decode_tokens={request.request_id: drafts}, + ) + assert bitmask is not None + assert _row_pattern(bitmask, [marker]) == expected_row_pattern + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_parallel_fill_masks_grammar_invalid_implicit_end( + tokenizer, backend: StructuredOutputsBackend +): + marker, end = _to_token_ids(tokenizer, (IMPLICIT_THINK_END, THINK_END)) + manager, request = _build_harness( + tokenizer, + backend, + reasoning_ended=False, + reasoning_parser_kwargs={"marker": marker, "end_marker": end}, + num_spec_tokens=0, + ) + manager.reasoner_cls = MockEngineReasoner + + def row_pattern() -> str: + bitmask = manager.grammar_bitmask( + requests={request.request_id: request}, + structured_output_request_ids=[request.request_id], + scheduled_spec_decode_tokens={}, + ) + assert bitmask is not None + return _row_pattern(bitmask, [marker]) + + manager.fill_bitmask_parallel_threshold = 0 + manager.fill_bitmask_parallel_batch_size = 16 + with ThreadPoolExecutor(max_workers=1) as executor: + manager.executor_for_fillmask = executor + assert row_pattern() == "R" + request.append_output_token_ids([end]) + assert manager.accept_tokens(request, [end]) + assert row_pattern() == "C" + + @pytest.mark.parametrize("backend", BACKENDS) @pytest.mark.parametrize( ("use_reasoner", "reasoning_ended", "enable_in_reasoning"), diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index 8d30403f30..0f1537403e 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/tests/v1/worker/test_attn_utils.py @@ -9,12 +9,14 @@ never addressed by the logical view. from types import SimpleNamespace from typing import Any +from unittest.mock import MagicMock import numpy as np import pytest import torch import vllm.v1.hisparse.binding as attn_utils_module +import vllm.v1.hisparse.runtime as hisparse_runtime_module from tests.v1.attention.utils import ( BatchSpec, create_common_attn_metadata, @@ -403,49 +405,63 @@ class _FakeSharedHostRegion: self.cleanup_calls += 1 -def test_profiling_cleanup_releases_tp_shared_region_once(monkeypatch): - """TP-shared profiling pools must use region-aware chunk cleanup.""" - region = _FakeSharedHostRegion() +def test_profiling_cleanup_keeps_host_pool_for_real_kv_cache(monkeypatch): + """The real KV cache must reuse the profiling host pool, not pin it again.""" + pinned = torch.empty(8, dtype=torch.int8) + allocate_pinned = MagicMock(return_value=(pinned, pinned)) + monkeypatch.setattr( + hisparse_runtime_module, "allocate_pinned_host_pool", allocate_pinned + ) + monkeypatch.setattr( + hisparse_runtime_module, "check_hisparse_host_memory", lambda size: None + ) + vllm_config = SimpleNamespace(instance_id="profiling-reuse") + kv_cache_config = SimpleNamespace(hisparse_shared_host_pool=False) + pool = attn_utils_module.get_hisparse_host_pool(vllm_config, kv_cache_config) + backing = pool.allocate(8) runtime = SimpleNamespace( - _host_cache=object(), - registered_host_pool=region.base_tensor, - hot_backing=object(), - shared_host_region=region, + _host_cache=backing, registered_host_pool=pinned, hot_backing=object() ) forward_context = { - "layer": SimpleNamespace( - hisparse_cache=SimpleNamespace(runtime=runtime), - ) + "layer": SimpleNamespace(hisparse_cache=SimpleNamespace(runtime=runtime)) } - released = [] - - def release_pinned_state(runtimes, pinned_host_pools, shared_host_region): - released.append((runtimes, pinned_host_pools, shared_host_region)) - - monkeypatch.setattr( - attn_utils_module, - "release_pinned_state", - release_pinned_state, - ) attn_utils_module.release_hisparse_profiling_cache(forward_context) - assert released == [([runtime], [], region)] + assert not hasattr(runtime, "_host_cache") + reused = attn_utils_module.get_hisparse_host_pool(vllm_config, kv_cache_config) + assert reused is pool + assert reused.allocate(8) is backing + allocate_pinned.assert_called_once_with(8) + hisparse_runtime_module.drop_hisparse_host_pool(vllm_config) @pytest.mark.parametrize("failure_phase", ["allocation", "binding", "buffers"]) -def test_init_hisparse_rolls_back_shared_region(monkeypatch, failure_phase): - """A failure after mmap allocation must not leak the shared registration.""" +@pytest.mark.parametrize("shared", [False, True]) +def test_init_hisparse_rolls_back_host_pool(monkeypatch, failure_phase, shared): + """Failed initialization must release shared or private registrations.""" region = _FakeSharedHostRegion() + registered = torch.empty(4096, dtype=torch.int8) + cudart = MagicMock() + cudart.cudaHostUnregister.return_value.value = 0 + monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None) + monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart) vllm_config = SimpleNamespace( cache_config=SimpleNamespace( get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC ), scheduler_config=SimpleNamespace(max_num_seqs=1), + instance_id="rollback", ) + pools = [] def allocate(*args): - args[-1].shared_region = region + pools.append(args[-1]) + if shared: + args[-1].shared_region = region + else: + args[-1].registered = registered + args[-1].backing = registered[:8] if failure_phase == "allocation": raise RuntimeError("initialization failed") return {} @@ -472,7 +488,13 @@ def test_init_hisparse_rolls_back_shared_region(monkeypatch, failure_phase): {}, SimpleNamespace(), ) - assert region.cleanup_calls == 1 + assert region.cleanup_calls == int(shared) + if shared: + cudart.cudaHostUnregister.assert_not_called() + else: + cudart.cudaHostUnregister.assert_called_once_with(registered.data_ptr()) + assert attn_utils_module.get_hisparse_host_pool(vllm_config, None) is not pools[0] + hisparse_runtime_module.drop_hisparse_host_pool(vllm_config) def test_reshape_padded_kv_cache_strides_by_padded_page(): diff --git a/tests/v1/worker/test_cpu_worker_kv_budget.py b/tests/v1/worker/test_cpu_worker_kv_budget.py new file mode 100644 index 0000000000..e0493d7278 --- /dev/null +++ b/tests/v1/worker/test_cpu_worker_kv_budget.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The CPU worker's automatic KV cache sizing must report the right cause. + +A worker whose own footprint exceeds its reservation and a node that is short +of free memory need opposite remedies, so they get separate errors. +""" + +from types import SimpleNamespace + +import pytest + +from vllm.platforms import current_platform +from vllm.utils.cpu_resource_utils import MemoryNodeInfo + +# Importing the CPU worker patches torch and the GPU buffer utilities for the +# whole process (vllm/v1/worker/cpu/shm.py), so the import is deferred into +# the tests and they run only on the CPU platform. +pytestmark = pytest.mark.skipif( + not current_platform.is_cpu(), reason="CPU worker test; needs the CPU platform" +) + +GiB = 1 << 30 + + +def _worker(monkeypatch, *, reserved, rss, total=8 * GiB, available=4 * GiB): + import vllm.v1.worker.cpu_worker as cpu_worker_mod + from vllm.v1.worker.cpu_worker import CPUWorker + + worker = CPUWorker.__new__(CPUWorker) + worker.requested_cpu_memory = reserved + worker.cache_config = SimpleNamespace(kv_cache_memory_bytes=None) + monkeypatch.setattr(CPUWorker, "_should_warm_up_model", lambda self: False) + monkeypatch.setattr( + cpu_worker_mod, + "get_allowed_cpu_list", + lambda: [SimpleNamespace(numa_node=0)], + ) + monkeypatch.setattr( + cpu_worker_mod, + "get_memory_node_info", + lambda node: MemoryNodeInfo(total_memory=total, available_memory=available), + ) + process = SimpleNamespace(memory_info=lambda: SimpleNamespace(rss=rss)) + monkeypatch.setattr(cpu_worker_mod.psutil, "Process", lambda pid: process) + return worker + + +def test_footprint_above_reservation_asks_for_a_larger_reservation(monkeypatch): + # 0.76 GiB already in use against a 0.16 GiB reservation (2% of 8 GiB). + worker = _worker(monkeypatch, reserved=int(0.16 * GiB), rss=int(0.76 * GiB)) + with pytest.raises(ValueError) as exc: + worker.determine_available_memory() + msg = str(exc.value) + assert "increase `--gpu-memory-utilization` to more than 0.10" in msg + assert "--kv-cache-memory-bytes" in msg + # Neither the misleading remedy nor a negative KV size. + assert "other processes" not in msg + assert "-0." not in msg + + +def test_short_node_memory_still_asks_to_free_memory(monkeypatch): + worker = _worker(monkeypatch, reserved=6 * GiB, rss=1 * GiB, available=2 * GiB) + with pytest.raises(ValueError, match="Reduce CPU memory used by other processes"): + worker.determine_available_memory() + + +def test_budget_that_fits_sizes_the_kv_cache(monkeypatch): + worker = _worker(monkeypatch, reserved=3 * GiB, rss=1 * GiB) + assert worker.determine_available_memory() == 2 * GiB diff --git a/tests/v1/worker/test_gpu_draft_lm_head.py b/tests/v1/worker/test_gpu_draft_lm_head.py new file mode 100644 index 0000000000..75ceee15bb --- /dev/null +++ b/tests/v1/worker/test_gpu_draft_lm_head.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import gc +from types import SimpleNamespace + +import pytest +import torch + +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization.online.lm_head import ( + quantized_lm_head_copy, + refresh_quantized_lm_heads, +) +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from vllm.model_executor.model_loader.reload import ( + finalize_layerwise_reload, + initialize_layerwise_reload, + record_metadata_for_reloading, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.platforms import current_platform +from vllm.v1.worker.gpu_worker import Worker + +# Not a multiple of the vocab padding, so the head has padded rows. +VOCAB, HIDDEN = 4000, 512 +requires_cuda = pytest.mark.skipif( + not torch.cuda.is_available(), reason="Marlin needs CUDA" +) +# NVFP4 needs a CUDA W4A16 kernel; ROCm only runs the fp8 cases. +requires_nvfp4 = pytest.mark.skipif( + not current_platform.is_cuda(), reason="nvfp4 needs a CUDA W4A16 kernel" +) +NVFP4 = pytest.param("nvfp4", marks=requires_nvfp4) + + +def _target_head(monkeypatch, seed: int = 0) -> ParallelLMHead: + # Online methods read the model dtype from the current config. + model_config = SimpleNamespace(dtype=torch.bfloat16) + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.online.fp8.get_current_vllm_config", + lambda: SimpleNamespace(model_config=model_config), + ) + head = ParallelLMHead(VOCAB, HIDDEN, disable_tp=True) + gen = torch.Generator().manual_seed(seed) + head.weight.data.copy_(torch.randn(head.weight.shape, generator=gen) * 0.02) + return head.to("cuda", torch.bfloat16) + + +def _tensors(module: torch.nn.Module) -> dict[str, torch.Tensor]: + return dict(module.named_parameters()) | dict(module.named_buffers()) + + +@requires_cuda +@pytest.mark.parametrize( + ("quantization", "max_rel_err", "min_top1"), + [("fp8", 0.05, 0.9), pytest.param("nvfp4", 0.15, 0.75, marks=requires_nvfp4)], +) +def test_tracks_bf16_head_through_logits_processor( + default_vllm_config, dist_init, monkeypatch, quantization, max_rel_err, min_top1 +): + head = _target_head(monkeypatch) + target_weight = head.weight.detach().clone() + draft_head = quantized_lm_head_copy(head, quantization) + + assert isinstance(draft_head, ParallelLMHead) + # The source head keeps its bf16 weight, and no bf16 copy is kept around. + assert torch.equal(head.weight, target_weight) + assert draft_head.weight.dtype != torch.bfloat16 + + logits_processor = LogitsProcessor(VOCAB) + hidden = torch.randn(2048, HIDDEN, device="cuda", dtype=torch.bfloat16) + ref = logits_processor(head, hidden).float() + out = logits_processor(draft_head, hidden).float() + assert out.shape == ref.shape == (2048, VOCAB) + rel_err = ((out - ref).norm() / ref.norm()).item() + top1 = (out.argmax(-1) == ref.argmax(-1)).float().mean().item() + assert rel_err < max_rel_err + assert top1 > min_top1 + + +@requires_cuda +@pytest.mark.parametrize("quantization", ["fp8", NVFP4]) +def test_refresh_follows_target_update( + default_vllm_config, dist_init, monkeypatch, quantization +): + """After the target's weights change, the copy is re-derived in place: it + equals a fresh quantization of the new target, keeps its storage, and + frees the temporary copy without waiting for the cycle collector.""" + head = _target_head(monkeypatch) + drafter = torch.nn.Module() + drafter.lm_head = quantized_lm_head_copy(head, quantization) + ptrs = {k: t.data_ptr() for k, t in _tensors(drafter.lm_head).items()} + + new_weight = _target_head(monkeypatch, seed=1).weight.detach() + head.weight.data.copy_(new_weight) + gc.disable() + try: + allocated = torch.accelerator.memory_allocated() + refresh_quantized_lm_heads(drafter) + assert torch.accelerator.memory_allocated() == allocated + finally: + gc.enable() + + expected = _tensors(quantized_lm_head_copy(head, quantization)) + for name, tensor in _tensors(drafter.lm_head).items(): + assert tensor.data_ptr() == ptrs[name] + if name != "workspace": + assert torch.equal(tensor, expected[name]), name + assert torch.equal(head.weight, new_weight) + + +def _draft_update_worker(vllm_config, drafter: torch.nn.Module, **engine) -> Worker: + worker = object.__new__(Worker) + worker.vllm_config = vllm_config + worker.weight_transfer_engine = SimpleNamespace( + supports_draft_weight_update=True, + reset_weight_update_target=lambda: None, + **engine, + ) + worker.model_runner = SimpleNamespace(get_draft_model=lambda: drafter) + worker._weight_update_active = False + worker._set_draft_weight_update_target = lambda: None + return worker + + +@requires_cuda +def test_draft_weight_update_refreshes_copy( + default_vllm_config, dist_init, monkeypatch +): + """A draft weight update can write a target weight the drafter shares, such + as an embedding tied to the lm_head; the copy must follow it.""" + head = _target_head(monkeypatch) + drafter = torch.nn.Module() + drafter.lm_head = quantized_lm_head_copy(head, "fp8") + worker = _draft_update_worker( + default_vllm_config, + drafter, + start_weight_update=lambda: None, + finish_weight_update=lambda: None, + ) + + worker.start_draft_weight_update() + head.weight.data.copy_(_target_head(monkeypatch, seed=1).weight) + worker.finish_weight_update() + + expected = _tensors(quantized_lm_head_copy(head, "fp8")) + for name, tensor in _tensors(drafter.lm_head).items(): + assert torch.equal(tensor, expected[name]), name + + +@requires_cuda +@pytest.mark.parametrize("quantization", ["fp8", NVFP4]) +def test_draft_weight_update_loads_lm_head( + default_vllm_config, dist_init, monkeypatch, quantization +): + """A draft update carrying lm_head.weight loads the shared target head, as + without quantization, and the copy is re-derived from it in place.""" + head = _target_head(monkeypatch) + record_metadata_for_reloading(head) + drafter = torch.nn.Module() + draft_head = quantized_lm_head_copy(head, quantization) + drafter.lm_head = draft_head + ptrs = {k: t.data_ptr() for k, t in _tensors(draft_head).items()} + new_weight = torch.randn(VOCAB, HIDDEN, device="cuda", dtype=torch.bfloat16) + + def load(update_info): + param = drafter.get_parameter("lm_head.weight") + getattr(param, "weight_loader", default_weight_loader)(param, new_weight) + + worker = _draft_update_worker( + default_vllm_config, + drafter, + start_weight_update=lambda: initialize_layerwise_reload(drafter), + update_weights=load, + finish_weight_update=lambda: finalize_layerwise_reload(drafter, None), + ) + worker.start_draft_weight_update() + worker.update_weights({}) + worker.finish_weight_update() + + assert drafter.lm_head is draft_head + assert torch.equal(head.weight[:VOCAB], new_weight) + expected = _tensors(quantized_lm_head_copy(head, quantization)) + for name, tensor in _tensors(draft_head).items(): + assert tensor.data_ptr() == ptrs[name] + if name != "workspace": + assert torch.equal(tensor, expected[name]), name + hidden = torch.randn(64, HIDDEN, device="cuda", dtype=torch.bfloat16) + assert LogitsProcessor(VOCAB)(draft_head, hidden).isfinite().all() + + +@requires_cuda +@requires_nvfp4 +def test_layerwise_reload_of_drafter_keeps_copy( + default_vllm_config, dist_init, monkeypatch +): + """A drafter weight update reloads the drafter layerwise; the derived copy + must come back unchanged, not be re-quantized.""" + head = _target_head(monkeypatch) + drafter = torch.nn.Module() + record_metadata_for_reloading(drafter) + drafter.lm_head = quantized_lm_head_copy(head, "nvfp4") + before = {k: t.clone() for k, t in _tensors(drafter.lm_head).items()} + + def fail(*args, **kwargs): + raise AssertionError("the copy was re-quantized") + + monkeypatch.setattr( + drafter.lm_head.quant_method, "process_weights_after_loading", fail + ) + initialize_layerwise_reload(drafter) + finalize_layerwise_reload(drafter, None) + + after = _tensors(drafter.lm_head) + assert after.keys() == before.keys() + for name, tensor in before.items(): + assert torch.equal(after[name], tensor), name diff --git a/tests/v1/worker/test_gpu_ubatch_slicing.py b/tests/v1/worker/test_gpu_ubatch_slicing.py index 127d289afa..56a365218b 100644 --- a/tests/v1/worker/test_gpu_ubatch_slicing.py +++ b/tests/v1/worker/test_gpu_ubatch_slicing.py @@ -302,10 +302,8 @@ def _sync_dp( reduced[4] = int(allow_ubatching) reduced[5] = 8 # num_reqs - with ( - patch.object(dp_utils.dist, "all_reduce", lambda t, group: t.copy_(reduced)), - patch.object(dp_utils, "get_dp_group", lambda: SimpleNamespace(cpu_group=None)), - ): + dp_group = SimpleNamespace(cpu_all_reduce_=lambda t: t.copy_(reduced)) + with patch.object(dp_utils, "get_dp_group", lambda: dp_group): return dp_utils.sync_cudagraph_and_dp_padding( cudagraph_manager=cudagraph_manager, desired_batch_desc=BatchExecutionDescriptor( diff --git a/tests/v1/worker/test_gpu_worker_weight_transfer.py b/tests/v1/worker/test_gpu_worker_weight_transfer.py index ec63c08c70..61d6f38e08 100644 --- a/tests/v1/worker/test_gpu_worker_weight_transfer.py +++ b/tests/v1/worker/test_gpu_worker_weight_transfer.py @@ -65,6 +65,9 @@ class _RecordingModelRunner: def reset_lora_state(self) -> None: self.reset_lora_calls += 1 + def get_draft_model(self) -> nn.Module | None: + return None + def _make_worker(engine: _RecordingEngine | None) -> Worker: worker = object.__new__(Worker) diff --git a/tests/v1/worker/test_kv_cache_allocation_scope.py b/tests/v1/worker/test_kv_cache_allocation_scope.py index 52cf2a3741..b5b288d9b5 100644 --- a/tests/v1/worker/test_kv_cache_allocation_scope.py +++ b/tests/v1/worker/test_kv_cache_allocation_scope.py @@ -62,7 +62,7 @@ def test_mrv2_kv_pool_only_wraps_backing_allocation(monkeypatch, hisparse) -> No monkeypatch.setattr( hisparse_binding, - "HiSparseHostPool", + "get_hisparse_host_pool", lambda vllm_config, kv_cache_config: host_pool, ) monkeypatch.setattr(hisparse_binding, "allocate_hisparse_kv_caches", allocate) diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index a8f49c1b88..b2e9e9054f 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -1257,6 +1257,8 @@ def test_hisparse_worker_shutdown_releases_pinned_state(monkeypatch): worker.cache_handles = [] worker.pinned_host_pools = [] worker.shared_host_region = object() + worker.vllm_config = SimpleNamespace(instance_id="shutdown") + pool = hisparse_runtime_module.get_hisparse_host_pool(worker.vllm_config, None) released = False def release_pinned_state(runtimes, pinned_host_pools, shared_host_region): @@ -1273,6 +1275,9 @@ def test_hisparse_worker_shutdown_releases_pinned_state(monkeypatch): worker.shutdown() assert released + new_pool = hisparse_runtime_module.get_hisparse_host_pool(worker.vllm_config, None) + assert new_pool is not pool + hisparse_runtime_module.drop_hisparse_host_pool(worker.vllm_config) class _TestReplaySSMMixer(MambaMixer2): diff --git a/tests/watermarking/test_gumbel.py b/tests/watermarking/test_gumbel.py index aa0ffa2799..b2d7aa006e 100644 --- a/tests/watermarking/test_gumbel.py +++ b/tests/watermarking/test_gumbel.py @@ -253,6 +253,8 @@ def test_philox_mixed_draft_sampling_caches_unprocessed_logits(): positions, apply_temperature=True, is_drafting=True, + logits_cache=cache, + logits_cache_col=cols, ) watermarked = philox_gumbel_sample(processed_logits, contexts, 42) expected = torch.where(skip_mask, ordinary, watermarked) diff --git a/tests/watermarking/test_watermarking.py b/tests/watermarking/test_watermarking.py index 384cc30d57..eec43845db 100644 --- a/tests/watermarking/test_watermarking.py +++ b/tests/watermarking/test_watermarking.py @@ -36,6 +36,7 @@ from vllm.v1.worker.gpu.sample.watermark import ( philox_gumbel_sample, repeated_context_mask, ) +from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator @@ -254,8 +255,8 @@ def test_speculative_decoding_uses_fixed_dual_key_roles(): assert draft is not None assert draft.watermarker.prf.key == derive_watermark_key(42, b"key_a") assert draft.num_speculative_steps == 3 - assert draft.deduplicate_contexts == "all" - assert draft.deduplicate_contexts_max_history is None + assert draft.state.deduplicate_contexts == "all" + assert draft.state.deduplicate_contexts_max_history is None def test_recovery_key_rejects_the_unsplit_dual_key_watermarker(): @@ -1324,8 +1325,8 @@ def test_draft_sampler_uses_draft_key_and_advances_context(monkeypatch): deduplicate_contexts="none", deduplicate_contexts_max_history=None, ) - draft_watermarker.contexts.copy_(torch.tensor([[1, 2], [3, 4]])) - draft_watermarker.enabled = torch.tensor([True, False]) + draft_watermarker.state.contexts.copy_(torch.tensor([[1, 2], [3, 4]])) + draft_watermarker.state.enabled = torch.tensor([True, False]) speculator.draft_watermarker = draft_watermarker monkeypatch.setattr( "vllm.v1.watermarking.watermarker.gumbel_sample", @@ -1343,7 +1344,7 @@ def test_draft_sampler_uses_draft_key_and_advances_context(monkeypatch): ) assert torch.equal(sampled, torch.tensor([7, 4])) - assert torch.equal(draft_watermarker.contexts, torch.tensor([[2, 7], [4, 4]])) + assert torch.equal(draft_watermarker.state.contexts, torch.tensor([[2, 7], [4, 4]])) def test_draft_sampler_uses_ordinary_samples_for_repeated_contexts(monkeypatch): @@ -1602,11 +1603,12 @@ def test_draft_watermarker_deduplicates_across_a_block_on_cuda( [True, True], ] assert torch.equal( - draft_watermarker.prior_contexts[:num_reqs].cpu(), + draft_watermarker.state.prior_contexts[:num_reqs].cpu(), torch.tensor([[[5, 6], [6, 6], [6, 6]], [[6, 3], [3, 6], [6, 6]]]), ) assert torch.equal( - draft_watermarker.contexts[:num_reqs].cpu(), torch.tensor([[6, 6], [6, 6]]) + draft_watermarker.state.contexts[:num_reqs].cpu(), + torch.tensor([[6, 6], [6, 6]]), ) for step, logits in enumerate(step_logits): # The 0-d draft step selects the logits cache column. @@ -1661,6 +1663,154 @@ def test_draft_watermarker_cuda_graph_reads_current_prompt_lens(): assert mask.item() +def test_dflash_draft_sampler_watermarks_drafts_in_step_order(monkeypatch): + class NextTokenWatermarker: + context_width = 2 + + @staticmethod + def _try_sample_block(*args, **kwargs): + return None + + @staticmethod + def sample(logits, contexts, random_sampler=None, skip_mask=None): + watermarked = contexts[:, -1] + 1 + return WatermarkSample( + torch.where(skip_mask, random_sampler(logits), watermarked), logits + ) + + def ordinary_sample(logits, idx_mapping, temperature, seed, pos, **kwargs): + assert kwargs["use_fp64"] + cache, col = kwargs["logits_cache"], kwargs["logits_cache_col"] + cache[idx_mapping, col] = kwargs["logits_cache_source"] + return seed[idx_mapping] + pos + + monkeypatch.setattr( + "vllm.v1.watermarking.watermarker.gumbel_sample", ordinary_sample + ) + num_reqs, num_steps, vocab_size = 3, 2, 4 + draft_watermarker = DraftWatermarker( + NextTokenWatermarker(), + max_num_reqs=8, + device=torch.device("cpu"), + num_speculative_steps=num_steps, + deduplicate_contexts="none", + deduplicate_contexts_max_history=None, + ) + # Batch rows hold request slots 2, 0, 1. Slot 0 is greedy and row 2 opted out. + draft_watermarker.state.contexts[:num_reqs] = torch.tensor([[1, 2], [3, 4], [5, 6]]) + draft_watermarker.state.enabled[:num_reqs] = torch.tensor([True, True, False]) + speculator = object.__new__(DFlashSpeculator) + speculator.model = SimpleNamespace( + compute_logits=lambda hidden_states: hidden_states + ) + speculator.draft_watermarker = draft_watermarker + speculator.use_fp64_gumbel = True + speculator.acceptance_estimator = None + batch_slots = torch.tensor([2, 0, 1]) + logits = torch.arange(num_reqs * num_steps * vocab_size, dtype=torch.float32) + draft_logits = torch.zeros(num_reqs, num_steps, vocab_size) + + sampled = speculator.sample_draft( + hidden_states=logits.view(-1, vocab_size), + sample_src_positions=torch.arange(num_reqs * num_steps), + idx_mapping=batch_slots.repeat_interleave(num_steps), + temperature=torch.tensor([0.0, 1.0, 0.5]), + seeds=torch.tensor([100, 200, 300]), + draft_step=torch.arange(num_steps).repeat(num_reqs), + draft_logits=draft_logits, + num_steps=num_steps, + ) + + assert torch.equal(sampled, torch.tensor([3, 4, 102, 103, 204, 205])) + assert torch.equal(draft_logits[batch_slots], logits.view(num_reqs, num_steps, -1)) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" +) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("deduplicate_contexts", ["none", "single_turn", "all"]) +def test_dflash_draft_sampler_matches_step_loop_on_cuda(deduplicate_contexts, dtype): + """The fused CUDA `sample_block` is bit-exact with sampling step by step. + + Covers repeated contexts, a NaN logit, and greedy, opted-out and padding + rows; padding rows are discarded, so only valid rows are compared. + """ + num_reqs, num_steps, vocab_size, max_num_reqs = 4, 5, 1500, 6 + generator = torch.Generator().manual_seed(0) + logits = torch.randn(num_reqs, num_steps, vocab_size, generator=generator) * 3 + logits[..., :3] += 6 + logits[0, 1, 5] = float("nan") + logits = logits.to(dtype).cuda() + idx_mapping = torch.tensor([4, 2, 0, -1], dtype=torch.int32) + valid = (idx_mapping >= 0).cuda() + idx_mapping = idx_mapping.repeat_interleave(num_steps).cuda() + temperature = torch.tensor([0.7, 1.0, 0.0, 1.3, 0.6, 0.9]).cuda() + seeds = torch.arange(max_num_reqs).cuda() * 1000 + pos = (torch.arange(num_steps).repeat(num_reqs) + 20).cuda() + draft_step = torch.arange(num_steps).repeat(num_reqs).cuda() + contexts = torch.randint(0, 3, (num_reqs, 2), generator=generator) + contexts[0, 0] = -1 + all_token_ids = torch.randint(0, 3, (max_num_reqs, 32), generator=generator) + + def run(fused: bool): + draft_watermarker = DraftWatermarker( + GumbelWatermarker(key=42, context_width=2), + max_num_reqs=max_num_reqs, + device=torch.device("cuda"), + num_speculative_steps=num_steps, + deduplicate_contexts=deduplicate_contexts, + deduplicate_contexts_max_history=None, + ) + draft_watermarker.prepare( + contexts=contexts.cuda(), + enabled=torch.tensor([True, False, True, True]).cuda(), + all_token_ids=all_token_ids.cuda(), + prompt_lens=torch.full((max_num_reqs,), 4).cuda(), + total_lens=torch.full((max_num_reqs,), 24).cuda(), + ) + draft_logits = torch.zeros(max_num_reqs, num_steps, vocab_size, dtype=dtype) + draft_logits = draft_logits.cuda() + args = (temperature, seeds) + if fused: + sampled = draft_watermarker.sample_block( + logits.view(-1, vocab_size), + idx_mapping, + *args, + pos, + logits_cache=draft_logits, + logits_cache_col=draft_step, + ) + else: + by_step = [t.view(num_reqs, num_steps) for t in (idx_mapping, pos)] + steps = draft_step.view(num_reqs, num_steps) + sampled = torch.stack( + [ + draft_watermarker.sample( + logits[:, step], + by_step[0][:, step], + *args, + by_step[1][:, step], + logits_cache=draft_logits, + logits_cache_col=steps[:, step], + apply_temperature=True, + is_drafting=True, + ) + for step in range(num_steps) + ], + dim=1, + ).view(-1) + return ( + sampled.view(num_reqs, num_steps)[valid], + draft_watermarker.state.contexts[:num_reqs][valid], + draft_watermarker.state.prior_contexts, + draft_logits, + ) + + for fused, reference in zip(run(fused=True), run(fused=False)): + torch.testing.assert_close(fused, reference, rtol=0, atol=0, equal_nan=True) + + def test_dspark_reduced_vocab_draft_sampler_applies_watermarking(monkeypatch): speculator = object.__new__(DSparkSpeculator) speculator.draft_logits = torch.zeros(2, 1, 8) diff --git a/tools/install_deepgemm.sh b/tools/install_deepgemm.sh index b617e6d567..b3f7ceb43f 100755 --- a/tools/install_deepgemm.sh +++ b/tools/install_deepgemm.sh @@ -9,8 +9,9 @@ set -e DEEPGEMM_GIT_REPO="https://github.com/vllm-project/DeepGEMM.git" # NOTE: This targets the vLLM fork's dev branch: upstream's 26/09/30 release # plus the sm120 and sm90 paged-MQA ports, the SwiGLU alpha/beta and SiTU -# Mega MoE activations, and NVFP4 Mega MoE. -DEEPGEMM_GIT_REF="6ce1c5ec6627a6bf11dbcc698b0b1e19c9248432" +# Mega MoE activations, NVFP4 Mega MoE, and BF16/FP8/NVFP4 swiglustep. +# Uses the PyTorch 2.11 stable ABI and CPython 3.10 limited API. +DEEPGEMM_GIT_REF="a7a927731b109fdfec770df7c2a5be5d7e2216b5" WHEEL_DIR="" # Parse command line arguments diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 9e01735441..0f57c9ecb4 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -11,6 +11,7 @@ import torch.distributed as dist from torch._ops import OpOverload import vllm.envs as envs +from vllm import _custom_ops as ops from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.utils.torch_utils import direct_register_custom_op @@ -526,6 +527,46 @@ def _rocm_aiter_topk_softmax_impl( ) +def _rocm_aiter_topk_gating_impl( + topk_weights: torch.Tensor, + topk_indices: torch.Tensor, + token_expert_indices: torch.Tensor, + gating_output: torch.Tensor, + renormalize: bool, + num_shared_experts: int = 0, + shared_expert_scoring_func: str = "", +) -> None: + # Callers pick this op only for launches AITER's topk_gating supports (see + # dispatch_topk_softmax_func); everything else goes to topk_softmax. + from aiter.ops.topk import topk_gating + + # AITER renormalizes the selected softmax mass inside this launch, + # preserving vLLM's `renormalize=True` semantics without a second + # pointwise kernel, and sigmoids the trailing shared-expert columns. + topk_gating( + topk_weights, + topk_indices, + gating_output, + correction_bias=None, + need_renorm=renormalize, + routed_scaling_factor=1.0, + score_func="softmax", + num_shared_experts=num_shared_experts, + ) + + +def _rocm_aiter_topk_gating_fake( + topk_weights: torch.Tensor, + topk_indices: torch.Tensor, + token_expert_indices: torch.Tensor, + gating_output: torch.Tensor, + renormalize: bool, + num_shared_experts: int = 0, + shared_expert_scoring_func: str = "", +) -> None: + pass + + def _rocm_aiter_topk_sigmoid_impl( topk_weights: torch.Tensor, topk_indices: torch.Tensor, @@ -970,6 +1011,39 @@ def _rocm_aiter_triton_gemm_a8w8_blockscale_fake( return Y +def _rocm_aiter_triton_gemm_a8w8_blockscale_preshuffle_impl( + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + output_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + from aiter.ops.triton.gemm_a8w8_blockscale import gemm_a8w8_blockscale_preshuffle + + n, k = B.shape + return gemm_a8w8_blockscale_preshuffle( + A, + B.reshape(n // 16, k * 16), + As, + Bs, + output_dtype, + None, # y + None, # config + False, # skip_reduce + False, # is_x_scale_transposed (misspelled in aiter) + ) + + +def _rocm_aiter_triton_gemm_a8w8_blockscale_preshuffle_fake( + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + output_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + return torch.empty(A.shape[0], B.shape[0], dtype=output_dtype, device=A.device) + + def _rocm_aiter_gemm_a8w8_blockscale_impl( A: torch.Tensor, B: torch.Tensor, @@ -2020,7 +2094,7 @@ class rocm_aiter_ops: Operations: - GEMM operations: gemm_a8w8, gemm_a8w8_blockscale - Fused MoE: fused_moe, asm_moe_tkw1 - - Routing: topk_softmax, biased_grouped_topk, grouped_topk + - Routing: topk_softmax, topk_gating, biased_grouped_topk, grouped_topk - MLA decode: mla_decode_fwd - Quantization: per_tensor_quant, per_token_quant, group_fp8_quant - Triton ops: triton_rotary_embed, triton_fp8_bmm, triton_gemm_a8w8_blockscale @@ -2198,6 +2272,18 @@ class rocm_aiter_ops: def is_fused_moe_enabled(cls) -> bool: return cls._AITER_ENABLED and cls._FMOE_ENABLED + @classmethod + @if_aiter_supported + def is_topk_gating_enabled(cls) -> bool: + """Use AITER ``topk_gating`` for softmax routing. + + CDNA3 and CDNA4 only (gfx942, gfx950): those are the targets it was + measured on. Other AITER targets keep the legacy ``topk_softmax`` launcher. + """ + from vllm.platforms.rocm import get_cdna_version + + return cls.is_fused_moe_enabled() and get_cdna_version() in (3, 4) + @classmethod @if_aiter_supported def is_fusion_moe_shared_experts_enabled(cls) -> bool: @@ -2541,6 +2627,32 @@ class rocm_aiter_ops: """ return cls._AITER_ENABLED and cls._gdn_triton_kernels_importable() + @staticmethod + def _kda_gluon_decode_importable() -> bool: + try: + from aiter.ops.triton.attention import kda # noqa: F401 + + return True + except (ImportError, ModuleNotFoundError): + return False + + @classmethod + @if_aiter_supported + def is_kda_gluon_decode_available(cls) -> bool: + """Check if AITER's Gluon KDA decode kernels can run here. + + ``aiter.ops.triton.attention.kda`` (conv + gated delta rule + gated + RMSNorm in one launch, plain and speculative decode) ships CDNA4+ + (gfx950, gfx1250) builds only, and is absent in older aiter releases. + """ + from vllm.platforms.rocm import get_cdna_version + + return ( + cls._AITER_ENABLED + and get_cdna_version() >= 4 + and cls._kda_gluon_decode_importable() + ) + @classmethod @functools.cache def is_gdn_flydsl_prefill_available(cls) -> bool: @@ -2851,6 +2963,14 @@ class rocm_aiter_ops: dispatch_key=current_platform.dispatch_key, ) + direct_register_custom_op( + op_name="rocm_aiter_topk_gating", + op_func=_rocm_aiter_topk_gating_impl, + mutates_args=["topk_weights", "topk_indices"], + fake_impl=_rocm_aiter_topk_gating_fake, + dispatch_key=current_platform.dispatch_key, + ) + direct_register_custom_op( op_name="rocm_aiter_topk_sigmoid", op_func=_rocm_aiter_topk_sigmoid_impl, @@ -2924,6 +3044,12 @@ class rocm_aiter_ops: fake_impl=_rocm_aiter_triton_gemm_a8w8_blockscale_fake, ) + direct_register_custom_op( + op_name="rocm_aiter_triton_gemm_a8w8_blockscale_preshuffle", + op_func=_rocm_aiter_triton_gemm_a8w8_blockscale_preshuffle_impl, + fake_impl=_rocm_aiter_triton_gemm_a8w8_blockscale_preshuffle_fake, + ) + direct_register_custom_op( op_name="rocm_aiter_gemm_a8w8_blockscale", op_func=_rocm_aiter_gemm_a8w8_blockscale_impl, @@ -3238,6 +3364,19 @@ class rocm_aiter_ops: A, B, As, Bs, output_dtype ) + @staticmethod + def triton_gemm_a8w8_blockscale_preshuffle( + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + block_size: list[int], + output_dtype: torch.dtype = torch.bfloat16, + ) -> torch.Tensor: + return torch.ops.vllm.rocm_aiter_triton_gemm_a8w8_blockscale_preshuffle( + A, B, As, Bs, output_dtype + ) + @staticmethod def gemm_a8w8_blockscale( A: torch.Tensor, @@ -3382,6 +3521,27 @@ class rocm_aiter_ops: ) return topk_weights, topk_indices + @staticmethod + def topk_gating( + topk_weights: torch.Tensor, + topk_indices: torch.Tensor, + token_expert_indices: torch.Tensor, + gating_output: torch.Tensor, + renormalize: bool, + num_shared_experts: int = 0, + shared_expert_scoring_func: str = "", + ) -> tuple[torch.Tensor, ...]: + torch.ops.vllm.rocm_aiter_topk_gating( + topk_weights, + topk_indices, + token_expert_indices, + gating_output, + renormalize, + num_shared_experts, + shared_expert_scoring_func, + ) + return topk_weights, topk_indices + @staticmethod def topk_sigmoid( topk_weights: torch.Tensor, @@ -3728,17 +3888,29 @@ class rocm_aiter_ops: v_scale: torch.Tensor, kv_cache_dtype: str, use_shuffle_layout: bool, + return_kv: bool, ) -> None: """Run the fused QK-norm+RoPE+KV-cache op on already-split k/v caches. - Shared by the AITER FA and unified-attention impls. The caller splits + Shared by the AITER FA, ROCM ATTN, and unified-attention impls. The caller splits kv_cache, since the unbind dim depends on the layout (e.g. the unified encoder-decoder path is K/V-first), and passes use_shuffle_layout - (unified reads NHD and must pass False). + (unified reads NHD and must pass False). k_scale/v_scale are the + per-tensor CPU scalars. return_kv says whether the impl consumes + k_out (FA and ROCM_ATTN prefill do; unified reads K from the cache and + passes False, so the kernel skips that store). For fp8 caches the + kernel emits k_out in the cache dtype, so a returned k_out is + converted to the activation dtype. """ - if kv_cache_dtype.startswith("fp8"): - key_cache = key_cache.view(current_platform.fp8_dtype()) - value_cache = value_cache.view(current_platform.fp8_dtype()) + is_fp8_cache = kv_cache_dtype.startswith("fp8") + dequant_k_out = return_kv and is_fp8_cache + kernel_k_out = k_out + if is_fp8_cache: + fp8_dtype = current_platform.fp8_dtype() + key_cache = key_cache.view(fp8_dtype) + value_cache = value_cache.view(fp8_dtype) + if dequant_k_out: + kernel_k_out = torch.empty_like(k_out, dtype=fp8_dtype) # Partial-rotary support (e.g. GLM-4.7 applies rotary to only a prefix # of each head's channel dim). rotary_dim = cos_sin_cache.shape[-1] @@ -3761,14 +3933,16 @@ class rocm_aiter_ops: slot_mapping=slot_mapping, k_scale=k_scale, v_scale=v_scale, - k_out=k_out, + k_out=kernel_k_out, v_out=None, - return_kv=True, + return_kv=return_kv, use_shuffle_layout=use_shuffle_layout, block_size=key_cache.shape[1], x=16 // key_cache.element_size(), rotary_dim=kernel_rotary_dim, ) + if dequant_k_out: + ops.convert_fp8(k_out, kernel_k_out, float(k_scale), kv_cache_dtype) @staticmethod def fused_qk_norm_mrope_and_cache( @@ -3918,6 +4092,90 @@ class rocm_aiter_ops: output_zeros=False, ) + @staticmethod + def kda_fused_recurrent_packed_decode( + mixed_qkv: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + lower_bound: float | None, + initial_state: torch.Tensor, + ssm_state_indices: torch.Tensor, + conv_state: torch.Tensor, + conv_weight: torch.Tensor, + out_gate: torch.Tensor, + norm_weight: torch.Tensor, + norm_eps: float, + out: torch.Tensor, + ) -> torch.Tensor: + from aiter.ops.triton.attention.kda import fused_recurrent_kda_packed_decode + + return fused_recurrent_kda_packed_decode( + mixed_qkv=mixed_qkv, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + initial_state=initial_state, + ssm_state_indices=ssm_state_indices, + conv_state=conv_state, + conv_weight=conv_weight, + out_gate=out_gate, + norm_weight=norm_weight, + norm_eps=norm_eps, + out=out, + ) + + @staticmethod + def kda_fused_recurrent_spec_decode( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + lower_bound: float | None, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, + ssm_state_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor, + conv_state: torch.Tensor, + conv_weight: torch.Tensor, + out_gate: torch.Tensor, + norm_weight: torch.Tensor, + norm_eps: float, + out: torch.Tensor, + ) -> None: + from aiter.ops.triton.attention.kda import fused_recurrent_kda + + fused_recurrent_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + initial_state=initial_state, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + num_accepted_tokens=num_accepted_tokens, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + pad_slot_guard=True, + conv_state=conv_state, + conv_weight=conv_weight, + out_gate=out_gate, + norm_weight=norm_weight, + norm_eps=norm_eps, + out=out, + ) + @staticmethod def batched_gemm_a16wfp4( X: torch.Tensor, @@ -3993,15 +4251,28 @@ class rocm_aiter_ops: except (AssertionError, ImportError): return False + @staticmethod + @if_aiter_supported + def is_triton_blockscale_bpreshuffle_gemm_tuned(n: int, k: int) -> bool: + from aiter.ops.triton.utils.gemm_config_utils import ( + get_gemm_config, + ) + + # Fake M value for now + _, is_tuned = get_gemm_config("GEMM-A8W8_BLOCKSCALE_PRESHUFFLED", 1, n, k) + return is_tuned + @staticmethod @functools.cache def is_triton_gemm_afp4wfp4_presh_ws_tuned(n: int, k_bytes: int) -> bool: if not current_platform.is_rocm(): return False - # Input = weight.shape[1] (bytes); *4 for fp4 and aiter config 2*k + # k_bytes = weight.shape[1] (two fp4 per byte). AITER names these + # config files by the logical K, which is what gemm_afp4wfp4_preshuffle + # itself looks up, so probe the same file. try: return _triton_gemm_config_is_tuned( - "GEMM-AFP4WFP4_PRESHUFFLED", n, 4 * k_bytes + "GEMM-AFP4WFP4_PRESHUFFLED", n, 2 * k_bytes ) except (AssertionError, ImportError): return False @@ -4250,8 +4521,8 @@ class rocm_aiter_ops: scale: float, K_QScale_hip: torch.Tensor, V_QScale_hip: torch.Tensor, - K_QScale_asm: torch.Tensor, - V_QScale_asm: torch.Tensor, + K_QScale_asm: torch.Tensor | None, + V_QScale_asm: torch.Tensor | None, out_: torch.Tensor, kv_cache_dtype: str, ): diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 1af864bc74..70fed78573 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -122,6 +122,7 @@ def paged_attention_rocm( v_scale: torch.Tensor, fp8_out_scale: torch.Tensor | None = None, mfma_type: str = "fp8" if envs.VLLM_ROCM_FP8_MFMA_PAGE_ATTN else "f16", + use_interleaved_v_cache: bool = False, ) -> None: torch.ops._rocm_C.paged_attention( out, @@ -144,6 +145,7 @@ def paged_attention_rocm( v_scale, fp8_out_scale, mfma_type, + use_interleaved_v_cache, ) @@ -2554,6 +2556,7 @@ def reshape_and_cache( kv_cache_dtype: str, k_scale: torch.Tensor, v_scale: torch.Tensor, + interleaved_v: bool = False, ) -> None: torch.ops._C_cache_ops.reshape_and_cache( key, @@ -2564,6 +2567,7 @@ def reshape_and_cache( kv_cache_dtype, k_scale, v_scale, + interleaved_v, ) diff --git a/vllm/compilation/passes/fusion/qk_norm_rope_kvcache_fusion.py b/vllm/compilation/passes/fusion/qk_norm_rope_kvcache_fusion.py index f418e8a7ef..35803e1347 100644 --- a/vllm/compilation/passes/fusion/qk_norm_rope_kvcache_fusion.py +++ b/vllm/compilation/passes/fusion/qk_norm_rope_kvcache_fusion.py @@ -893,6 +893,7 @@ class QkNormRopeKvCacheFusionPass(VllmPatternMatcherPass): layer.head_size, ) continue + layer.impl.set_interleaved_v_cache() if supports_rope: for epsilon in [1e-5, 1e-6]: for neox in [True, False]: @@ -903,7 +904,6 @@ class QkNormRopeKvCacheFusionPass(VllmPatternMatcherPass): is_neox=neox, quant_query=quant_q, ).register(self.patterns) - if supports_mrope: for section, is_interleaved in self.mrope_configs: rotary_dim = 2 * sum(section) diff --git a/vllm/config/attention.py b/vllm/config/attention.py index 67605f3c7e..fc3ecb0dfc 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -103,6 +103,10 @@ class AttentionConfig: candidate pool instead, with fp32 logits, and keeps the masked dense walk for steps whose contexts are too short for the pool to pay.""" + indexer_tp_shard_min_rows: int = Field(default=65536, ge=1) + """Minimum total prefill rows for the sparse-attention indexer to split + them across TP ranks instead of scoring every row on each rank.""" + hisparse_config: HiSparseConfig | None = None """HiSparse host-resident KV configuration. Setting this enables experimental Model Runner V2-only HiSparse sparse-MLA decode hot-buffering. It is inferred diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 5c4b877fd7..a1b141237f 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -139,6 +139,7 @@ MoEBackend = Literal[ "aiter", "aiter_mega_moe", "aiter_triton_mxfp4_bf16", + "aiter_triton_mxfp4_mxfp4", "flydsl", "hpc", "emulation", @@ -329,6 +330,8 @@ class KernelConfig: DeepSeek-V4 only) - "aiter_triton_mxfp4_bf16": Use the AITER Triton MXFP4 W4A16 (moe_gemm_a16w4) MoE kernel (ROCm gfx942/gfx950/gfx1250) + - "aiter_triton_mxfp4_mxfp4": Use the AITER Triton/Gluon MXFP4 W4A4 + (moe_gemm_a4w4) MoE kernel (ROCm gfx1250) - "flydsl": Use AMD FlyDSL kernels (ROCm only) - "rdna3": Use the fused RDNA3 W4A16 HIP kernel (ROCm gfx1100 only) - "hpc": Use HPC kernels (FP8 and Hopper only) diff --git a/vllm/config/model.py b/vllm/config/model.py index c4016e82bc..7a883eac54 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1534,8 +1534,8 @@ class ModelConfig: parallel_config: ParallelConfig, ) -> None: if parallel_config.enable_hc_sp and parallel_config.tensor_parallel_size > 1: - if not current_platform.is_cuda(): - raise ValueError("--enable-hc-sp currently requires CUDA") + if not current_platform.is_cuda_alike(): + raise ValueError("--enable-hc-sp currently requires CUDA or ROCm") if not self.registry.is_hc_sp_supported_model(self.architectures, self): raise ValueError("--enable-hc-sp is not supported for this model") if parallel_config.pipeline_parallel_size != 1: @@ -1622,8 +1622,8 @@ class ModelConfig: def get_inputs_embeds_size(self) -> int: # The size of inputs_embeds is usually identical to the size # of the hidden states, however there are exceptions, such as - # embedding models like CLIP and SigLIP - names = ("projection_dim", "projection_size") + # OPT's projected token embeddings and embedding models like CLIP/SigLIP. + names = ("word_embed_proj_dim", "projection_dim", "projection_size") return getattr_iter( self.hf_text_config, names, default_factory=self.get_hidden_size ) diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 6716d5a21b..69677ee9d6 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -15,6 +15,7 @@ from torch.distributed import ProcessGroup, ReduceOp, Store import vllm.envs as envs from vllm.config.fault_tolerance import FaultToleranceConfig from vllm.config.utils import config +from vllm.distributed.utils import cpu_all_reduce_ from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.utils.network_utils import get_open_ports_list @@ -183,7 +184,7 @@ class ParallelConfig: """Whether the deployed model is MoE (if known).""" enable_hc_sp: bool = False """Enable sequence parallelism for HyperConnection (HC) computation. - Currently supported by Qwen4Exp on CUDA with TP>1 and PP=1 (ignored when TP=1). + Currently supported by Qwen4Exp on CUDA/ROCm with TP>1 and PP=1 (ignored when TP=1). Does not change MoE parallelism. MoE SP may enable HC SP independently of this flag. """ enable_expert_parallel: bool = False @@ -887,7 +888,7 @@ class ParallelConfig: tensor = torch.tensor( [int(has_unfinished), int(pending_pause)], dtype=torch.int32, device="cpu" ) - torch.distributed.all_reduce(tensor, op=ReduceOp.SUM, group=dp_group) + cpu_all_reduce_(tensor, dp_group) dp_size = dp_group.size() pause_count = tensor[1].item() has_unfinished_global = tensor[0].item() > 0 or pause_count % dp_size != 0 @@ -979,11 +980,11 @@ class ParallelConfig: "Elastic EP is not supported with pipeline parallelism " f"(pipeline_parallel_size={self.pipeline_parallel_size})." ) - if self.data_parallel_external_lb or self.data_parallel_hybrid_lb: + if self.data_parallel_hybrid_lb: raise NotImplementedError( - "Elastic EP is not compatible with data_parallel_external_lb " - "or data_parallel_hybrid_lb. Elastic EP relies on a single API " - "server and core client to coordinate scale up/down." + "Elastic EP is not compatible with data_parallel_hybrid_lb. " + "Elastic EP supports a single API server/core client path " + "or data_parallel_external_lb, but not hybrid load balancing." ) if self.eplb_config.use_async: from vllm.distributed.nixl_utils import is_nixl_available diff --git a/vllm/config/scheduler.py b/vllm/config/scheduler.py index 1dbe4b72eb..5ad38c0a46 100644 --- a/vllm/config/scheduler.py +++ b/vllm/config/scheduler.py @@ -201,9 +201,12 @@ class SchedulerConfig: (the default) disables the watermark.""" prefill_schedule_interval: int = Field(default=1, ge=1) - """For data-parallel deployments, only admit new prefill requests - once every N engine steps, aligned across DP ranks, to better balance - per-step forward-pass times.""" + """Only admit prefill once every N engine steps, leaving the steps in + between to decode alone. 1 (the default) admits prefill on every step. + + A deferred step carries no prefill at all, including chunks of an + already-admitted prompt, so a long prompt can take up to N times as many + steps to finish prefilling.""" async_scheduling: bool | None = None """If set to False, disable async scheduling. Async scheduling helps to diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index c80016fd36..0205393f69 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -452,6 +452,12 @@ class SpeculativeConfig: speculative methods. This is an experimental option for measuring the acceptance-rate impact of reusing that block. It does not disable the speculative drafter itself.""" + draft_lm_head_quantization: Literal["fp8", "nvfp4"] | None = None + """Quantize the drafter's copy of a target lm_head it shares (EAGLE and + MTP). "fp8" uses per-row weight and per-token activation scales on the + platform's FP8 linear kernel; "nvfp4" is weight-only on the platform's + W4A16 kernel. The target verifies with its own head, so only the + acceptance length can change. Model Runner V2 only.""" use_local_argmax_reduction: bool = False """Use vocab-parallel local argmax instead of all-gathering full logits for draft token generation. Reduces communication from O(vocab_size) to @@ -1780,6 +1786,20 @@ class SpeculativeConfig: "n_predict parameter." ) + if self.draft_lm_head_quantization is not None: + # Gemma-4 MTP scores tied embedding rows directly; no lm_head to share. + if self.method not in ("eagle", "eagle3", "mtp") or self.use_gemma4_mtp(): + raise ValueError( + "draft_lm_head_quantization is not supported for method " + f"{self.method!r}." + ) + draft = self.draft_model_config + if draft is not None and draft.head_dtype != draft.dtype: + raise ValueError( + "draft_lm_head_quantization needs the lm_head to run in the " + f"model dtype, got head_dtype={draft.head_dtype}." + ) + if self.num_speculative_tokens <= 0: raise ValueError( "Expected num_speculative_tokens to be greater " @@ -1853,19 +1873,21 @@ class SpeculativeConfig: ==================== ============= ======== ================ EAGLE3 eagle3 No 0 P-EAGLE eagle3 Yes K - 1 - DFlash dflash Yes K - DSpark dspark Yes K - 1 + DFlash dflash Yes 0 + DSpark dspark Yes 0 MTP mtp No 0 N-gram ngram No 0 Draft model draft_model No 1 PARD draft_model Yes K ==================== ============= ======== ================ + + DFlash and DSpark draft in a separate query pass rather than adding + rows to the target batch. """ num_draft_tokens = self.num_speculative_tokens - if self.use_dflash(): - # DFlash uses one bonus query followed by K mask queries. - return num_draft_tokens + if self.use_dflash() or self.use_dspark(): + return 0 if self.parallel_drafting: if self.uses_draft_model(): diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 7c1710f89f..85c8e772eb 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1325,15 +1325,26 @@ class VllmConfig: ) if ( speculative_config.parallel_drafting - and speculative_config.method != "dspark" + and speculative_config.method not in ("dflash", "dspark") ): raise ValueError( "Parallel speculative drafting is not supported with watermarking." ) - if speculative_config.method not in ("dspark", "eagle", "eagle3", "mtp"): + if speculative_config.method not in ( + "dflash", + "dspark", + "eagle", + "eagle3", + "mtp", + ): raise ValueError( - "Watermarking supports only autoregressive model-based " - "speculative decoding." + "Watermarking supports only dflash, dspark, eagle, eagle3 and " + "mtp speculative decoding." + ) + if self._is_dflash_candidate_draft(): + raise ValueError( + "Watermarking does not support DFlash2 and LiLiCorr " + "candidate-head drafters yet." ) if ( not watermark_config.allow_target_only_watermarking @@ -2307,9 +2318,6 @@ class VllmConfig: worker is configured to handle. """ if self.speculative_config is not None: - scheduled_token_delta = ( - self.speculative_config.max_num_new_slots_for_drafting - ) max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens if self.scheduler_config.max_num_scheduled_tokens is None: self.scheduler_config.max_num_scheduled_tokens = max_num_batched_tokens @@ -2333,12 +2341,34 @@ class VllmConfig: " num_speculative_tokens.", ) - if max_num_batched_tokens <= scheduled_token_delta: + # num_lookahead_tokens is at least every method's per-request draft + # slots, and covers the DFlash/DSpark draft pass's query rows. + if max_num_batched_tokens <= self.num_lookahead_tokens: raise ValueError( "VllmConfig does not have enough slots to schedule a token and" " support the speculative decoding settings." - f" Got {max_num_batched_tokens=} and {scheduled_token_delta=}." + f" Got {max_num_batched_tokens=} and" + f" num_lookahead_tokens={self.num_lookahead_tokens}." ) + spec = self.speculative_config + if spec.use_dflash() or spec.use_dspark(): + max_num_active_seqs = ( + self.scheduler_config.max_num_active_seqs + or self.scheduler_config.max_num_seqs + ) + max_draft_seqs = max_num_batched_tokens // self.num_lookahead_tokens + if max_draft_seqs < max_num_active_seqs: + logger.warning_once( + "The %s draft pass needs %d query tokens per request within " + "max_num_batched_tokens=%d. Overriding max_num_active_seqs " + "from %d to %d.", + spec.method, + self.num_lookahead_tokens, + max_num_batched_tokens, + max_num_active_seqs, + max_draft_seqs, + ) + self.scheduler_config.max_num_active_seqs = max_draft_seqs def _set_cudagraph_sizes(self): """VLLM defines the default candidate list of batch sizes for CUDA graph @@ -3025,6 +3055,8 @@ class VllmConfig: unsupported.append("dspark speculative decoding") if self.speculative_config.enable_adaptive_verification: unsupported.append("adaptive draft verification") + if self.speculative_config.draft_lm_head_quantization is not None: + unsupported.append("draft lm_head quantization") # Mixed sliding/full DFlash drafts need multiple KV groups (V2 only). if self._dflash_needs_multi_kv_group(): diff --git a/vllm/distributed/elastic_ep/elastic_state.py b/vllm/distributed/elastic_ep/elastic_state.py index 5ba1891b7a..f9869b4da4 100644 --- a/vllm/distributed/elastic_ep/elastic_state.py +++ b/vllm/distributed/elastic_ep/elastic_state.py @@ -11,6 +11,9 @@ from vllm.config import ParallelConfig from vllm.distributed import ( stateless_destroy_torch_distributed_process_group, ) +from vllm.distributed.elastic_ep.external_elastic_ep import ( + ExternalElasticEPScaleCoordinator, +) from vllm.distributed.utils import get_cached_tcp_store_client from vllm.logger import init_logger from vllm.v1.engine import ( @@ -256,10 +259,23 @@ class ElasticEPScalingState: if state == ScaleDownRemovingEngineState.PREPARE: assert self.old_dp_group.rank() > 0 self._commit_scale_down(removing=True) - self.state = ScaleDownRemovingEngineState.COMPLETE + external_lb = self.vllm_config.parallel_config.data_parallel_external_lb + if external_lb: + # Preserve completion delivery: client cleanup on exit must not + # cancel the receiver before it records the notification. + parallel_config = self.new_parallel_config + store = get_cached_tcp_store_client( + parallel_config.data_parallel_master_ip, + parallel_config._coord_store_port, + ) + key = ExternalElasticEPScaleCoordinator.key + epoch = store.get(key("current_epoch")).decode() self.engine_core._eep_send_engine_core_notification( EEPNotificationType.SHUTDOWN_COMPLETE ) + if external_lb: + store.wait([key(epoch, "shutdown_complete", self.engine_core.dp_rank)]) + self.state = ScaleDownRemovingEngineState.COMPLETE return True else: @@ -384,7 +400,10 @@ class ElasticEPScalingState: def _send_reconfigure_finished(self): assert self.new_dp_group is not None - if self.new_dp_group.rank() == 0: + if ( + self.new_dp_group.rank() == 0 + or self.vllm_config.parallel_config.data_parallel_external_lb + ): self.engine_core._eep_send_engine_core_notification( EEPNotificationType.RECONFIGURE_FINISHED ) diff --git a/vllm/distributed/elastic_ep/external_elastic_ep.py b/vllm/distributed/elastic_ep/external_elastic_ep.py new file mode 100644 index 0000000000..6db6e5f4b1 --- /dev/null +++ b/vllm/distributed/elastic_ep/external_elastic_ep.py @@ -0,0 +1,897 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import asyncio +import contextlib +import enum +import uuid +from dataclasses import dataclass +from threading import Event, Thread +from typing import TYPE_CHECKING, Any + +import msgspec +import msgspec.msgpack +import zmq + +from vllm.logger import init_logger +from vllm.utils.network_utils import make_zmq_socket +from vllm.v1.engine import ( + EEPNotificationType, + ReconfigureDistributedRequest, + ReconfigureRankType, +) +from vllm.v1.engine.utils import EngineHandshakeMetadata, EngineZmqAddresses +from vllm.v1.utils import get_engine_client_zmq_addr + +if TYPE_CHECKING: + from vllm.v1.engine.core_client import DPAsyncMPClient + +logger = init_logger(__name__) +POLL_BACKOFF_STEPS_S = (0.1, 0.2, 0.4, 0.8) + + +class ExternalElasticEPScalePhase(str, enum.Enum): + """Lifecycle phases shared for an epoch-scoped external EEP operation. + + These values are stored in the control and reconfiguration stores and + exposed through the status API so external orchestrators can observe the + same operation across API processes and reconnects. COMPLETED means the + topology commit and rank completion barrier have finished; API processes + may clear their process-local request gate shortly afterward. + """ + + IDLE = "idle" + PREPARING = "preparing" + COMMITTING = "committing" + FAILED = "failed" + COMPLETED = "completed" + + +@dataclass +class ExternalElasticEPScaleUpHandshakeMetadata: + engine_metadata: EngineHandshakeMetadata + num_redundant_experts: int + elastic_ep_max_dp_size: int + + +class ExternalElasticEPScaleUpHandshakeServer: + """Temporary rank-0 handshake server for external EEP scale-up. + + During normal external startup the global handshake listener only exists + while the rank starts. Scale-up needs the same handshake contract again for + newly launched ranks, so rank 0 re-opens a temporary listener for the + duration of the current scale operation. + + The handshake gives new ranks the information they cannot infer locally: + front-end and DP coordinator ZMQ addresses, the new DP master address/ports, + target DP size, target expert redundancy, and the coordination store port + used for EEP reconfiguration. + """ + + def __init__( + self, + *, + handshake_address: str, + expected_new_ranks: list[int], + addresses: EngineZmqAddresses, + bootstrap: ReconfigureDistributedRequest, + num_redundant_experts: int, + elastic_ep_max_dp_size: int, + ) -> None: + self.handshake_address = handshake_address + self.expected_new_ranks = set(expected_new_ranks) + self.addresses = addresses + self.bootstrap = bootstrap + self.num_redundant_experts = num_redundant_experts + self.elastic_ep_max_dp_size = elastic_ep_max_dp_size + self.started_event = Event() + self._stop_event = Event() + self._thread = Thread( + target=self._run, + name="ExternalElasticEPHandshakeServer", + daemon=True, + ) + self._error: Exception | None = None + + def start(self) -> None: + self._thread.start() + started = self.started_event.wait(timeout=5) + if self._error is not None: + raise self._error + if not started: + raise TimeoutError( + "Timed out waiting for external EEP handshake server to " + f"start listening at {self.handshake_address}" + ) + + def stop(self) -> None: + self._stop_event.set() + if self._thread.is_alive(): + self._thread.join(timeout=5) + if self._error is not None: + raise self._error + + def raise_if_failed(self) -> None: + if self._error is not None: + raise self._error + + def _run(self) -> None: + pending_ready_ranks = self.expected_new_ranks.copy() + try: + with ( + zmq.Context() as ctx, + make_zmq_socket( + ctx, + self.handshake_address, + zmq.ROUTER, + bind=True, + linger=0, + router_handover=True, + ) as handshake_socket, + ): + self.started_event.set() + poller = zmq.Poller() + poller.register(handshake_socket, zmq.POLLIN) + + while not self._stop_event.is_set(): + if not pending_ready_ranks: + return + + events = dict(poller.poll(timeout=1000)) + if handshake_socket not in events: + continue + + engine_identity, payload = handshake_socket.recv_multipart() + engine_rank = int.from_bytes(engine_identity, "little") + if engine_rank not in self.expected_new_ranks: + raise RuntimeError( + "Received scale-up handshake from unexpected " + f"dp rank {engine_rank}" + ) + + message = msgspec.msgpack.decode(payload) + status = message["status"] + if status == "HELLO": + b = self.bootstrap + parallel_config: dict[str, int | str | list[int]] = { + "data_parallel_master_ip": b.new_data_parallel_master_ip, + "data_parallel_master_port": ( + b.new_data_parallel_master_port + ), + "_data_parallel_master_port_list": ( + b.new_data_parallel_master_port_list + ), + "data_parallel_size": b.new_data_parallel_size, + "_coord_store_port": b.coord_store_port, + } + init_message = msgspec.msgpack.encode( + ExternalElasticEPScaleUpHandshakeMetadata( + engine_metadata=EngineHandshakeMetadata( + addresses=self.addresses, + parallel_config=parallel_config, + ), + num_redundant_experts=self.num_redundant_experts, + elastic_ep_max_dp_size=self.elastic_ep_max_dp_size, + ) + ) + handshake_socket.send_multipart( + (engine_identity, init_message), copy=False + ) + elif status == "READY": + pending_ready_ranks.discard(engine_rank) + else: + raise RuntimeError( + f"Unexpected handshake status {status} from dp rank " + f"{engine_rank}" + ) + except Exception as e: + self._error = e + self.started_event.set() + + +@dataclass +class _PreparedExternalElasticEPScale: + control_store: Any + reconfig_store: Any + epoch: str + bootstrap: ReconfigureDistributedRequest + dp_rank: int + cur_data_parallel_size: int + num_redundant_experts: int + scale_up: bool + handshake_server: ExternalElasticEPScaleUpHandshakeServer | None + + +class ExternalElasticEPScaleCoordinator: + def __init__(self, client: "DPAsyncMPClient") -> None: + self.client = client + self.active_reconfig_store: tuple[str, int] | None = None + # Keep the previous rank-0 TCPStore server alive while the next scale + # operation publishes bootstrap metadata through it. Overwriting + # client._coord_store with the new store can otherwise close the old + # server while other ranks are still polling it. + self.control_store_ref: Any | None = None + self.reconfig_store_ref: Any | None = None + self.active_epoch: str | None = None + self.prepared_scale: _PreparedExternalElasticEPScale | None = None + + @staticmethod + def key(*parts: str | int) -> str: + return "/".join(["elastic_ep/external", *[str(part) for part in parts]]) + + @staticmethod + async def _sleep_with_backoff(backoff_step: int) -> int: + max_step = len(POLL_BACKOFF_STEPS_S) - 1 + backoff_step = max(0, min(backoff_step, max_step)) + await asyncio.sleep(POLL_BACKOFF_STEPS_S[backoff_step]) + return min(backoff_step + 1, max_step) + + def _update_parallel_config( + self, + bootstrap: ReconfigureDistributedRequest, + num_redundant_experts: int, + ) -> None: + parallel_config = self.client.vllm_config.parallel_config + parallel_config.data_parallel_size = bootstrap.new_data_parallel_size + parallel_config.data_parallel_master_ip = bootstrap.new_data_parallel_master_ip + parallel_config.data_parallel_master_port = ( + bootstrap.new_data_parallel_master_port + ) + parallel_config._data_parallel_master_port_list = ( + bootstrap.new_data_parallel_master_port_list.copy() + ) + parallel_config._coord_store_port = bootstrap.coord_store_port + parallel_config.eplb_config.num_redundant_experts = num_redundant_experts + + def _calculate_num_redundant_experts( + self, + cur_data_parallel_size: int, + new_data_parallel_size: int, + ) -> int: + parallel_config = self.client.vllm_config.parallel_config + num_experts = self.client.vllm_config.model_config.get_num_experts() + num_redundant_experts = ( + num_experts + parallel_config.eplb_config.num_redundant_experts + ) * new_data_parallel_size // cur_data_parallel_size - num_experts + if num_redundant_experts < 0: + raise ValueError( + "Cannot scale external Elastic EP from DP size " + f"{cur_data_parallel_size} to {new_data_parallel_size}: " + "the target topology cannot retain all model experts." + ) + return num_redundant_experts + + def _get_reconfig_store(self): + from vllm.distributed.utils import get_cached_tcp_store_client + + if self.active_reconfig_store is not None: + store_addr = self.active_reconfig_store + else: + parallel_config = self.client.vllm_config.parallel_config + if not parallel_config._coord_store_port: + raise RuntimeError( + "External Elastic EP requires an active reconfiguration " + "coordination store." + ) + store_addr = ( + parallel_config.data_parallel_master_ip, + parallel_config._coord_store_port, + ) + + return get_cached_tcp_store_client(*store_addr) + + def _get_existing_engine_zmq_address(self) -> EngineZmqAddresses: + coordinator_input = self.client.coordinator_input_address + coordinator_output = self.client.coordinator_output_address + if coordinator_input is None or coordinator_output is None: + raise RuntimeError( + "External Elastic EP scale-up requires DP coordinator addresses." + ) + + return EngineZmqAddresses( + inputs=[self.client.input_address], + outputs=[self.client.output_address], + coordinator_input=coordinator_input, + coordinator_output=coordinator_output, + frontend_stats_publish_address=self.client.stats_update_address, + ) + + def _setup_reconfig_bootstrap(self) -> tuple[str, int]: + from vllm.distributed.utils import create_tcp_store + + self.control_store_ref = getattr(self.client, "_coord_store", None) + ip = self.client.vllm_config.parallel_config.data_parallel_master_ip + # Keep the serving configuration on the old store until commit succeeds. + self.reconfig_store_ref = create_tcp_store( + ip, 0, is_master=True, world_size=-1, wait_for_workers=False + ) + return ip, self.reconfig_store_ref.port + + def _get_error(self, store: Any, epoch: str) -> str | None: + error_key = self.key(epoch, "error") + if not store.check([error_key]): + return None + return store.get(error_key).decode() + + def get_status(self) -> dict[str, str | int | None]: + store = self._get_reconfig_store() + current_epoch_key = self.key("current_epoch") + if not store.check([current_epoch_key]): + return { + "phase": ExternalElasticEPScalePhase.IDLE.value, + "epoch": None, + "error": None, + "requested_data_parallel_size": None, + } + + epoch = store.get(current_epoch_key).decode() + error = self._get_error(store, epoch) + if error is not None: + phase = ExternalElasticEPScalePhase.FAILED + elif store.check([self.key(epoch, "completed")]): + phase = ExternalElasticEPScalePhase.COMPLETED + elif store.check([self.key(epoch, "phase")]): + phase = ExternalElasticEPScalePhase( + store.get(self.key(epoch, "phase")).decode() + ) + else: + phase = ExternalElasticEPScalePhase.IDLE + requested_dp_size = None + bootstrap_key = self.key(epoch, "bootstrap") + if store.check([bootstrap_key]): + bootstrap = msgspec.msgpack.decode( + store.get(bootstrap_key), type=ReconfigureDistributedRequest + ) + requested_dp_size = bootstrap.new_data_parallel_size + return { + "phase": phase.value, + "epoch": epoch, + "error": error, + "requested_data_parallel_size": requested_dp_size, + } + + def _set_phase( + self, store: Any, epoch: str, phase: ExternalElasticEPScalePhase + ) -> None: + store.set(self.key(epoch, "phase"), phase.value.encode()) + + async def _wait_for_bootstrap( + self, + store: Any, + requested_new_dp_size: int, + timeout_s: float = 300, + ) -> tuple[str, ReconfigureDistributedRequest]: + """Wait for rank 0 to publish matching scale bootstrap metadata. + Non-zero ranks poll the old control store until prepared or errored.""" + loop = asyncio.get_running_loop() + start = loop.time() + current_epoch_key = self.key("current_epoch") + previous_epoch = self.active_epoch + if store.check([current_epoch_key]): + current_epoch = store.get(current_epoch_key).decode() + if store.check([self.key(current_epoch, "completed")]): + previous_epoch = current_epoch + backoff_step = 0 + while True: + if store.check([current_epoch_key]): + epoch = store.get(current_epoch_key).decode() + if epoch != previous_epoch: + error = self._get_error(store, epoch) + if error is not None: + self.active_epoch = epoch + raise RuntimeError(error) + + bootstrap_key = self.key(epoch, "bootstrap") + if store.check([bootstrap_key]): + bootstrap = msgspec.msgpack.decode( + store.get(bootstrap_key), + type=ReconfigureDistributedRequest, + ) + prepared_key = self.key(epoch, "prepared") + completed = store.check([self.key(epoch, "completed")]) + if bootstrap.new_data_parallel_size != requested_new_dp_size: + if store.check([prepared_key]) and not completed: + raise RuntimeError( + "A different external Elastic EP scaling " + "operation is already in progress for target " + f"dp size {bootstrap.new_data_parallel_size}." + ) + elif store.check([prepared_key]): + return epoch, bootstrap + + now = loop.time() + if now - start > timeout_s: + raise TimeoutError( + "Timed out waiting for rank 0 to publish external Elastic EP " + "bootstrap metadata." + ) + backoff_step = await self._sleep_with_backoff(backoff_step) + + def _prepare_reconfig_bootstrap( + self, + store: Any, + cur_data_parallel_size: int, + new_data_parallel_size: int, + ) -> tuple[str, ReconfigureDistributedRequest]: + from vllm.utils.network_utils import get_open_ports_list + + current_epoch_key = self.key("current_epoch") + if store.check([current_epoch_key]): + current_epoch = store.get(current_epoch_key).decode() + current_error = self._get_error(store, current_epoch) + if current_error is None and not store.check( + [self.key(current_epoch, "completed")] + ): + raise RuntimeError( + "Another external Elastic EP scaling operation is already active." + ) + + ip, coord_store_port = self._setup_reconfig_bootstrap() + epoch = uuid.uuid4().hex + master_ports = get_open_ports_list(5) + bootstrap = ReconfigureDistributedRequest( + new_data_parallel_size=new_data_parallel_size, + new_data_parallel_rank=ReconfigureRankType.KEEP_CURRENT_RANK, + new_data_parallel_rank_local=ReconfigureRankType.KEEP_CURRENT_RANK, + new_data_parallel_master_ip=ip, + new_data_parallel_master_port=master_ports.pop(), + new_data_parallel_master_port_list=master_ports, + coord_store_port=coord_store_port, + ) + + bootstrap_key = self.key(epoch, "bootstrap") + bootstrap_payload = msgspec.msgpack.encode(bootstrap) + stores = [store] + if self.reconfig_store_ref is not None: + stores.append(self.reconfig_store_ref) + for target_store in stores: + target_store.set(bootstrap_key, bootstrap_payload) + self._set_phase(target_store, epoch, ExternalElasticEPScalePhase.PREPARING) + # Publish the epoch only after its metadata is ready for readers. + target_store.set(current_epoch_key, epoch.encode()) + return epoch, bootstrap + + def _start_scale_up_handshake_server( + self, + bootstrap: ReconfigureDistributedRequest, + cur_data_parallel_size: int, + num_redundant_experts: int, + elastic_ep_max_dp_size: int, + ) -> ExternalElasticEPScaleUpHandshakeServer: + handshake_server = ExternalElasticEPScaleUpHandshakeServer( + handshake_address=get_engine_client_zmq_addr( + False, + bootstrap.new_data_parallel_master_ip, + self.client.vllm_config.parallel_config.data_parallel_rpc_port, + ), + expected_new_ranks=list( + range( + cur_data_parallel_size, + bootstrap.new_data_parallel_size, + ) + ), + addresses=self._get_existing_engine_zmq_address(), + bootstrap=bootstrap, + num_redundant_experts=num_redundant_experts, + elastic_ep_max_dp_size=elastic_ep_max_dp_size, + ) + handshake_server.start() + return handshake_server + + async def _wait_for_ready_ranks( + self, + control_store: Any, + reconfig_store: Any, + epoch: str, + new_data_parallel_size: int, + handshake_server: ExternalElasticEPScaleUpHandshakeServer | None = None, + timeout_s: float = 300, + ) -> None: + """Wait until every rank in the new DP group is ready to commit.""" + loop = asyncio.get_running_loop() + start = loop.time() + ready_keys = [f"eep_ready/{rank}" for rank in range(new_data_parallel_size)] + + backoff_step = 0 + while True: + if handshake_server is not None: + handshake_server.raise_if_failed() + + error = self._get_error(control_store, epoch) + if error is not None: + raise RuntimeError( + "External Elastic EP scaling failed while waiting for " + f"ranks to prepare: {error}" + ) + + if reconfig_store.check(ready_keys): + return + + now = loop.time() + if now - start > timeout_s: + raise TimeoutError( + "Timed out waiting for external Elastic EP ranks to prepare." + ) + backoff_step = await self._sleep_with_backoff(backoff_step) + + async def _wait_for_local_reconfig_finished( + self, + control_store: Any, + reconfig_store: Any, + epoch: str, + bootstrap: ReconfigureDistributedRequest, + dp_rank: int, + scale_up: bool, + timeout_s: float = 600, + ) -> None: + """Wait until this rank's EngineCore finishes the scale transition.""" + if not scale_up and dp_rank >= bootstrap.new_data_parallel_size: + wait_key = self.key(epoch, "shutdown_complete", dp_rank) + timeout_msg = ( + "Timed out waiting for local external Elastic EP scale-down " + "shutdown to finish." + ) + else: + wait_key = self.key(epoch, "old_rank_finished", dp_rank) + timeout_msg = ( + "Timed out waiting for local external Elastic EP " + "reconfiguration to finish." + ) + + loop = asyncio.get_running_loop() + start = loop.time() + backoff_step = 0 + while True: + error = self._get_error(control_store, epoch) + if error is not None: + raise RuntimeError( + error or "External Elastic EP scaling failed on another rank." + ) + + if reconfig_store.check([wait_key]): + return + + now = loop.time() + if now - start > timeout_s: + raise TimeoutError(timeout_msg) + backoff_step = await self._sleep_with_backoff(backoff_step) + + async def _wait_for_all_old_ranks( + self, + prepared: _PreparedExternalElasticEPScale, + timeout_s: float = 600, + ) -> None: + bootstrap = prepared.bootstrap + keys = [ + self.key( + prepared.epoch, + ( + "shutdown_complete" + if not prepared.scale_up + and rank >= bootstrap.new_data_parallel_size + else "old_rank_finished" + ), + rank, + ) + for rank in range(prepared.cur_data_parallel_size) + ] + loop = asyncio.get_running_loop() + start = loop.time() + backoff_step = 0 + while True: + error = self._get_error(prepared.control_store, prepared.epoch) + if error is not None: + raise RuntimeError(error) + if prepared.reconfig_store.check(keys): + return + if loop.time() - start > timeout_s: + raise TimeoutError( + "Timed out waiting for all old ranks to finish external " + "Elastic EP scaling." + ) + backoff_step = await self._sleep_with_backoff(backoff_step) + + @staticmethod + def _stop_handshake_server( + server: ExternalElasticEPScaleUpHandshakeServer | None, + suppress_errors: bool, + ) -> None: + if server is None: + return + if suppress_errors: + with contextlib.suppress(Exception): + server.stop() + else: + server.stop() + + def _publish_error( + self, + prepared: _PreparedExternalElasticEPScale, + error: BaseException, + ) -> None: + error_key = self.key(prepared.epoch, "error") + error_payload = (str(error) or type(error).__name__).encode() + for store in (prepared.control_store, prepared.reconfig_store): + with contextlib.suppress(Exception): + store.set(error_key, error_payload) + self._set_phase( + store, prepared.epoch, ExternalElasticEPScalePhase.FAILED + ) + + def abort(self, error: str) -> None: + prepared = self.prepared_scale + if prepared is not None: + self._publish_error(prepared, RuntimeError(error)) + self._stop_handshake_server(prepared.handshake_server, suppress_errors=True) + self.prepared_scale = None + + async def _wait_for_commit_ready( + self, prepared: _PreparedExternalElasticEPScale, timeout_s: float = 600 + ) -> None: + # No engine may enter the collective pause until all old API ranks + # have finished draining; a drain failure must reach waiting peers. + store = prepared.reconfig_store + await asyncio.to_thread( + store.set, self.key(prepared.epoch, "commit_ready", prepared.dp_rank), b"1" + ) + keys = [ + self.key(prepared.epoch, "commit_ready", rank) + for rank in range(prepared.cur_data_parallel_size) + ] + + def is_ready() -> bool: + error = self._get_error(prepared.control_store, prepared.epoch) + if error is not None: + raise RuntimeError(error) + return store.check(keys) + + start = asyncio.get_running_loop().time() + backoff_step = 0 + while True: + if await asyncio.to_thread(is_ready): + return + if asyncio.get_running_loop().time() - start > timeout_s: + raise TimeoutError("Timed out waiting for external Elastic EP drains.") + backoff_step = await self._sleep_with_backoff(backoff_step) + + async def prepare( + self, cur_data_parallel_size: int, new_data_parallel_size: int + ) -> None: + from vllm.distributed.utils import get_cached_tcp_store_client + + if self.prepared_scale is not None: + if ( + self.prepared_scale.bootstrap.new_data_parallel_size + == new_data_parallel_size + ): + return + raise RuntimeError("External Elastic EP scaling is already prepared.") + + parallel_config = self.client.vllm_config.parallel_config + dp_rank = parallel_config.data_parallel_rank + scale_up = new_data_parallel_size > cur_data_parallel_size + if new_data_parallel_size > parallel_config.elastic_ep_max_dp_size: + raise ValueError( + f"Cannot scale to data_parallel_size {new_data_parallel_size}; " + "--elastic-ep-max-dp-size is " + f"{parallel_config.elastic_ep_max_dp_size}." + ) + # Validate expert capacity on every rank: rank 0 publishes no epoch + # when validation fails, so peers must reject before polling for one. + num_redundant_experts = self._calculate_num_redundant_experts( + cur_data_parallel_size, + new_data_parallel_size, + ) + if not parallel_config._coord_store_port: + raise RuntimeError( + "External Elastic EP requires a runtime coordination store port." + ) + control_store = get_cached_tcp_store_client( + parallel_config.data_parallel_master_ip, + parallel_config._coord_store_port, + ) + handshake_server: ExternalElasticEPScaleUpHandshakeServer | None = None + bootstrap: ReconfigureDistributedRequest | None = None + epoch: str | None = None + reconfig_store: Any | None = None + + try: + if dp_rank == 0: + epoch, bootstrap = self._prepare_reconfig_bootstrap( + control_store, + cur_data_parallel_size, + new_data_parallel_size, + ) + control_store.set( + self.key(epoch, "num_redundant_experts"), + str(num_redundant_experts).encode(), + ) + if scale_up: + handshake_server = self._start_scale_up_handshake_server( + bootstrap, + cur_data_parallel_size, + num_redundant_experts, + parallel_config.elastic_ep_max_dp_size, + ) + control_store.set(self.key(epoch, "prepared"), b"1") + else: + epoch, bootstrap = await self._wait_for_bootstrap( + control_store, new_data_parallel_size + ) + num_redundant_experts = int( + control_store.get(self.key(epoch, "num_redundant_experts")).decode() + ) + + assert epoch is not None + assert bootstrap is not None + self.active_epoch = epoch + self.active_reconfig_store = ( + bootstrap.new_data_parallel_master_ip, + bootstrap.coord_store_port, + ) + reconfig_store = self._get_reconfig_store() + reconfig_store.set(self.key("current_epoch"), epoch.encode()) + self._set_phase( + reconfig_store, epoch, ExternalElasticEPScalePhase.PREPARING + ) + + if scale_up or dp_rank < bootstrap.new_data_parallel_size: + await self.client.call_utility_async( + "reinitialize_distributed", bootstrap + ) + + await self._wait_for_ready_ranks( + control_store, + reconfig_store, + epoch, + bootstrap.new_data_parallel_size, + handshake_server, + ) + self.prepared_scale = _PreparedExternalElasticEPScale( + control_store=control_store, + reconfig_store=reconfig_store, + epoch=epoch, + bootstrap=bootstrap, + dp_rank=dp_rank, + cur_data_parallel_size=cur_data_parallel_size, + num_redundant_experts=num_redundant_experts, + scale_up=scale_up, + handshake_server=handshake_server, + ) + except BaseException as e: + if epoch is not None: + self.active_epoch = epoch + error_key = self.key(epoch, "error") + error_payload = (str(e) or type(e).__name__).encode() + with contextlib.suppress(Exception): + control_store.set(error_key, error_payload) + self._set_phase( + control_store, epoch, ExternalElasticEPScalePhase.FAILED + ) + failed_reconfig_store = ( + reconfig_store + if reconfig_store is not None + else self.reconfig_store_ref + ) + if failed_reconfig_store is not None: + with contextlib.suppress(Exception): + failed_reconfig_store.set(error_key, error_payload) + self._set_phase( + failed_reconfig_store, + epoch, + ExternalElasticEPScalePhase.FAILED, + ) + self.active_reconfig_store = None + self._stop_handshake_server(handshake_server, suppress_errors=True) + raise + + async def commit(self) -> None: + prepared = self.prepared_scale + if prepared is None: + raise RuntimeError("External Elastic EP scaling has not been prepared.") + + bootstrap = prepared.bootstrap + remaining = prepared.scale_up or ( + prepared.dp_rank < bootstrap.new_data_parallel_size + ) + try: + await self._wait_for_commit_ready(prepared) + for store in (prepared.control_store, prepared.reconfig_store): + self._set_phase( + store, prepared.epoch, ExternalElasticEPScalePhase.COMMITTING + ) + await self.client.pause_scheduler_async( + mode=self.client._eep_commit_pause_mode() if remaining else "abort", + clear_cache=False, + ) + if remaining: + await self.client.call_utility_async("commit_prepared_elastic_ep") + else: + reconfig_request = msgspec.structs.replace( + bootstrap, + new_data_parallel_rank=ReconfigureRankType.SHUTDOWN_CURRENT_RANK, + ) + await self.client.call_utility_async( + "reinitialize_distributed", reconfig_request + ) + + await self._wait_for_local_reconfig_finished( + prepared.control_store, + prepared.reconfig_store, + prepared.epoch, + bootstrap, + prepared.dp_rank, + prepared.scale_up, + ) + if prepared.dp_rank == 0: + await self._wait_for_all_old_ranks(prepared) + completed_key = self.key(prepared.epoch, "completed") + prepared.control_store.set(completed_key, b"1") + prepared.reconfig_store.set(completed_key, b"1") + for store in (prepared.control_store, prepared.reconfig_store): + self._set_phase( + store, prepared.epoch, ExternalElasticEPScalePhase.COMPLETED + ) + + if remaining: + self._update_parallel_config(bootstrap, prepared.num_redundant_experts) + if prepared.dp_rank == 0: + self.client._coord_store = self.reconfig_store_ref + self.client._ensure_stats_update_task() + await self.client.first_req_send_socket.send( + msgspec.msgpack.encode( + ("SCALE_ELASTIC_EP", bootstrap.new_data_parallel_size) + ) + ) + await self.client.resume_scheduler_async() + self._stop_handshake_server( + prepared.handshake_server, suppress_errors=False + ) + except BaseException as e: + self._publish_error(prepared, e) + self._stop_handshake_server(prepared.handshake_server, suppress_errors=True) + raise + finally: + self.prepared_scale = None + + async def process_engine_core_notification( + self, notification_data: tuple[str, int] + ) -> None: + """Record scale notifications emitted by EngineCore processes. + The stored keys are later polled by the external scale coordinator.""" + parallel_config = self.client.vllm_config.parallel_config + if not ( + parallel_config.enable_elastic_ep + and parallel_config.data_parallel_external_lb + ): + return + + notification_type_str, dp_rank = notification_data + notification_type = EEPNotificationType(notification_type_str) + if not parallel_config._coord_store_port: + logger.warning( + "Ignoring external Elastic EP notification %s because coord " + "store metadata is not available yet.", + notification_type.value, + ) + return + + epoch = self.active_epoch + reconfig_store = self._get_reconfig_store() + if epoch is None: + current_epoch_key = self.key("current_epoch") + if not reconfig_store.check([current_epoch_key]): + logger.warning( + "Ignoring external Elastic EP notification %s because " + "active epoch metadata is not available yet.", + notification_type.value, + ) + return + epoch = reconfig_store.get(current_epoch_key).decode() + + if notification_type == EEPNotificationType.RECONFIGURE_FINISHED: + key = self.key(epoch, "old_rank_finished", dp_rank) + reconfig_store.set(key, b"1") + elif notification_type == EEPNotificationType.SHUTDOWN_COMPLETE: + key = self.key(epoch, "shutdown_complete", dp_rank) + reconfig_store.set(key, b"1") diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py index c53be27079..6e3d3415c3 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py @@ -28,6 +28,7 @@ from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy from vllm.v1.hisparse.layout import HISPARSE_HOT_SUFFIX from vllm.v1.hisparse.runtime import ( HiSparseCacheHandle, + drop_hisparse_host_pool, release_pinned_state, update_hisparse_residency, ) @@ -254,6 +255,7 @@ class HiSparseConnectorWorker: pinned_host_pools, shared_host_region, ) + drop_hisparse_host_pool(self.vllm_config) raise def initialize( @@ -972,4 +974,5 @@ class HiSparseConnectorWorker: self.pinned_host_pools, self.shared_host_region, ) + drop_hisparse_host_pool(self.vllm_config) self._initialized = False diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index ed0e6eafe6..4e86000bb1 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -82,6 +82,7 @@ from vllm.utils.math_utils import cdiv from vllm.utils.network_utils import make_zmq_path from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.kv_cache_interface import ( + AttentionSpec, CircularBufferSpec, FullAttentionSpec, KVCacheLayout, @@ -2671,6 +2672,22 @@ class NixlBaseConnectorWorker: "Heterogeneous TP head-dimension splitting requires contiguous heads. " "Use a block-contiguous layout (e.g. LBHNC) on the prefill side." ) + # NVFP4 pages store all heads' data before all heads' scales, so a + # head's bytes are not a contiguous slice of the page. + if ( + not self.use_mla + and self.transfer_topo.local_physical_heads + != max(1, self.transfer_topo.total_num_kv_heads // remote_tp_size) + and any( + isinstance(spec, AttentionSpec) and spec.kv_quant_mode.is_nvfp4 + for spec in self._layer_specs.values() + ) + ): + raise RuntimeError( + "NIXL KV-head re-sharding is not supported for NVFP4 KV cache " + f"(local TP {self.transfer_topo.tp_size}, remote TP " + f"{remote_tp_size}). Use matching TP sizes on prefill and decode." + ) # Per-region block_len validation enforcing the P/D invariant. # REPLICATE regions (MLA, or a whole-model MLA / replicated-KV transfer) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/__init__.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/__init__.py new file mode 100644 index 0000000000..06eee4bfab --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/__init__.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from .connector import UMBPStoreConnector +from .data import ( + BlockIdentityCodec, + BlockTransferPlan, + KVLayoutDescriptor, + KVLayoutPlanner, + KVRange, + KVRegion, + LoadSpec, + RankTopology, + RequestTracker, + TransferJobState, + TransferJobStatus, + UMBPConnectorMetadata, + UMBPConnectorWorkerMetadata, + UMBPNamespace, +) +from .runtime import UMBPRuntimeConfig, UMBPRuntimeFactory + +__all__ = [ + "UMBPStoreConnector", + "BlockIdentityCodec", + "BlockTransferPlan", + "KVLayoutPlanner", + "KVLayoutDescriptor", + "KVRange", + "KVRegion", + "LoadSpec", + "RequestTracker", + "RankTopology", + "TransferJobState", + "TransferJobStatus", + "UMBPConnectorMetadata", + "UMBPConnectorWorkerMetadata", + "UMBPNamespace", + "UMBPRuntimeConfig", + "UMBPRuntimeFactory", +] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/connector.py new file mode 100644 index 0000000000..177d9fca11 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/connector.py @@ -0,0 +1,220 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""vLLM-facing shared UMBP connector.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import torch + +from vllm.config import VllmConfig +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorBase_V1, + KVConnectorMetadata, + KVConnectorRole, + KVConnectorTransferResults, +) +from vllm.forward_context import ForwardContext +from vllm.v1.attention.backend import AttentionMetadata +from vllm.v1.core.sched.output import SchedulerOutput + +from .data import ( + BlockIdentityCodec, + KVLayoutPlanner, + RankTopology, + UMBPConnectorMetadata, + UMBPConnectorWorkerMetadata, + UMBPNamespace, +) +from .runtime import UMBPRuntimeConfig, UMBPRuntimeFactory +from .scheduler import UMBPStoreConnectorScheduler +from .worker import UMBPStoreConnectorWorker + +if TYPE_CHECKING: + from vllm.v1.core.kv_cache_manager import KVCacheBlocks + from vllm.v1.kv_cache_interface import KVCacheConfig + from vllm.v1.request import Request + + +class UMBPStoreConnector(KVConnectorBase_V1): + """vLLM connector lifecycle shared by every UMBP runtime mode.""" + + def __init__( + self, + vllm_config: VllmConfig, + role: KVConnectorRole, + kv_cache_config: KVCacheConfig, + ) -> None: + super().__init__(vllm_config, role, kv_cache_config) + topology = RankTopology.from_vllm_config(vllm_config) + runtime_config = UMBPRuntimeConfig.from_vllm( + vllm_config + ).resolve_for_rank_count(topology.rank_count) + if topology.pp_size > 1: + raise NotImplementedError( + f"{runtime_config.mode} UMBP does not support pipeline parallelism" + ) + runtime = UMBPRuntimeFactory.build(runtime_config) + namespace = UMBPNamespace.from_vllm_config(vllm_config, kv_cache_config) + codec = BlockIdentityCodec( + namespace=namespace, + tp_rank=topology.tp_rank, + pp_rank=topology.pp_rank, + pcp_rank=topology.pcp_rank, + dcp_rank=topology.dcp_rank, + ) + layout = KVLayoutPlanner.from_kv_cache_config(kv_cache_config) + layout_descriptor = layout.describe(topology) + self._runtime = runtime + self._store_submitted = False + self.connector_scheduler: UMBPStoreConnectorScheduler | None = None + self.connector_worker: UMBPStoreConnectorWorker | None = None + if role == KVConnectorRole.SCHEDULER: + self.connector_scheduler = UMBPStoreConnectorScheduler( + vllm_config, + kv_cache_config, + runtime.create_scheduler_handle( + namespace.value, topology, layout_descriptor + ), + codec, + topology, + ) + else: + self.connector_worker = UMBPStoreConnectorWorker( + runtime.create_worker_handle( + namespace.value, topology, layout_descriptor + ), + layout, + codec=codec, + ) + + def get_num_new_matched_tokens( + self, request: Request, num_computed_tokens: int + ) -> tuple[int | None, bool]: + assert self.connector_scheduler is not None + return self.connector_scheduler.get_num_new_matched_tokens( + request, num_computed_tokens + ) + + def update_state_after_alloc( + self, + request: Request, + blocks: KVCacheBlocks, + num_external_tokens: int, + ) -> None: + assert self.connector_scheduler is not None + self.connector_scheduler.update_state_after_alloc( + request, blocks, num_external_tokens + ) + + def bind_gpu_block_pool(self, gpu_block_pool: Any) -> None: + assert self.connector_scheduler is not None + self.connector_scheduler.bind_gpu_block_pool(gpu_block_pool) + + def build_connector_meta( + self, scheduler_output: SchedulerOutput + ) -> KVConnectorMetadata: + assert self.connector_scheduler is not None + return self.connector_scheduler.build_connector_meta(scheduler_output) + + def request_finished( + self, request: Request, block_ids: list[int] + ) -> tuple[bool, dict[str, Any] | None]: + assert self.connector_scheduler is not None + return self.connector_scheduler.request_finished(request, (block_ids,)) + + def update_connector_output(self, connector_output: Any) -> None: + assert self.connector_scheduler is not None + self.connector_scheduler.update_connector_output(connector_output) + + def has_pending_push_work(self) -> bool: + assert self.connector_scheduler is not None + return self.connector_scheduler.has_pending_push_work() + + def reset_cache(self) -> bool | None: + if self.connector_scheduler is None: + return None + return self.connector_scheduler.reset_store() + + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: + assert self.connector_worker is not None + self.connector_worker.register_kv_caches(kv_caches) + + def bind_connector_metadata(self, connector_metadata: KVConnectorMetadata) -> None: + super().bind_connector_metadata(connector_metadata) + self._store_submitted = False + + def clear_connector_metadata(self) -> None: + # No-forward steps skip the model runner's wait_for_save hook. + if self.connector_worker is not None and self.has_connector_metadata(): + self.wait_for_save() + super().clear_connector_metadata() + + def handle_preemptions(self, kv_connector_metadata: KVConnectorMetadata) -> None: + assert self.connector_worker is not None + assert isinstance(kv_connector_metadata, UMBPConnectorMetadata) + self.connector_worker.handle_preemptions(kv_connector_metadata) + + def start_load_kv(self, forward_context: ForwardContext, **kwargs: Any) -> None: + del kwargs + assert self.connector_worker is not None + metadata = self._get_connector_metadata() + assert isinstance(metadata, UMBPConnectorMetadata) + self.connector_worker.start_load_kv(forward_context, metadata) + + def wait_for_layer_load(self, layer_name: str) -> None: + assert self.connector_worker is not None + self.connector_worker.wait_for_layer_load(layer_name) + + def save_kv_layer( + self, + layer_name: str, + kv_layer: torch.Tensor, + attn_metadata: AttentionMetadata, + **kwargs: Any, + ) -> None: + # Stores are submitted once per step, from wait_for_save. + return + + def wait_for_save(self) -> None: + if self._store_submitted: + return + assert self.connector_worker is not None + metadata = self._get_connector_metadata() + assert isinstance(metadata, UMBPConnectorMetadata) + self.connector_worker.enqueue_stores(metadata) + self.connector_worker.wait_for_save() + self._store_submitted = True + + def get_finished( + self, finished_req_ids: set[str] + ) -> tuple[set[str] | None, set[str] | None]: + assert self.connector_worker is not None + return self.connector_worker.get_finished(finished_req_ids) + + def get_transfer_results( + self, finished_req_ids: set[str] + ) -> KVConnectorTransferResults: + assert self.connector_worker is not None + finished_sending, finished_recving = self.connector_worker.get_finished( + finished_req_ids + ) + return KVConnectorTransferResults( + finished_sending=set(finished_sending or ()), + finished_recving=set(finished_recving or ()), + ) + + def get_block_ids_with_load_errors(self) -> set[int]: + assert self.connector_worker is not None + return self.connector_worker.get_block_ids_with_load_errors() + + def build_connector_worker_meta(self) -> UMBPConnectorWorkerMetadata: + assert self.connector_worker is not None + return self.connector_worker.build_connector_worker_meta() + + def shutdown(self) -> None: + if self.connector_scheduler is not None: + self.connector_scheduler.close() + if self.connector_worker is not None: + self.connector_worker.close() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py new file mode 100644 index 0000000000..4fa25d4071 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/data.py @@ -0,0 +1,620 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared data types for the UMBP KV connector. + +The types in this module deliberately do not depend on a UMBP client. They +describe vLLM's logical KV objects and the physical ranges a runtime adapter +will transfer. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Sequence +from dataclasses import dataclass, field, replace +from enum import Enum +from typing import Any, NamedTuple + +import torch + +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorMetadata, + KVConnectorWorkerMetadata, +) +from vllm.v1.kv_cache_interface import KVCacheConfig + + +@dataclass(frozen=True) +class RankTopology: + """Model-parallel ranks that participate in one KV object.""" + + tp_rank: int = 0 + tp_size: int = 1 + pp_rank: int = 0 + pp_size: int = 1 + pcp_rank: int = 0 + pcp_size: int = 1 + dcp_rank: int = 0 + dcp_size: int = 1 + + def __post_init__(self) -> None: + for name in ("tp_size", "pp_size", "pcp_size", "dcp_size"): + if getattr(self, name) <= 0: + raise ValueError(f"{name} must be positive") + for rank, size, name in ( + (self.tp_rank, self.tp_size, "tp_rank"), + (self.pp_rank, self.pp_size, "pp_rank"), + (self.pcp_rank, self.pcp_size, "pcp_rank"), + (self.dcp_rank, self.dcp_size, "dcp_rank"), + ): + if rank < 0 or rank >= size: + raise ValueError(f"{name}={rank} is outside [0, {size})") + + @classmethod + def from_vllm_config(cls, vllm_config: Any) -> RankTopology: + parallel = vllm_config.parallel_config + tp_size = getattr(parallel, "tensor_parallel_size", 1) + pp_size = getattr(parallel, "pipeline_parallel_size", 1) + pcp_size = getattr(parallel, "prefill_context_parallel_size", 1) + # An engine's workers are laid out as [PP, PCP, TP], TP innermost + # (initialize_model_parallel). + rank = getattr(parallel, "rank", 0) + tp_rank = rank % tp_size + dcp_size = getattr(parallel, "decode_context_parallel_size", 1) + return cls( + tp_rank=tp_rank, + tp_size=tp_size, + pp_rank=rank // (tp_size * pcp_size) % pp_size, + pp_size=pp_size, + pcp_rank=rank // tp_size % pcp_size, + pcp_size=pcp_size, + # DCP partitions the TP workers; it does not create a second + # Cartesian worker dimension. ParallelConfig is constructed before + # process groups, so derive this stable identity from TP rank. + dcp_rank=tp_rank % dcp_size, + dcp_size=dcp_size, + ) + + @property + def local_namespace(self) -> tuple[int, int, int, int]: + return (self.tp_rank, self.pcp_rank, self.dcp_rank, self.pp_rank) + + @property + def rank_count(self) -> int: + return len(self.all_namespaces()) + + def all_namespaces(self) -> tuple[tuple[int, int, int, int], ...]: + if self.dcp_size > 1: + if self.tp_size % self.dcp_size: + raise ValueError("dcp_size must divide tp_size") + return tuple( + (tp, pcp, tp % self.dcp_size, pp) + for pp in range(self.pp_size) + for pcp in range(self.pcp_size) + for tp in range(self.tp_size) + ) + return tuple( + (tp, pcp, dcp, pp) + for pp in range(self.pp_size) + for dcp in range(self.dcp_size) + for pcp in range(self.pcp_size) + for tp in range(self.tp_size) + ) + + +@dataclass(frozen=True) +class UMBPNamespace: + """Compatibility namespace for objects stored outside the engine.""" + + value: str + + @classmethod + def from_vllm_config( + cls, + vllm_config: Any, + kv_cache_config: KVCacheConfig, + ) -> UMBPNamespace: + extra = vllm_config.kv_transfer_config.kv_connector_extra_config + configured = extra.get("key_namespace", "auto") + if configured != "auto": + if not isinstance(configured, str) or not configured: + raise ValueError("key_namespace must be a non-empty string or 'auto'") + return cls(configured) + + model_config = vllm_config.model_config + parallel_config = vllm_config.parallel_config + # Draft-model KV (e.g. DSpark/EAGLE layers) is stored alongside the + # target's, so a different drafter must not reuse it. + speculative_config = getattr(vllm_config, "speculative_config", None) + draft_model_config = getattr(speculative_config, "draft_model_config", None) + speculative = ( + None + if speculative_config is None + else { + "method": getattr(speculative_config, "method", None), + "model": getattr(draft_model_config, "model", None), + "revision": getattr(draft_model_config, "revision", None), + } + ) + fields = { + "model": getattr(model_config, "model", ""), + "revision": getattr(model_config, "revision", None), + "block_size": vllm_config.cache_config.block_size, + "cache_layout": getattr(kv_cache_config, "kv_cache_layout", None), + "groups": [ + { + "layers": group.layer_names, + "block_size": group.kv_cache_spec.block_size, + "kind": type(group.kv_cache_spec).__name__, + "dtype": str(getattr(group.kv_cache_spec, "dtype", None)), + } + for group in kv_cache_config.kv_cache_groups + ], + "tp": parallel_config.tensor_parallel_size, + "pp": parallel_config.pipeline_parallel_size, + "pcp": getattr(parallel_config, "prefill_context_parallel_size", 1), + "dcp": getattr(parallel_config, "decode_context_parallel_size", 1), + "layout_version": 1, + } + if speculative is not None: + # Added only when set, so namespaces without a drafter are unchanged. + fields["speculative"] = speculative + payload = json.dumps(fields, sort_keys=True, separators=(",", ":")).encode() + return cls(hashlib.sha256(payload).hexdigest()[:32]) + + +@dataclass(frozen=True) +class BlockIdentityCodec: + """Build stable rank-local keys from vLLM block hashes.""" + + namespace: UMBPNamespace + tp_rank: int = 0 + pp_rank: int = 0 + pcp_rank: int = 0 + dcp_rank: int = 0 + + def key(self, block_hash: bytes, group_id: int = 0) -> str: + return self.key_for_namespace( + block_hash, + group_id, + (self.tp_rank, self.pcp_rank, self.dcp_rank, self.pp_rank), + ) + + def key_for_namespace( + self, + block_hash: bytes, + group_id: int, + namespace: tuple[int, int, int, int], + ) -> str: + if not isinstance(block_hash, bytes) or not block_hash: + raise ValueError("block_hash must be non-empty bytes") + if group_id < 0: + raise ValueError("group_id must be non-negative") + tp_rank, pcp_rank, dcp_rank, pp_rank = namespace + return ( + f"umbp:vllm:v1:{self.namespace.value}:" + f"tp{tp_rank}:pcp{pcp_rank}:dcp{dcp_rank}:" + f"pp{pp_rank}:g{group_id}:" + f"{block_hash.hex()}" + ) + + def keys_for_topology( + self, + block_hash: bytes, + topology: RankTopology, + group_ids: Sequence[int], + ) -> tuple[str, ...]: + return tuple( + self.key_for_namespace(block_hash, group_id, namespace) + for namespace in topology.all_namespaces() + for group_id in group_ids + ) + + +@dataclass(frozen=True) +class KVRegion: + """One deterministic region in a rank-local KV object.""" + + layer_name: str + group_id: int + block_stride: int + block_bytes: int + object_offset: int + block_size: int = 1 + + def __post_init__(self) -> None: + if self.block_stride <= 0 or self.block_bytes <= 0: + raise ValueError("KV region sizes must be positive") + if self.block_size <= 0: + raise ValueError("KV region block_size must be positive") + if self.block_bytes > self.block_stride: + raise ValueError( + f"KV region bytes ({self.block_bytes}) exceed its stride " + f"({self.block_stride}) for {self.layer_name!r}" + ) + + +@dataclass(frozen=True) +class KVLayoutDescriptor: + """Backend-neutral description of one rank-local KV object layout.""" + + regions: tuple[KVRegion, ...] + topology: RankTopology + layout_format: str = "rank_local" + + def __post_init__(self) -> None: + if not self.layout_format: + raise ValueError("layout_format must not be empty") + names = [region.layer_name for region in self.regions] + if len(names) != len(set(names)): + raise ValueError("KV layout contains duplicate layer regions") + + @property + def object_size(self) -> int: + return sum(region.block_bytes for region in self.regions) + + +class KVLayoutPlanner: + """Create a deterministic all-layer object layout from KVCacheConfig.""" + + def __init__(self, regions: Sequence[KVRegion]) -> None: + self.regions = tuple(regions) + self._regions_by_group = { + group_id: tuple(r for r in self.regions if r.group_id == group_id) + for group_id in {r.group_id for r in self.regions} + } + self._base_addresses: dict[str, int] = {} + self._physical_strides: dict[str, int] = {} + + @classmethod + def from_kv_cache_config(cls, config: KVCacheConfig) -> KVLayoutPlanner: + group_by_layer = { + layer: group_id + for group_id, group in enumerate(config.kv_cache_groups) + for layer in group.layer_names + if group.enable_kv_transfer + } + regions: list[KVRegion] = [] + offset = 0 + for tensor in config.kv_cache_tensors: + for layer_name in tensor.layers: + group_id = group_by_layer.get(layer_name) + if group_id is None: + continue + spec = config.kv_cache_groups[group_id].kv_cache_spec + block_bytes = spec.page_size_bytes + regions.append( + KVRegion( + layer_name=layer_name, + group_id=group_id, + block_stride=tensor.block_stride, + block_bytes=block_bytes, + object_offset=offset, + block_size=spec.block_size, + ) + ) + offset += block_bytes + return cls(regions) + + @property + def object_size(self) -> int: + return sum(region.block_bytes for region in self.regions) + + def describe( + self, + topology: RankTopology, + layout_format: str = "rank_local", + ) -> KVLayoutDescriptor: + return KVLayoutDescriptor( + regions=self.regions, + topology=topology, + layout_format=layout_format, + ) + + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: + """Capture stable base addresses for the planned KV regions.""" + required = {region.layer_name for region in self.regions} + missing = required - kv_caches.keys() + if missing: + raise ValueError(f"missing KV cache layers: {sorted(missing)}") + addresses: dict[str, int] = {} + for layer_name in required: + cache = kv_caches[layer_name] + if cache.ndim == 0: + raise ValueError(f"KV cache {layer_name!r} must be a tensor") + block_stride = cache.stride(0) * cache.element_size() + addresses[layer_name] = cache.data_ptr() + for region in self.regions: + if region.layer_name != layer_name: + continue + if ( + block_stride > region.block_stride + or region.block_stride % block_stride != 0 + or ( + region.block_bytes > block_stride + and region.block_bytes % block_stride != 0 + ) + ): + raise ValueError( + f"KV block stride mismatch for {layer_name!r}: " + f"config={region.block_stride}, tensor={block_stride}" + ) + self._physical_strides[layer_name] = block_stride + self._base_addresses = addresses + + def plan_registered_block( + self, + key: str, + block_id: int, + *, + group_id: int | None = None, + request_id: str | None = None, + ) -> BlockTransferPlan: + if not self._base_addresses: + raise RuntimeError("KV caches must be registered before planning") + return self.plan_for_block( + key, + block_id, + self._base_addresses, + group_id=group_id, + request_id=request_id, + ) + + def plan_for_block( + self, + key: str, + block_id: int, + base_addresses: dict[str, int], + *, + group_id: int | None = None, + request_id: str | None = None, + ) -> BlockTransferPlan: + """Lower one GPU block into a deterministic scatter/gather plan.""" + return BlockTransferPlan( + key=key, + block_id=block_id, + ranges=self._ranges_for_block(block_id, base_addresses, group_id), + request_id=request_id, + ) + + def materialize(self, plan: BlockTransferPlan) -> BlockTransferPlan: + if not self._base_addresses: + raise RuntimeError("KV caches must be registered before planning") + return replace( + plan, + ranges=self._ranges_for_block( + plan.block_id, self._base_addresses, plan.group_id + ), + ) + + def _ranges_for_block( + self, + block_id: int, + base_addresses: dict[str, int], + group_id: int | None, + ) -> tuple[KVRange, ...]: + if block_id < 0: + raise ValueError("block_id must be non-negative") + regions = ( + self.regions + if group_id is None + else self._regions_by_group.get(group_id, ()) + ) + if group_id is not None and not regions: + raise ValueError(f"KV layout does not contain cache group {group_id}") + ranges: list[KVRange] = [] + object_offset = 0 + for region in regions: + try: + base_address = base_addresses[region.layer_name] + except KeyError as exc: + raise ValueError( + f"missing base address for layer {region.layer_name!r}" + ) from exc + physical_stride = self._physical_strides.get( + region.layer_name, region.block_stride + ) + if region.block_bytes <= physical_stride: + ranges.append( + KVRange( + layer_name=region.layer_name, + group_id=region.group_id, + block_id=block_id, + base_address=base_address + block_id * physical_stride, + stride=physical_stride, + length=region.block_bytes, + object_offset=( + object_offset + if group_id is not None + else region.object_offset + ), + ) + ) + object_offset += region.block_bytes + continue + physical_parts = max(1, region.block_bytes // physical_stride) + for part in range(physical_parts): + part_start = part * physical_stride + part_end = min(part_start + physical_stride, region.block_bytes) + physical_block_id = block_id * physical_parts + part + ranges.append( + KVRange( + layer_name=region.layer_name, + group_id=region.group_id, + block_id=block_id, + base_address=base_address + physical_block_id * physical_stride, + stride=physical_stride, + length=part_end - part_start, + object_offset=( + object_offset + part_start + if group_id is not None + else region.object_offset + part * physical_stride + ), + ) + ) + object_offset += region.block_bytes + return tuple(ranges) + + +class KVRange(NamedTuple): + """A source/destination range in a runtime transfer.""" + + layer_name: str + group_id: int + block_id: int + base_address: int + stride: int + length: int + object_offset: int + + +@dataclass(frozen=True) +class BlockTransferPlan: + """One complete logical block transfer.""" + + key: str + block_id: int + ranges: tuple[KVRange, ...] = () + request_id: str | None = None + group_id: int | None = None + + +@dataclass +class LoadSpec: + """Scheduler decision for one externally loaded prefix.""" + + local_tokens: int + external_tokens: int + + @property + def num_tokens_to_load(self) -> int: + return max(self.external_tokens - self.local_tokens, 0) + + +@dataclass +class RequestTracker: + """Core request state shared by every runtime mode.""" + + saved_tokens: int = 0 + block_ids: tuple[list[int], ...] = () + load_spec: LoadSpec | None = None + + def reset(self) -> None: + self.saved_tokens = 0 + self.block_ids = () + self.load_spec = None + + def update_blocks(self, block_ids: tuple[list[int], ...]) -> None: + if not self.block_ids: + self.block_ids = tuple(list(group) for group in block_ids) + return + if len(self.block_ids) != len(block_ids): + raise ValueError("KV cache group count changed for a request") + for current, new in zip(self.block_ids, block_ids, strict=True): + if new[: len(current)] == current: + current.extend(new[len(current) :]) + elif not new or current[-len(new) :] == new: + continue + else: + current.extend(new) + + def replace_blocks(self, block_ids: tuple[list[int], ...]) -> None: + """Replace stale pre-preemption tables with resumed full tables.""" + self.block_ids = tuple(list(group) for group in block_ids) + self.saved_tokens = 0 + + def mark_saved(self, token_count: int, block_size: int) -> int: + """Advance the save watermark only to a complete block boundary.""" + complete_tokens = token_count // block_size * block_size + if complete_tokens < self.saved_tokens: + return self.saved_tokens + self.saved_tokens = complete_tokens + return complete_tokens + + +class TransferJobStatus(str, Enum): + PENDING = "pending" + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + + +@dataclass +class TransferJobState: + """Per-key completion state; one failed key does not fail its siblings.""" + + plans: tuple[BlockTransferPlan, ...] + status: TransferJobStatus = TransferJobStatus.PENDING + completed_keys: set[str] = field(default_factory=set) + failed_keys: set[str] = field(default_factory=set) + error: str | None = None + + @property + def keys(self) -> Sequence[str]: + return tuple(plan.key for plan in self.plans) + + def start(self) -> None: + if self.status != TransferJobStatus.PENDING: + raise RuntimeError(f"cannot start job in state {self.status}") + self.status = TransferJobStatus.RUNNING + + def complete(self, keys: Sequence[str] = ()) -> None: + self.completed_keys.update(keys or self.keys) + self._finish_if_done() + + def fail(self, keys: Sequence[str], error: str | None = None) -> None: + self.failed_keys.update(keys) + self.error = error + self._finish_if_done() + + @property + def failed_block_ids(self) -> set[int]: + return {plan.block_id for plan in self.plans if plan.key in self.failed_keys} + + def _finish_if_done(self) -> None: + keys = set(self.keys) + if self.completed_keys | self.failed_keys >= keys: + self.status = ( + TransferJobStatus.FAILED + if self.failed_keys + else TransferJobStatus.COMPLETED + ) + + +@dataclass +class UMBPConnectorMetadata(KVConnectorMetadata): + """Scheduler-to-worker metadata for one engine step.""" + + async_load: bool = False + store_plans: list[BlockTransferPlan] = field(default_factory=list) + load_requests: dict[str, list[BlockTransferPlan]] = field(default_factory=dict) + store_requests: dict[str, list[BlockTransferPlan]] = field(default_factory=dict) + preempted_request_ids: set[str] = field(default_factory=set) + store_event: int = -1 + + +@dataclass +class StoreEventResult: + """Number of workers that finished one store event.""" + + completed_workers: int = 0 + + def merge(self, other: StoreEventResult) -> None: + self.completed_workers += other.completed_workers + + +@dataclass +class UMBPConnectorWorkerMetadata(KVConnectorWorkerMetadata): + """Worker completion metadata aggregated across ranks.""" + + store_events: dict[int, StoreEventResult] = field(default_factory=dict) + + def aggregate( + self, other: KVConnectorWorkerMetadata + ) -> UMBPConnectorWorkerMetadata: + if not isinstance(other, UMBPConnectorWorkerMetadata): + raise TypeError("cannot aggregate incompatible UMBP worker metadata") + for event, result in other.store_events.items(): + self.store_events.setdefault(event, StoreEventResult()).merge(result) + return self diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/__init__.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/__init__.py new file mode 100644 index 0000000000..7b2ce33ebf --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/__init__.py @@ -0,0 +1,13 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from .base import IUMBPRuntime, UMBPSchedulerHandle, UMBPWorkerHandle +from .factory import UMBPRuntimeConfig, UMBPRuntimeFactory + +__all__ = [ + "IUMBPRuntime", + "UMBPSchedulerHandle", + "UMBPWorkerHandle", + "UMBPRuntimeConfig", + "UMBPRuntimeFactory", +] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/base.py new file mode 100644 index 0000000000..22277492d3 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/base.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Backend-neutral runtime contract for UMBP.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Protocol + +import torch + +from ..data import ( + BlockTransferPlan, + KVLayoutDescriptor, + RankTopology, + TransferJobState, +) + + +class UMBPSchedulerHandle(Protocol): + def lookup(self, keys: Sequence[str]) -> Sequence[bool]: + """Return authoritative existence for each key.""" + + def clear(self) -> bool: + """Clear all objects visible to this scheduler handle.""" + + def close(self) -> None: + """Release scheduler-side resources.""" + + +class UMBPWorkerHandle(Protocol): + def register_buffers(self, kv_caches: dict[str, torch.Tensor]) -> None: + """Register each unique KV allocation with the runtime.""" + + def load(self, plans: Sequence[BlockTransferPlan]) -> TransferJobState: + """Start loading the requested plans.""" + + def store(self, plans: Sequence[BlockTransferPlan]) -> TransferJobState: + """Start storing the requested plans.""" + + def wait(self, job: TransferJobState) -> TransferJobState: + """Block until the job is final and its buffers are safe to reuse.""" + + def poll(self, job: TransferJobState) -> TransferJobState | None: + """Return a finished job without blocking, or None if still pending.""" + + def publish(self, job: TransferJobState) -> None: + """Accept a completed store; a runtime may defer visibility until then.""" + + def close(self) -> None: + """Release worker-side resources.""" + + +class IUMBPRuntime(Protocol): + def create_scheduler_handle( + self, + namespace: str, + topology: RankTopology, + layout: KVLayoutDescriptor, + ) -> UMBPSchedulerHandle: + """Create a scheduler-side handle.""" + + def create_worker_handle( + self, + namespace: str, + topology: RankTopology, + layout: KVLayoutDescriptor, + ) -> UMBPWorkerHandle: + """Create a worker-side handle.""" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/factory.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/factory.py new file mode 100644 index 0000000000..b62c2587fc --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/factory.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Runtime configuration and adapter factory for UMBP.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, replace +from typing import Any + +from vllm.config import VllmConfig + +from .base import IUMBPRuntime + + +@dataclass(frozen=True) +class UMBPRuntimeConfig: + """Shared connector options passed to the selected runtime adapter.""" + + mode: str + options: dict[str, Any] + rank_count: int = 1 + dp_index: int = 0 + + @classmethod + def from_vllm(cls, vllm_config: VllmConfig) -> UMBPRuntimeConfig: + transfer_config = vllm_config.kv_transfer_config + if transfer_config is None: + raise ValueError("UMBP requires kv_transfer_config") + options = dict(transfer_config.kv_connector_extra_config) + mode = options.get("mode", "embedded") + if mode not in {"embedded"}: + raise ValueError(f"unknown UMBP mode: {mode!r}") + namespace = options.get("key_namespace", "auto") + if namespace != "auto" and (not isinstance(namespace, str) or not namespace): + raise ValueError("key_namespace must be a non-empty string or 'auto'") + return cls( + mode=mode, + options=options, + dp_index=int( + getattr(vllm_config.parallel_config, "data_parallel_index", 0) or 0 + ), + ) + + def resolve_for_rank_count(self, rank_count: int) -> UMBPRuntimeConfig: + """Pass topology to the adapter without interpreting its capacity options.""" + if rank_count <= 0: + raise ValueError("rank_count must be positive") + return replace(self, rank_count=rank_count) + + +RuntimeBuilder = Callable[[UMBPRuntimeConfig], IUMBPRuntime] + + +class UMBPRuntimeFactory: + """Registry used by UMBP adapters without importing MORI in vLLM.""" + + _builders: dict[str, RuntimeBuilder] = {} + + @classmethod + def register(cls, mode: str, builder: RuntimeBuilder) -> None: + if mode not in {"embedded"}: + raise ValueError(f"unknown UMBP mode: {mode!r}") + if mode in cls._builders: + raise ValueError(f"UMBP runtime mode {mode!r} is already registered") + cls._builders[mode] = builder + + @classmethod + def build(cls, config: UMBPRuntimeConfig) -> IUMBPRuntime: + try: + builder = cls._builders[config.mode] + except KeyError as exc: + raise RuntimeError( + f"no UMBP runtime adapter is registered for mode {config.mode!r}" + ) from exc + return builder(config) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/scheduler.py new file mode 100644 index 0000000000..8ca17a99c1 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/scheduler.py @@ -0,0 +1,518 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared scheduler-side UMBP connector behavior.""" + +from __future__ import annotations + +import logging +from collections.abc import Sequence +from dataclasses import dataclass, field +from functools import lru_cache +from typing import Any + +from vllm.config import VllmConfig +from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata +from vllm.logger import init_logger +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID +from vllm.v1.core.block_pool import BlockPool +from vllm.v1.core.kv_cache_manager import KVCacheBlocks +from vllm.v1.core.kv_cache_utils import ( + resolve_dcp_kv_block_size, + resolve_kv_cache_block_sizes, +) +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.kv_cache_interface import ( + KVCacheConfig, +) +from vllm.v1.outputs import KVConnectorOutput +from vllm.v1.request import Request + +from .data import ( + BlockIdentityCodec, + BlockTransferPlan, + LoadSpec, + RankTopology, + RequestTracker, + StoreEventResult, + UMBPConnectorMetadata, + UMBPConnectorWorkerMetadata, +) +from .runtime import UMBPSchedulerHandle + +logger = init_logger(__name__) + + +@lru_cache(maxsize=8192) +def _rank_keys( + codec: BlockIdentityCodec, + topology: RankTopology, + block_hash: bytes, + group_id: int, +) -> tuple[str, ...]: + """Reuse immutable object identities, not runtime residency.""" + return codec.keys_for_topology(block_hash, topology, (group_id,)) + + +@dataclass +class _StoreEvent: + plans: tuple[BlockTransferPlan, ...] + result: StoreEventResult = field(default_factory=StoreEventResult) + + +class UMBPStoreConnectorScheduler: + """Own vLLM prefix matching while the runtime owns lookup transport.""" + + def __init__( + self, + vllm_config: VllmConfig, + kv_cache_config: KVCacheConfig, + runtime: UMBPSchedulerHandle, + codec: BlockIdentityCodec, + topology: RankTopology | None = None, + ) -> None: + self.block_size, resolved_hash_block_size = resolve_kv_cache_block_sizes( + kv_cache_config, vllm_config + ) + self.kv_cache_config = kv_cache_config + dcp_size = vllm_config.parallel_config.decode_context_parallel_size + self.group_block_sizes = { + group_id: resolve_dcp_kv_block_size(group.kv_cache_spec, dcp_size) + for group_id, group in enumerate(kv_cache_config.kv_cache_groups) + } + self.runtime = runtime + self.codec = codec + self.topology = topology or RankTopology() + transfer_config = vllm_config.kv_transfer_config + assert transfer_config is not None + extra = transfer_config.kv_connector_extra_config + self.load_async = bool(extra.get("load_async", True)) + self.enable_lookup = bool(extra.get("enable_lookup", True)) + self.save_decode_cache = bool(extra.get("save_decode_cache", False)) + # A kv_consumer only restores from the pool, as with other store + # connectors; it never offloads its own KV. + self.store_enabled = getattr(transfer_config, "kv_role", None) != "kv_consumer" + # Request.block_hashes are computed at this size; keys index into them. + self.hash_block_size = resolved_hash_block_size + speculative_config = getattr(vllm_config, "speculative_config", None) + self.use_eagle = bool( + speculative_config is not None and speculative_config.use_eagle_block_drop() + ) + self._pending_loads: dict[str, list[BlockTransferPlan]] = {} + self._load_specs: dict[str, LoadSpec] = {} + self._lookup_units = { + group_id: self.group_block_sizes[group_id] + for group_id in kv_cache_config.prefix_cacheable_group_ids + } + self._requests: dict[str, Request] = {} + self._request_trackers: dict[str, RequestTracker] = {} + self._gpu_block_pool: BlockPool | None = None + self._store_events: dict[int, _StoreEvent] = {} + self._num_workers = getattr(vllm_config.parallel_config, "world_size", 1) + self._store_event_counter = 0 + + def bind_gpu_block_pool(self, gpu_block_pool: BlockPool) -> None: + self._gpu_block_pool = gpu_block_pool + + def get_num_new_matched_tokens( + self, request: Request, num_computed_tokens: int + ) -> tuple[int | None, bool]: + if not self.enable_lookup: + return 0, False + hashes = request.block_hashes + if not hashes or request.num_tokens < self.block_size: + return 0, False + if num_computed_tokens % self.block_size != 0: + return 0, False + group_ids = self.kv_cache_config.prefix_cacheable_group_ids + # vLLM must execute at least the final prompt token. Returning a hit + # for the entire prompt leaves the scheduler with no new token to run. + max_external_token = request.num_tokens - 1 + if max_external_token <= num_computed_tokens: + self._load_specs.pop(request.request_id, None) + return 0, False + logical_objects_by_group, keys = self._lookup_objects( + request, num_computed_tokens, hashes + ) + try: + hits = list(self.runtime.lookup(keys)) + except Exception as exc: + logger.debug("UMBP lookup failed request=%s: %s", request.request_id, exc) + return 0, False + if len(hits) != len(keys): + logger.debug( + "UMBP lookup returned an invalid result length request=%s", + request.request_id, + ) + return 0, False + per_rank = self.topology.rank_count + offset = 0 + logical_hits_by_group: dict[int, list[bool]] = {} + for group_id in group_ids: + group_size = len(logical_objects_by_group[group_id]) * per_rank + group_hits = hits[offset : offset + group_size] + offset += group_size + logical_hits_by_group[group_id] = ( + group_hits + if per_rank == 1 + else [ + all(group_hits[index : index + per_rank]) + for index in range(0, len(group_hits), per_rank) + ] + ) + group_id = group_ids[0] + matched_blocks = 0 + for hit in logical_hits_by_group[group_id]: + if not hit: + break + matched_blocks += 1 + if self.use_eagle: + # As core's coordinator: the drafter recomputes the last block. + matched_blocks = max(0, matched_blocks - 1) + need_to_load = matched_blocks * self._lookup_units[group_id] + if logger.isEnabledFor(logging.DEBUG): + logger.debug( + "UMBP lookup result request=%s keys=%s group_hits=%s load_tokens=%d", + request.request_id, + keys, + logical_hits_by_group, + need_to_load, + ) + matched_tokens = num_computed_tokens + need_to_load + if need_to_load <= 0: + return 0, False + self._load_specs[request.request_id] = LoadSpec( + local_tokens=num_computed_tokens, + external_tokens=matched_tokens, + ) + return need_to_load, self.load_async + + def _lookup_objects( + self, request: Request, local_tokens: int, hashes: Sequence[bytes] + ) -> tuple[dict[int, list[tuple[bytes, tuple[str, ...]]]], list[str]]: + objects: dict[int, list[tuple[bytes, tuple[str, ...]]]] = {} + keys: list[str] = [] + for group_id, unit in self._lookup_units.items(): + group_objects = [] + for token_end in range(local_tokens + unit, request.num_tokens, unit): + block_hash = self._object_hash_at_token_end(hashes, token_end) + rank_keys = _rank_keys(self.codec, self.topology, block_hash, group_id) + group_objects.append((block_hash, rank_keys)) + keys.extend(rank_keys) + objects[group_id] = group_objects + return objects, keys + + def update_state_after_alloc( + self, + request: Request, + blocks: KVCacheBlocks, + num_external_tokens: int, + ) -> None: + self._requests[request.request_id] = request + block_groups = blocks.get_block_ids( + group_ids=self.kv_cache_config.prefix_cacheable_group_ids + ) + if not block_groups: + return + tracker = self._tracker_for_request(request.request_id) + # Core passes the request's whole table on every admission. A request + # whose load failed is recomputed into newly allocated blocks. + tracker.replace_blocks(tuple(list(group) for group in block_groups)) + if num_external_tokens <= 0: + self._pending_loads.pop(request.request_id, None) + self._load_specs.pop(request.request_id, None) + return + spec = self._load_specs.get(request.request_id) + if spec is not None: + tracker.load_spec = spec + # A request admitted without its load would wait for it forever. + assert num_external_tokens == spec.num_tokens_to_load, ( + f"UMBP matched {spec.num_tokens_to_load} external tokens for " + f"request {request.request_id} but {num_external_tokens} were " + "allocated" + ) + plans = self._load_plans_for_external_tokens( + request, + tracker, + block_groups, + num_external_tokens, + ) + self._pending_loads[request.request_id] = plans + + def build_connector_meta( + self, scheduler_output: SchedulerOutput + ) -> KVConnectorMetadata: + meta = UMBPConnectorMetadata(async_load=self.load_async) + for request_id in scheduler_output.finished_req_ids: + self._pending_loads.pop(request_id, None) + self._load_specs.pop(request_id, None) + self._request_trackers.pop(request_id, None) + self._requests.pop(request_id, None) + for request in scheduler_output.scheduled_new_reqs: + load_plans = self._pending_loads.pop(request.req_id, []) + self._load_specs.pop(request.req_id, None) + if load_plans: + meta.load_requests[request.req_id] = load_plans + else: + tracker = self._tracker_for_request(request.req_id) + total_tokens = ( + request.num_computed_tokens + + scheduler_output.num_scheduled_tokens[request.req_id] + ) + tracked_request = self._requests.get(request.req_id) + selected_block_ids = tuple( + request.block_ids[group_id] + for group_id in (self.kv_cache_config.prefix_cacheable_group_ids) + ) + store_plans = ( + [] + if not self.store_enabled or tracked_request is None + else self._store_plans( + tracked_request, + tracker, + total_tokens, + block_ids_override=selected_block_ids, + ) + ) + logger.debug( + "UMBP store planning request=%s tokens=%d plans=%d keys=%s", + request.req_id, + total_tokens, + len(store_plans), + tuple(plan.key for plan in store_plans), + ) + meta.store_plans.extend(store_plans) + if store_plans: + meta.store_requests[request.req_id] = store_plans + + block_state = getattr(scheduler_output, "kv_connector_block_state", None) + cached_reqs = scheduler_output.scheduled_cached_reqs + resumed_ids: set[str] = set(getattr(cached_reqs, "resumed_req_ids", ())) + for request_id in cached_reqs.req_ids: + load_plans = self._pending_loads.pop(request_id, []) + self._load_specs.pop(request_id, None) + if load_plans: + meta.load_requests[request_id] = load_plans + elif self.store_enabled: + cached_request = self._requests.get(request_id) + cached_tracker = self._request_trackers.get(request_id) + if cached_request is not None and cached_tracker is not None: + request_index = cached_reqs.req_ids.index(request_id) + num_computed_tokens = cached_reqs.num_computed_tokens[request_index] + # num_tokens counts sampled tokens, so a decode step is still + # short of it; only the prompt bounds the prefill. + is_prefill = num_computed_tokens < cached_request.num_prompt_tokens + if not is_prefill and not self.save_decode_cache: + continue + new_block_ids = cached_reqs.new_block_ids[request_index] + if new_block_ids: + selected = tuple( + new_block_ids[group_id] + for group_id in ( + self.kv_cache_config.prefix_cacheable_group_ids + ) + ) + if request_id in resumed_ids: + cached_tracker.replace_blocks(selected) + else: + cached_tracker.update_blocks(selected) + # The appended table keeps blocks core has since freed, such + # as sliding-window blocks out of the window, which a store + # would read; core's current table does not. + current = ( + block_state.get_block_ids(request_id) + if block_state is not None + else None + ) + if current is not None: + cached_tracker.block_ids = tuple( + list(current[group_id]) + for group_id in ( + self.kv_cache_config.prefix_cacheable_group_ids + ) + ) + total_tokens = ( + num_computed_tokens + + scheduler_output.num_scheduled_tokens[request_id] + ) + if not self.save_decode_cache: + total_tokens = min( + total_tokens, cached_request.num_prompt_tokens + ) + store_plans = self._store_plans( + cached_request, + cached_tracker, + total_tokens, + block_ids_override=cached_tracker.block_ids, + ) + meta.store_plans.extend(store_plans) + if store_plans: + meta.store_requests[request_id] = store_plans + + # Async loads are admitted without scheduling model tokens, so those + # requests are absent from both scheduled request collections above. + # Forward their plans explicitly so workers can start the transfer and + # eventually move the request out of WAITING_FOR_REMOTE_KVS. + for request_id, load_plans in list(self._pending_loads.items()): + self._pending_loads.pop(request_id, None) + self._load_specs.pop(request_id, None) + if load_plans: + meta.load_requests[request_id] = load_plans + + for request_id in scheduler_output.preempted_req_ids or set(): + meta.preempted_request_ids.add(request_id) + self._load_specs.pop(request_id, None) + if preempted_tracker := self._request_trackers.get(request_id): + preempted_tracker.reset() + + event_plans = tuple(meta.store_plans) + if event_plans: + event = self._store_event_counter + self._store_event_counter += 1 + meta.store_event = event + self._store_events[event] = _StoreEvent(event_plans) + if self._gpu_block_pool is not None: + self._gpu_block_pool.touch( + [ + self._gpu_block_pool.blocks[block_id] + for block_id in dict.fromkeys( + plan.block_id for plan in event_plans + ) + ] + ) + return meta + + def _tracker_for_request(self, request_id: str) -> RequestTracker: + return self._request_trackers.setdefault(request_id, RequestTracker()) + + def _load_plans_for_external_tokens( + self, + request: Any, + tracker: RequestTracker, + block_groups: tuple[list[int], ...], + num_external_tokens: int, + ) -> list[BlockTransferPlan]: + """Plan the logical objects selected by the external hit window.""" + group_ids = self.kv_cache_config.prefix_cacheable_group_ids + spec = tracker.load_spec + local_tokens = spec.local_tokens if spec is not None else 0 + end_tokens = local_tokens + num_external_tokens + plans: list[BlockTransferPlan] = [] + for group_index, group_id in enumerate(group_ids): + group_block_size = self.group_block_sizes[group_id] + start_block = local_tokens // group_block_size + num_full_blocks = end_tokens // group_block_size + for block_index in range(start_block, num_full_blocks): + token_end = (block_index + 1) * group_block_size + block_hash = self._object_hash_at_token_end( + request.block_hashes, token_end + ) + plans.append( + BlockTransferPlan( + key=self.codec.key(block_hash, group_id=group_id), + block_id=block_groups[group_index][block_index], + request_id=request.request_id, + group_id=group_id, + ) + ) + return plans + + def _object_hash_at_token_end( + self, hashes: Sequence[bytes], token_end: int + ) -> bytes: + hash_index = token_end // self.hash_block_size - 1 + if hash_index < 0 or hash_index >= len(hashes): + raise ValueError( + f"request has {len(hashes)} hashes, needs index {hash_index}" + ) + return hashes[hash_index] + + def _store_plans( + self, + request: Any, + tracker: RequestTracker, + token_count: int, + block_ids_override: tuple[list[int], ...] | None = None, + ) -> list[BlockTransferPlan]: + """Describe full blocks produced by a scheduled prefill.""" + request_id = getattr(request, "request_id", None) or request.req_id + group_ids = self.kv_cache_config.prefix_cacheable_group_ids + # Scheduled speculative tokens are not yet verified and have no block + # hash; a block is storable only once its hash exists. + token_count = min(token_count, len(request.block_hashes) * self.hash_block_size) + previous_saved_tokens = tracker.saved_tokens + # Hash-aligned rather than block-aligned: a group with smaller blocks + # than the largest must not wait for that group's block to fill. + save_to = tracker.mark_saved(token_count, self.hash_block_size) + plans: list[BlockTransferPlan] = [] + block_groups = ( + block_ids_override + if block_ids_override is not None + else tuple(request.block_ids[group_id] for group_id in group_ids) + ) + block_ranges = { + group_id: range( + previous_saved_tokens // self.group_block_sizes[group_id], + min(save_to // self.group_block_sizes[group_id], len(block_ids)), + ) + for group_id, block_ids in zip(group_ids, block_groups) + } + for group_id, block_ids in zip(group_ids, block_groups): + if group_id not in block_ranges: + continue + group_block_size = self.group_block_sizes[group_id] + for index in block_ranges[group_id]: + if block_ids[index] == NULL_BLOCK_ID: + continue + token_end = (index + 1) * group_block_size + block_hash = self._object_hash_at_token_end( + request.block_hashes, token_end + ) + plans.append( + BlockTransferPlan( + key=self.codec.key(block_hash, group_id=group_id), + block_id=block_ids[index], + request_id=request_id, + group_id=group_id, + ) + ) + return plans + + def request_finished( + self, request: Request, block_ids: tuple[list[int], ...] + ) -> tuple[bool, dict[str, Any] | None]: + self._pending_loads.pop(request.request_id, None) + self._load_specs.pop(request.request_id, None) + self._requests.pop(request.request_id, None) + return False, None + + def update_connector_output(self, output: KVConnectorOutput) -> None: + metadata = output.kv_connector_worker_meta + if not isinstance(metadata, UMBPConnectorWorkerMetadata): + return + pool = self._gpu_block_pool + for event_id, result in metadata.store_events.items(): + event = self._store_events.get(event_id) + if event is None: + continue + event.result.merge(result) + if event.result.completed_workers < self._num_workers: + continue + del self._store_events[event_id] + if pool is not None: + block_ids = dict.fromkeys(plan.block_id for plan in event.plans) + pool.free_blocks( + pool.blocks[block_id] for block_id in reversed(block_ids) + ) + + def has_pending_push_work(self) -> bool: + return bool(self._store_events) + + def reset_store(self) -> bool: + if self._store_events: + return False + self._load_specs.clear() + self._pending_loads.clear() + return self.runtime.clear() + + def close(self) -> None: + self.runtime.close() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/worker.py new file mode 100644 index 0000000000..34101f294c --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/worker.py @@ -0,0 +1,240 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared worker-side UMBP connector behavior.""" + +from __future__ import annotations + +from dataclasses import dataclass, replace + +import torch + +from vllm.forward_context import ForwardContext +from vllm.logger import init_logger + +from .data import ( + BlockIdentityCodec, + BlockTransferPlan, + KVLayoutPlanner, + StoreEventResult, + TransferJobState, + TransferJobStatus, + UMBPConnectorMetadata, + UMBPConnectorWorkerMetadata, +) +from .runtime import UMBPWorkerHandle + +logger = init_logger(__name__) + + +@dataclass +class _StoreBatch: + event_id: int + jobs: dict[str, TransferJobState] + + +class UMBPStoreConnectorWorker: + """Translate shared metadata into runtime load/store jobs.""" + + def __init__( + self, + runtime: UMBPWorkerHandle, + layout: KVLayoutPlanner | None = None, + *, + codec: BlockIdentityCodec | None = None, + ) -> None: + self.runtime = runtime + self.layout = layout + self.codec = codec + self._load_jobs: dict[str, dict[str | None, TransferJobState]] = {} + self._store_jobs: dict[str, TransferJobState] = {} + self._pending_stores: list[_StoreBatch] = [] + self._worker_meta = UMBPConnectorWorkerMetadata() + self._finished_recving: set[str] = set() + self._load_error_block_ids: set[int] = set() + self._report_load_completions = False + self._active_store_event = -1 + + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: + if self.layout is not None and self.layout.regions: + self.layout.register_kv_caches(kv_caches) + self.runtime.register_buffers(kv_caches) + + def _materialize_plans( + self, plans: list[BlockTransferPlan] + ) -> list[BlockTransferPlan]: + plans = [self._localize_plan(plan) for plan in plans] + if self.layout is None or not self.layout.regions: + return plans + materialized: list[BlockTransferPlan] = [] + for plan in plans: + if plan.ranges: + materialized.append(plan) + continue + materialized.append(self.layout.materialize(plan)) + return materialized + + def _localize_plan(self, plan: BlockTransferPlan) -> BlockTransferPlan: + if self.codec is None or plan.group_id is None: + return plan + key = self._localize_key(plan.key, plan.group_id) + if key == plan.key: + return plan + return replace(plan, key=key) + + def _localize_key(self, key: str, group_id: int) -> str: + if self.codec is None: + return key + try: + block_hash = bytes.fromhex(key.rsplit(":", 1)[1]) + except (IndexError, ValueError): + return key + return self.codec.key(block_hash, group_id) + + def start_load_kv( + self, forward_context: ForwardContext, metadata: UMBPConnectorMetadata + ) -> None: + del forward_context + self._report_load_completions = metadata.async_load + for request_id, plans in metadata.load_requests.items(): + if plans: + self._submit_loads(request_id, plans) + + def _submit_loads(self, request_id: str, plans: list[BlockTransferPlan]) -> None: + materialized = self._materialize_plans(plans) + logger.debug( + "UMBP load submitted request=%s plans=%d ranges=%d bytes=%d", + request_id, + len(materialized), + sum(len(plan.ranges) for plan in materialized), + sum(item.length for plan in materialized for item in plan.ranges), + ) + self._load_jobs[request_id] = {None: self.runtime.load(materialized)} + + def _finish_job( + self, request_id: str, result: TransferJobState, *, is_load: bool + ) -> bool: + operation = "load" if is_load else "store" + succeeded = result.status == TransferJobStatus.COMPLETED + if not succeeded: + logger.debug( + "UMBP %s failed request=%s status=%s error=%s", + operation, + request_id, + result.status.value, + result.error, + ) + if is_load: + self._load_error_block_ids.update(result.failed_block_ids) + elif succeeded: + self.runtime.publish(result) + return succeeded + + def wait_for_layer_load(self, layer_name: str) -> None: + if self._report_load_completions: + # Asynchronous loads fill blocks of requests that do not run until + # get_finished reports them, so no forward pass reads those blocks. + return + self._drain_load_jobs(wait=True) + + def wait_for_save(self) -> None: + if self._store_jobs or self._active_store_event >= 0: + self._pending_stores.append( + _StoreBatch(self._active_store_event, self._store_jobs) + ) + self._store_jobs = {} + self._active_store_event = -1 + self._drain_store_jobs(wait=False) + + def _drain_store_jobs(self, *, wait: bool) -> None: + for batch in list(self._pending_stores): + for request_id, job in list(batch.jobs.items()): + result = self.runtime.wait(job) if wait else self.runtime.poll(job) + if result is None: + continue + # A partly failed job is not published, including its stored keys. + self._finish_job(request_id, result, is_load=False) + del batch.jobs[request_id] + if not batch.jobs: + if batch.event_id >= 0: + self._worker_meta.store_events[batch.event_id] = StoreEventResult( + completed_workers=1 + ) + self._pending_stores.remove(batch) + + def _submit_store_plans( + self, plans: list[BlockTransferPlan], request_id: str + ) -> None: + if not plans: + return + materialized = self._materialize_plans(plans) + self._store_jobs[request_id] = self.runtime.store(materialized) + + def enqueue_stores(self, metadata: UMBPConnectorMetadata) -> None: + self._active_store_event = metadata.store_event + store_requests = { + request_id: list(plans) + for request_id, plans in metadata.store_requests.items() + } + grouped = {plan for plans in store_requests.values() for plan in plans} + for plan in metadata.store_plans: + if plan not in grouped: + store_requests.setdefault( + plan.request_id or "__umbp_batch__", [] + ).append(plan) + grouped.add(plan) + for request_id, plans in store_requests.items(): + self._submit_store_plans(plans, request_id) + + def handle_preemptions(self, metadata: UMBPConnectorMetadata) -> None: + """Finish request-local jobs before vLLM reuses their GPU blocks.""" + preempted = metadata.preempted_request_ids + if not preempted: + return + for request_id in preempted: + for job in self._load_jobs.pop(request_id, {}).values(): + self.runtime.wait(job) + for jobs in [self._store_jobs, *(batch.jobs for batch in self._pending_stores)]: + for request_id, job in jobs.items(): + if request_id in preempted or any( + plan.request_id in preempted for plan in job.plans + ): + jobs[request_id] = self.runtime.wait(job) + + def get_finished( + self, finished_req_ids: set[str] + ) -> tuple[set[str] | None, set[str] | None]: + del finished_req_ids + self._drain_load_jobs(wait=False) + self._drain_store_jobs(wait=False) + recving = set(self._finished_recving) + self._finished_recving.difference_update(recving) + return None, recving or None + + def _drain_load_jobs(self, *, wait: bool) -> None: + for request_id, jobs in list(self._load_jobs.items()): + for layer, job in list(jobs.items()): + result = self.runtime.wait(job) if wait else self.runtime.poll(job) + if result is None: + continue + del jobs[layer] + self._finish_job(request_id, result, is_load=True) + if not jobs: + del self._load_jobs[request_id] + if self._report_load_completions: + self._finished_recving.add(request_id) + + def build_connector_worker_meta(self) -> UMBPConnectorWorkerMetadata: + result = self._worker_meta + self._worker_meta = UMBPConnectorWorkerMetadata() + return result + + def get_block_ids_with_load_errors(self) -> set[int]: + failed = self._load_error_block_ids + self._load_error_block_ids = set() + return failed + + def close(self) -> None: + self.wait_for_save() + self._drain_store_jobs(wait=True) + self._drain_load_jobs(wait=True) + self.runtime.close() diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index fea0c4e3c7..e4b0222b0b 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -50,6 +50,7 @@ from vllm.distributed.device_communicators.base_device_communicator import ( ) from vllm.distributed.utils import ( StatelessProcessGroup, + cpu_all_reduce_, create_tcp_store, get_cached_tcp_store_client, ) @@ -1376,6 +1377,10 @@ class GroupCoordinator: """ torch.distributed.barrier(group=self.cpu_group) + def cpu_all_reduce_(self, tensor: torch.Tensor) -> torch.Tensor: + """In-place SUM all-reduce of a small CPU tensor over `cpu_group`.""" + return cpu_all_reduce_(tensor, self.cpu_group) + def send(self, tensor: torch.Tensor, dst: int | None = None) -> None: """Sends a tensor to the destination rank in a blocking way.""" """NOTE: `dst` is the local rank of the destination rank.""" diff --git a/vllm/distributed/utils.py b/vllm/distributed/utils.py index 3618361f2a..010d517d46 100644 --- a/vllm/distributed/utils.py +++ b/vllm/distributed/utils.py @@ -55,6 +55,70 @@ def ensure_divisibility(numerator, denominator): ) +# P2P tags [base, base + log2(group size) + 1] are reserved for cpu_all_reduce_. +_CPU_ALL_REDUCE_TAG_BASE = 27100 + + +def cpu_all_reduce_(tensor: torch.Tensor, group: ProcessGroup) -> torch.Tensor: + """Low-latency in-place SUM all-reduce of a small CPU tensor over a gloo group. + + Recursive doubling within power-of-two blocks, then one concurrent exchange + between blocks. Results are bitwise identical across ranks. Works with + stateless (unregistered) process groups. + """ + world_size = group.size() + if world_size == 1: + return tensor + block_size = world_size & -world_size + num_blocks = world_size // block_size + if tensor.device.type != "cpu" or not tensor.is_contiguous(): + torch.distributed.all_reduce(tensor, group=group) + return tensor + + rank = group.rank() + tag = _CPU_ALL_REDUCE_TAG_BASE + + # Recursive doubling within the block (peers stay in the aligned block). + if block_size > 1: + # This branch applies only to even world sizes. + scratch = torch.empty_like(tensor) + mask = 1 + while mask < block_size: + peer = rank ^ mask + # Wait for both ops before touching either buffer. + recv_work = group.recv([scratch], peer, tag) + send_work = group.send([tensor], peer, tag) + recv_work.wait() + send_work.wait() + tensor.add_(scratch) + mask <<= 1 + tag += 1 + + # Exchange block sums with the same lane in every other block at once. + if num_blocks > 1: + # This branch applies only when world_size is not a power of two. + lane, my_block = rank % block_size, rank // block_size + block_sums = [ + tensor if b == my_block else torch.empty_like(tensor) + for b in range(num_blocks) + ] + works = [] + for b, buf in enumerate(block_sums): + if b != my_block: + works.append(group.recv([buf], b * block_size + lane, tag)) + for b in range(num_blocks): + if b != my_block: + works.append(group.send([tensor], b * block_size + lane, tag)) + for work in works: + work.wait() + # Sum in block order so every rank computes the identical result. + total = block_sums[0] + for buf in block_sums[1:]: + total.add_(buf) + tensor.copy_(total) + return tensor + + def divide(numerator, denominator): """Ensure that numerator is divisible by the denominator and return the division value.""" diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py index 15f8f8b744..492f8a2ce1 100644 --- a/vllm/engine/protocol.py +++ b/vllm/engine/protocol.py @@ -60,6 +60,10 @@ class EngineClient(ABC): @abstractmethod def errored(self) -> bool: ... + @property + def should_keep_api_server_alive(self) -> bool: + return False + @property @abstractmethod def dead_error(self) -> BaseException: ... @@ -262,6 +266,12 @@ class EngineClient(ABC): """Scale the engine.""" raise NotImplementedError + async def get_external_elastic_ep_status( + self, + ) -> dict[str, str | int | None] | None: + """Return shared external Elastic EP operation status, if enabled.""" + return None + async def collective_rpc( self, method: str, diff --git a/vllm/entrypoints/cohere/protocol.py b/vllm/entrypoints/cohere/protocol.py index 45e0d310dc..577462f592 100644 --- a/vllm/entrypoints/cohere/protocol.py +++ b/vllm/entrypoints/cohere/protocol.py @@ -44,7 +44,7 @@ from cohere.types import ( ToolV2, UserChatMessageV2, ) -from pydantic import BaseModel, Field, field_validator +from pydantic import BaseModel, Field, field_validator, model_validator import vllm.envs as envs from vllm.exceptions import VLLMValidationError @@ -211,6 +211,16 @@ class CohereChatV2Request(BaseModel): ) return v + @model_validator(mode="after") + def _validate_tool_choice(self) -> CohereChatV2Request: + if self.tool_choice == "REQUIRED" and not self.tools: + raise VLLMValidationError( + 'tool_choice="REQUIRED" requires at least one tool in tools.', + parameter="tool_choice", + value=self.tool_choice, + ) + return self + @field_validator("messages", mode="before") @classmethod def _normalize_message_roles(cls, v: Any) -> Any: diff --git a/vllm/entrypoints/generate/base/serving.py b/vllm/entrypoints/generate/base/serving.py index a0eda6d05f..632f397711 100644 --- a/vllm/entrypoints/generate/base/serving.py +++ b/vllm/entrypoints/generate/base/serving.py @@ -356,6 +356,32 @@ def decode_token_ids( return out +def chat_top_logprobs_limit( + top_logprobs: int | None, return_all: bool = False +) -> int | None: + """How many of a position's logprob candidates `/v1/chat/completions` + returns, in the engine's order (sampled token first): the first + ``top_logprobs``, none when it is unset, and all of them (``None``) for + ``top_logprobs=-1`` or ``logprob_token_ids`` (``return_all``). Shared with + chat derender so the two cannot drift.""" + if return_all or top_logprobs == -1: + return None + return top_logprobs or 0 + + +def completion_top_logprobs_limit( + logprobs: int | None, return_all: bool = False +) -> int | None: + """How many of a position's logprob candidates `/v1/completions` returns: + the first ``logprobs + 1`` (the OpenAI completions convention), and all of + them (``None``) for ``logprobs=-1``, ``logprob_token_ids`` (``return_all``) + or when ``logprobs`` is unset. Shared with completions derender so the two + cannot drift.""" + if return_all or logprobs is None or logprobs == -1: + return None + return max(logprobs + 1, 0) + + def clamp_prompt_logprobs( prompt_logprobs: PromptLogprobs | None, ) -> PromptLogprobs | None: diff --git a/vllm/entrypoints/launchers/launcher.py b/vllm/entrypoints/launchers/launcher.py index 44161439b6..1db294cfc6 100644 --- a/vllm/entrypoints/launchers/launcher.py +++ b/vllm/entrypoints/launchers/launcher.py @@ -328,7 +328,10 @@ def terminate_if_errored(server: uvicorn.Server, engine: EngineClient): for this request. """ engine_errored = engine.errored and not engine.is_running - if not envs.VLLM_KEEP_ALIVE_ON_ENGINE_DEATH and engine_errored: + keep_alive = ( + envs.VLLM_KEEP_ALIVE_ON_ENGINE_DEATH or engine.should_keep_api_server_alive + ) + if not keep_alive and engine_errored: server.should_exit = True diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 007c961ca8..9ddc20ed30 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio +import itertools import time from collections.abc import AsyncGenerator, AsyncIterator from collections.abc import Sequence as GenericSequence @@ -22,6 +23,7 @@ from vllm.entrypoints.generate.base.serving import ( GenerateBaseServing, build_per_request_timing_metrics, build_spec_decoding_metrics, + chat_top_logprobs_limit, clamp_prompt_logprobs, format_token_id_placeholder, ) @@ -1231,10 +1233,9 @@ class OpenAIServingChat(GenerateBaseServing): logprob=max(p[1].logprob, -9999.0), bytes=list(token.encode("utf-8", errors="replace")), ) - for i, p in enumerate(logprobs.items()) - if return_all - or top_logprobs == -1 - or (top_logprobs is not None and i < top_logprobs) + for p in itertools.islice( + logprobs.items(), chat_top_logprobs_limit(top_logprobs, return_all) + ) ] def _create_chat_logprobs( diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index e5ba20c52f..0306909237 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio +import itertools import time from collections.abc import AsyncGenerator, AsyncIterator from collections.abc import Sequence as GenericSequence @@ -19,6 +20,7 @@ from vllm.entrypoints.generate.base.serving import ( build_per_request_timing_metrics, build_spec_decoding_metrics, clamp_prompt_logprobs, + completion_top_logprobs_limit, format_token_id_placeholder, ) from vllm.entrypoints.openai.completion.protocol import ( @@ -721,10 +723,12 @@ class OpenAIServingCompletion(GenerateBaseServing): tokenizer, return_as_token_id=should_return_as_token_id, ): max(top_lp[1].logprob, -9999.0) - for i, top_lp in enumerate(step_top_logprobs.items()) - if logprob_token_ids - or num_output_top_logprobs == -1 - or num_output_top_logprobs >= i + for top_lp in itertools.islice( + step_top_logprobs.items(), + completion_top_logprobs_limit( + num_output_top_logprobs, bool(logprob_token_ids) + ), + ) } ) diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py index f09ee6fad5..02a6810866 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -424,14 +424,12 @@ class ServingTokens(GenerateBaseServing): logprobs = self._create_text_logprobs( token_ids=token_ids, top_logprobs=out_logprobs, - num_output_top_logprobs=sampling_params.logprobs, tokenizer=tokenizer, ) else: logprobs = self._create_tokens_logprobs( token_ids=token_ids, top_logprobs=out_logprobs, - num_output_top_logprobs=sampling_params.logprobs, ) else: logprobs = None @@ -599,14 +597,12 @@ class ServingTokens(GenerateBaseServing): logprobs = self._create_text_logprobs( token_ids=delta_token_ids, top_logprobs=out_logprobs, - num_output_top_logprobs=sampling_params.logprobs, tokenizer=tokenizer, ) else: logprobs = self._create_tokens_logprobs( token_ids=delta_token_ids, top_logprobs=out_logprobs, - num_output_top_logprobs=sampling_params.logprobs, ) else: logprobs = None @@ -728,14 +724,15 @@ class ServingTokens(GenerateBaseServing): self, token_ids: GenericSequence[int], top_logprobs: GenericSequence[dict[int, Logprob] | None], - num_output_top_logprobs: int | None = None, tokenizer: TokenizerLike | None = None, ) -> ChatCompletionLogProbs: """Create OpenAI-style logprobs for ``output_mode="text"``. With a ``tokenizer`` the tokens are decoded strings and carry ``bytes``; without one (``--return-tokens-as-token-ids``) they are ``token_id:N`` - placeholders. + placeholders. Like tokens mode, ``top_logprobs`` carries every candidate + the engine returned (the sampled token first, then the top k), not cut + to the requested count. """ logprobs_content: list[ChatCompletionLogProbsContent] = [] @@ -757,14 +754,7 @@ class ServingTokens(GenerateBaseServing): bytes=token_bytes, top_logprobs=[ _top_logprob(top_id, logprob, tokenizer) - for i, (top_id, logprob) in enumerate( - step_top_logprobs.items() - ) - if num_output_top_logprobs is not None - and ( - num_output_top_logprobs == -1 - or i < max(num_output_top_logprobs, 1) - ) + for top_id, logprob in step_top_logprobs.items() ], ) ) @@ -775,10 +765,16 @@ class ServingTokens(GenerateBaseServing): self, token_ids: GenericSequence[int], top_logprobs: GenericSequence[dict[int, Logprob] | None], - num_output_top_logprobs: int | None = None, ) -> GenerateLogProbs: """Create generate-shaped logprobs (integer token ids, no tokenizer). + ``top_logprobs`` carries every candidate the engine returned for the + position: the sampled token first, then the top k, so k entries when the + sampled token is in the top k and k + 1 when it is not. It is not cut to + the requested count: `/v1/chat/completions` and + `/v1/completions` cut differently, and derender applies each endpoint's + cut the way the coupled server does. + The engine reports rank 0 for a sampled token whose logprob is NaN; that is not a rank, so it is sent as ``None`` (the logprob is clamped). """ @@ -806,14 +802,7 @@ class ServingTokens(GenerateBaseServing): logprob=_clamp_logprob(top_logprob.logprob), rank=top_logprob.rank or None, ) - for rank_index, (top_token_id, top_logprob) in enumerate( - step_top_logprobs.items() - ) - if num_output_top_logprobs is not None - and ( - num_output_top_logprobs == -1 - or rank_index < max(num_output_top_logprobs, 1) - ) + for top_token_id, top_logprob in step_top_logprobs.items() ], ) ) diff --git a/vllm/entrypoints/serve/elastic_ep/api_router.py b/vllm/entrypoints/serve/elastic_ep/api_router.py index 1dbfc96b9b..5a839948cd 100644 --- a/vllm/entrypoints/serve/elastic_ep/api_router.py +++ b/vllm/entrypoints/serve/elastic_ep/api_router.py @@ -8,6 +8,7 @@ from http import HTTPStatus from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request from fastapi.responses import JSONResponse +from vllm.distributed.elastic_ep.external_elastic_ep import ExternalElasticEPScalePhase from vllm.engine.protocol import EngineClient from vllm.entrypoints.serve.elastic_ep.middleware import get_scaling_elastic_ep from vllm.entrypoints.serve.engine.protocol import ErrorResponse @@ -82,7 +83,35 @@ async def scale_elastic_ep(raw_request: Request): @router.post("/is_scaling_elastic_ep") async def is_scaling_elastic_ep(raw_request: Request): - return JSONResponse({"is_scaling_elastic_ep": get_scaling_elastic_ep()}) + # External operation status comes from the shared store. Middleware gating + # remains process-local, so scaling requests must reach every old-rank API. + try: + status = await engine_client(raw_request).get_external_elastic_ep_status() + except Exception as e: + logger.warning("Failed to query external Elastic EP status: %s", e) + raise HTTPException( + status_code=503, + detail="External Elastic EP status is temporarily unavailable", + ) from e + if status is None: + # Non-external EEP modes retain the process-local middleware state. + is_scaling = get_scaling_elastic_ep() + status = { + "phase": ( + ExternalElasticEPScalePhase.COMMITTING.value + if is_scaling + else ExternalElasticEPScalePhase.IDLE.value + ), + "epoch": None, + "error": None, + "requested_data_parallel_size": None, + } + else: + is_scaling = status["phase"] in ( + ExternalElasticEPScalePhase.PREPARING.value, + ExternalElasticEPScalePhase.COMMITTING.value, + ) + return JSONResponse({"is_scaling_elastic_ep": is_scaling, **status}) def attach_router(app: FastAPI): diff --git a/vllm/entrypoints/serve/elastic_ep/middleware.py b/vllm/entrypoints/serve/elastic_ep/middleware.py index dcc62aea2c..d7f81fd7ef 100644 --- a/vllm/entrypoints/serve/elastic_ep/middleware.py +++ b/vllm/entrypoints/serve/elastic_ep/middleware.py @@ -6,8 +6,12 @@ from collections.abc import Awaitable from fastapi.responses import JSONResponse from starlette.types import ASGIApp, Receive, Scope, Send -# Global variable to track scaling state +# Process-local states used to pause serving during scaling and after this rank +# has been removed from the external-LB topology. A retired API process remains +# available for health and scaling-status queries until its orchestrator exits it. _scaling_elastic_ep = False +_elastic_ep_rank_retired = False +_SCALING_OBSERVABILITY_PATHS = frozenset({"/health", "/is_scaling_elastic_ep"}) def get_scaling_elastic_ep(): @@ -19,12 +23,20 @@ def set_scaling_elastic_ep(value): _scaling_elastic_ep = value -class ScalingMiddleware: - """Middleware that checks if the model is currently scaling and - returns a 503 Service Unavailable response if it is. +def get_elastic_ep_rank_retired(): + return _elastic_ep_rank_retired - This middleware applies to all HTTP requests and prevents - processing when the model is in a scaling state. + +def set_elastic_ep_rank_retired(value): + global _elastic_ep_rank_retired + _elastic_ep_rank_retired = value + + +class ScalingMiddleware: + """Pause serving while scaling or after the local rank is retired. + + Health and scaling-status requests remain available so an external + orchestrator can observe the operation and remove a retired process. """ def __init__(self, app: ASGIApp) -> None: @@ -34,13 +46,18 @@ class ScalingMiddleware: if scope["type"] != "http": return self.app(scope, receive, send) - # Check global scaling state - if get_scaling_elastic_ep(): - # Return 503 Service Unavailable response + # Keep observability available while serving is paused or this rank is + # waiting for the external orchestrator to remove its process. + rank_retired = get_elastic_ep_rank_retired() + serving_paused = get_scaling_elastic_ep() or rank_retired + if serving_paused and scope["path"] not in _SCALING_OBSERVABILITY_PATHS: + error = ( + "This data-parallel rank has been retired and is awaiting shutdown." + if rank_retired + else "The model is currently scaling. Please try again later." + ) response = JSONResponse( - content={ - "error": "The model is currently scaling. Please try again later." - }, + content={"error": error}, status_code=503, ) return response(scope, receive, send) diff --git a/vllm/forward_context.py b/vllm/forward_context.py index 9ddc8e9255..95bb83dee8 100644 --- a/vllm/forward_context.py +++ b/vllm/forward_context.py @@ -327,10 +327,8 @@ def set_forward_context( vllm_config.parallel_config, num_tokens or 0, num_tokens_across_dp ) - # Convenience: if cudagraph is used and num_tokens is given, we can just - # create a batch descriptor here if not given (there's no harm since if it - # doesn't match in the wrapper it'll fall through). - if cudagraph_runtime_mode != CUDAGraphMode.NONE and num_tokens is not None: + # Preserve the batch size before model-level sequence parallelism. + if num_tokens is not None: batch_descriptor = batch_descriptor or BatchDescriptor(num_tokens=num_tokens) additional_kwargs = current_platform.set_additional_forward_context( diff --git a/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py b/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py index aeeb2f63b4..0e3ae398f6 100644 --- a/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py +++ b/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py @@ -427,7 +427,7 @@ def sparse_mqa_logits_paged_decode( q: [rows, 1, H, D] packed Q values (MXFP4 viewed as int8). q_scale: [rows, 1, H] int32 packed UE8M0 Q scales. kv_cache: 4D [num_pages, page_kv, 1, head_bytes] uint8 fused cache - view; the page stride must be 512B-aligned. + view; the page stride must be 256B-aligned. weights: [rows, H] bf16 per-head weights. context_lens: [rows] int32 per-row (compressed) context lengths. block_table: [rows, P] int32 per-row page table, ``stride(-1) == 1``. @@ -451,8 +451,8 @@ def sparse_mqa_logits_paged_decode( rows = q.shape[0] page_kv = kv_cache.shape[1] assert weights.dtype == torch.bfloat16, weights.dtype - assert kv_cache.stride(0) % 512 == 0, ( - "paged cache page stride must be 512B-aligned: " + assert kv_cache.stride(0) % 256 == 0, ( + "paged cache page stride must be 256B-aligned: " f"shape={tuple(kv_cache.shape)} stride={kv_cache.stride()}" ) assert page_kv % sparse_block_kv == 0, (page_kv, sparse_block_kv) diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index ec3059ea7f..3743dfc1a4 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -88,6 +88,9 @@ from vllm.model_executor.kernels.linear.mxfp4.humming import ( from vllm.model_executor.kernels.linear.mxfp4.marlin import ( MarlinMxFp4LinearKernel, ) +from vllm.model_executor.kernels.linear.mxfp4.pytorch import ( + TorchMxFp4LinearKernel, +) from vllm.model_executor.kernels.linear.mxfp4.xpu import ( XPUMxFp4LinearKernel, ) @@ -341,6 +344,7 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = { ChannelWiseTorchFP8ScaledMMLinearKernel, RowWiseTorchFP8ScaledMMLinearKernel, BlockWiseTorchFP8ScaledMMLinearKernel, + TorchMxFp4LinearKernel, TorchMxFp8LinearKernel, TorchNvFp4LinearKernel, }, @@ -624,6 +628,7 @@ _POSSIBLE_MXFP4_KERNELS: dict[PlatformEnum, list[type[MxFp4LinearKernel]]] = { ], PlatformEnum.XPU: [ XPUMxFp4LinearKernel, + TorchMxFp4LinearKernel, ], } diff --git a/vllm/model_executor/kernels/linear/mxfp4/pytorch.py b/vllm/model_executor/kernels/linear/mxfp4/pytorch.py new file mode 100644 index 0000000000..5902b66114 --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp4/pytorch.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + xpu_mxfp4_quantize, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp4Dynamic, +) +from vllm.model_executor.utils import replace_parameter +from vllm.platforms import current_platform + +from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig + +_E8M0 = torch.float8_e8m0fnu +_FP4X2 = torch.float4_e2m1fn_x2 + + +class TorchMxFp4LinearKernel(MxFp4LinearKernel): + """MXFP4 W4A4 GEMM on XPU using the native ``torch._scaled_mm`` dispatch.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_xpu(): + return False, "TorchMxFp4 only supports XPU" + return True, None + + @classmethod + def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + if config.activation_quant_key != kMxfp4Dynamic: + return False, "only supports MXFP4 dynamic activation" + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # weight: [N, K_packed] packed-fp4; weight_scale: [N, K//32] E8M0. + # oneDNN wants un-swizzled scales and a column-major weight (B operand). + weight = layer.weight.view(_FP4X2) + replace_parameter(layer, "weight", weight.data.t()) + + weight_scale = layer.weight_scale.view(_E8M0) + weight_scale = weight_scale.t().contiguous() + replace_parameter(layer, "weight_scale", weight_scale.data) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + out_dtype = x.dtype + weight = layer.weight # column-major [K_packed, N] + N = weight.shape[1] + + input_2d = x.reshape(-1, x.shape[-1]) + x_fp4, x_scale = xpu_mxfp4_quantize(input_2d) + x_scale = x_scale.view(_E8M0) + + out = torch._scaled_mm( + x_fp4, + weight, + scale_a=x_scale, + scale_b=layer.weight_scale, + bias=bias, + out_dtype=out_dtype, + ) + + return out.reshape(*x.shape[:-1], N) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py index 03e9be51cf..23a5886ab7 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py @@ -9,6 +9,7 @@ from vllm._aiter_ops import ( rocm_aiter_ops, ) from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 from vllm.model_executor.layers.quantization.utils.fp8_utils import ( _upcast_e8m0_to_fp32, ) @@ -399,6 +400,26 @@ class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): rocm_aiter_ops.is_triton_gemm_w8a8_tuned(n, k) or _on_gfx1250 ) + self.use_bpreshuffle = ( + _on_gfx1250 + and n % 16 == 0 + and k % 128 == 0 + and bool(rocm_aiter_ops.is_triton_blockscale_bpreshuffle_gemm_tuned(n, k)) + ) + if self.use_bpreshuffle: + # The tuned preshuffle configs expect a column-major group scale. + act_scale_descriptor = config.activation_quant_key.scale + self.quant_fp8 = QuantFP8( + static=act_scale_descriptor.static, + group_shape=act_scale_descriptor.group_shape, + num_token_padding=self.get_output_padding(), + use_ue8m0=False, + column_major_scales=True, + ) + logger.info_once( + "Using aiter A8W8 blockscale B-preshuffle GEMM for N=%d K=%d", n, k + ) + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: super().process_weights_after_loading(layer) @@ -410,6 +431,18 @@ class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): if ws is not None and ws.dtype == torch.float8_e8m0fnu: replace_parameter(layer, attr, _upcast_e8m0_to_fp32(ws).contiguous()) + if getattr(layer, "is_bmm", False): + self.use_bpreshuffle = False + + if self.use_bpreshuffle: + w = FP8BlockParams.from_layer(layer).weight + replace_parameter( + layer, + FP8BlockParams.WEIGHT, + rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)), + ) + layer.aiter_bpreshuffled = True + @classmethod def is_supported(cls, compute_capability=None): if ( @@ -464,7 +497,11 @@ class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): Bs = Bs.to(torch.float32) out_dtype = self.config.out_dtype - if self.use_triton: + if self.use_bpreshuffle: + gemm_a8w8_blockscale_op = ( + rocm_aiter_ops.triton_gemm_a8w8_blockscale_preshuffle + ) + elif self.use_triton: gemm_a8w8_blockscale_op = rocm_aiter_ops.triton_gemm_a8w8_blockscale else: gemm_a8w8_blockscale_op = rocm_aiter_ops.gemm_a8w8_blockscale diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index e79700ae44..e8bd822bec 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -249,6 +249,8 @@ class FusedMoEQuantConfig: _w1: FusedMoEQuantDesc _w2: FusedMoEQuantDesc is_scale_swizzled: bool = True + # Emit block-quant activation scales column-major (physical [G, M] layout). + is_scale_transposed: bool = False # MXFP4-specific TRTLLM parameters for SwiGLU activation clamping. # These correspond to gemm1_alpha, gemm1_beta, gemm1_clamp_limit @@ -503,6 +505,7 @@ class FusedMoEQuantConfig: w2_zp: torch.Tensor | None = None, weight_dtype: torch.dtype | str | None = None, is_scale_swizzled: bool = True, + is_scale_transposed: bool = False, gemm1_alpha: float | None = None, gemm1_beta: float | None = None, gemm1_clamp_limit: float | None = None, @@ -537,6 +540,9 @@ class FusedMoEQuantConfig: - is_scale_swizzled: Whether the activation scale-factor layout is swizzled. Pass through to the underlying quantization kernel for dtypes that distinguish layouts (nvfp4, mxfp8). Defaults to True. + - is_scale_transposed: Whether block-quantized activation scales are + emitted column-major (physical [num_groups, num_tokens] layout) by + the quantize step. Defaults to False. - gemm1_alpha: Optional MXFP4 TRTLLM SwiGLU alpha parameter. - gemm1_beta: Optional MXFP4 TRTLLM SwiGLU beta parameter. - gemm1_clamp_limit: Optional MXFP4 TRTLLM SwiGLU clamp limit. @@ -573,6 +579,7 @@ class FusedMoEQuantConfig: weight_dtype, w_shape, w2_scale, g2_alphas, w2_zp, w2_bias ), is_scale_swizzled=is_scale_swizzled, + is_scale_transposed=is_scale_transposed, gemm1_alpha=gemm1_alpha, gemm1_beta=gemm1_beta, gemm1_clamp_limit=gemm1_clamp_limit, @@ -597,6 +604,7 @@ def fp8_w8a8_moe_quant_config( a2_gscale: torch.Tensor | None = None, g1_alphas: torch.Tensor | None = None, g2_alphas: torch.Tensor | None = None, + is_scale_transposed: bool = False, gemm1_alpha: float | None = None, gemm1_beta: float | None = None, gemm1_clamp_limit: float | None = None, @@ -617,6 +625,7 @@ def fp8_w8a8_moe_quant_config( per_act_token_quant=per_act_token_quant, per_out_ch_quant=per_out_ch_quant, block_shape=block_shape, + is_scale_transposed=is_scale_transposed, gemm1_alpha=gemm1_alpha, gemm1_beta=gemm1_beta, gemm1_clamp_limit=gemm1_clamp_limit, diff --git a/vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py b/vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py index b86820f672..69a234b081 100644 --- a/vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py +++ b/vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py @@ -1,170 +1,559 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""DeepGEMM MegaMoE backends: fused EP dispatch, expert GEMMs and combine.""" +"""DeepGEMM MegaMoE experts: fused EP dispatch, expert GEMMs and combine. +Each ``MEGA_MOE_*`` format is one row of the SM100 feature matrix. MX weights +and activations carry one UE8M0 scale per 1x32 block. NVFP4 weights carry one +E4M3 scale per 1x16 block plus a per-expert FP32 global scale; NVFP4 dispatch +scales each token by its own FP32 scale. +""" + +import math +from collections.abc import Callable +from dataclasses import dataclass from inspect import signature -from typing import Any +from typing import Any, ClassVar, Literal import torch import torch.nn as nn +import vllm.envs as envs +from vllm.config import VllmConfig +from vllm.distributed import get_ep_group +from vllm.distributed.eplb.eplb_state import EplbLayerState +from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.prepare_megamoe import ( + prepare_megamoe_inputs, + prepare_megamoe_nvfp4_inputs, +) +from vllm.model_executor.layers.fused_moe.router.base_router import ( + eplb_map_to_physical_and_record, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kMxfp8Dynamic, + kNvfp4DynamicToken, ) +from vllm.model_executor.models.utils import extract_layer_index +from vllm.model_executor.utils import set_weight_attrs from vllm.platforms import current_platform +from vllm.v1.worker.ubatching import dbo_current_ubatch_id logger = init_logger(__name__) -class DeepGemmMegaMoEBackend: - """Abstract MegaMoE backend that hides deep_gemm-specific details. +@dataclass(frozen=True) +class MegaMoEActivation: + """Gated activation fused into the MegaMoE L1 epilogue. - A backend advertises its expected hidden-state quantization layout via - ``hidden_quant`` and its expected MMA type via ``mma_type``. It owns the - kernel weight layout: :meth:`create_weights` registers parameters in that - layout and :meth:`load_weight` writes each checkpoint tensor into it, so - loading (and reloading) needs no post-load transform. It also runs the - MegaMoE kernel. Runtime capability checks live in - :func:`get_deep_gemm_mega_moe_backend`: obtaining a backend already implies - that the current device and shapes are supported. + ``swiglu`` computes ``gate * sigmoid(alpha * gate) * (up + beta)`` after + clamping gate from above and up symmetrically to ``clamp``; this covers + plain, DeepSeek-V4 (``clamp``) and GPT-OSS style (``alpha``/``beta``) + experts. ``situ`` is Kimi K3's SiTU with gate tanh scale ``alpha`` and + linear tanh scale ``beta``; ``beta=0`` leaves the linear branch unscaled. """ - # Human-readable MMA type tag used by deep_gemm helpers. SM100 default. - mma_type: str = "fp8xfp4" - hidden_quant: QuantKey = kMxfp8Dynamic + name: Literal["swiglu", "situ"] = "swiglu" + clamp: float | None = None + alpha: float = 1.0 + beta: float = 0.0 - def create_weights( + def __post_init__(self) -> None: + if self.name not in ("swiglu", "situ"): + raise ValueError(f"Unsupported MegaMoE activation {self.name!r}.") + if not (math.isfinite(self.alpha) and math.isfinite(self.beta)): + raise ValueError(f"MegaMoE activation scales must be finite: {self}.") + if self.clamp is not None and not self.clamp >= 0: + raise ValueError(f"MegaMoE activation clamp must be >= 0: {self}.") + if self.name == "situ" and ( + self.clamp is not None or self.alpha <= 0 or self.beta < 0 + ): + raise ValueError( + "SiTU takes no clamp, a positive alpha and a non-negative beta; " + f"got {self}." + ) + + def kernel_kwargs(self) -> dict[str, Any]: + # Defaults stay implicit so DeepGEMM builds without alpha/beta still work. + kwargs: dict[str, Any] = {"activation_clamp": self.clamp} + if self.name != "swiglu": + kwargs["activation"] = self.name + if self.alpha != 1.0: + kwargs["activation_alpha"] = self.alpha + if self.beta != 0.0: + kwargs["activation_beta"] = self.beta + return kwargs + + +@dataclass(frozen=True) +class MegaMoEWeightSpec: + """Storage of one MegaMoE weight in the kernel layout.""" + + dtype: torch.dtype + # Weight values packed per storage element along K (2 for FP4). + values_per_element: int = 1 + # K values per block scale; None for unscaled weights. + scale_group_k: int | None = None + # Block scale dtype: UE8M0 (MX) or E4M3 (NVFP4). + scale_dtype: torch.dtype | None = None + + @property + def has_global_scale(self) -> bool: + """E4M3 block scales pair with an FP32 per-tensor scale.""" + return self.scale_dtype == torch.float8_e4m3fn + + +_MXFP4_WEIGHT = MegaMoEWeightSpec( + torch.int8, + values_per_element=2, + scale_group_k=32, + scale_dtype=torch.float8_e8m0fnu, +) +_MXFP8_WEIGHT = MegaMoEWeightSpec( + torch.float8_e4m3fn, scale_group_k=32, scale_dtype=torch.float8_e8m0fnu +) +_NVFP4_WEIGHT = MegaMoEWeightSpec( + torch.int8, + values_per_element=2, + scale_group_k=16, + scale_dtype=torch.float8_e4m3fn, +) +_BF16_WEIGHT = MegaMoEWeightSpec(torch.bfloat16) + + +@dataclass(frozen=True) +class MegaMoEWeightFormat: + """One routed-expert precision of the MegaMoE feature matrix.""" + + name: str + # DeepGEMM MMA tag; it also keys the symmetric-buffer layout. + mma_type: str + kernel: str + routed: MegaMoEWeightSpec + # Precisions a shared expert may have to run inside the same launch. + shared: tuple[MegaMoEWeightSpec, ...] + # Quantization of the dispatched hidden states; None dispatches BF16. + hidden_quant: QuantKey | None + activations: frozenset[str] + + @property + def uses_nvfp4_api(self) -> bool: + """``nvfp4_mega_moe`` has its own symmetric buffer, BLOCK_M and scales.""" + return self.kernel == "nvfp4_mega_moe" + + +MEGA_MOE_MXFP4 = MegaMoEWeightFormat( + name="mxfp4", + mma_type="fp8xfp4", + kernel="fp8_fp4_mega_moe", + routed=_MXFP4_WEIGHT, + shared=(_MXFP8_WEIGHT,), + hidden_quant=kMxfp8Dynamic, + activations=frozenset({"swiglu", "situ"}), +) +MEGA_MOE_MXFP8 = MegaMoEWeightFormat( + name="mxfp8", + mma_type="fp8xfp8", + kernel="fp8_fp4_mega_moe", + routed=_MXFP8_WEIGHT, + shared=(_MXFP8_WEIGHT,), + hidden_quant=kMxfp8Dynamic, + activations=frozenset({"swiglu", "situ"}), +) +MEGA_MOE_NVFP4 = MegaMoEWeightFormat( + name="nvfp4", + mma_type="nvfp4", + kernel="nvfp4_mega_moe", + routed=_NVFP4_WEIGHT, + shared=(_MXFP8_WEIGHT, _BF16_WEIGHT), + hidden_quant=kNvfp4DynamicToken, + activations=frozenset({"swiglu"}), +) +MEGA_MOE_BF16 = MegaMoEWeightFormat( + name="bf16", + mma_type="bf16xbf16", + kernel="bf16_mega_moe", + routed=_BF16_WEIGHT, + shared=(_BF16_WEIGHT,), + hidden_quant=None, + activations=frozenset({"swiglu"}), +) +MEGA_MOE_WEIGHT_FORMATS = { + fmt.name: fmt + for fmt in (MEGA_MOE_MXFP4, MEGA_MOE_MXFP8, MEGA_MOE_NVFP4, MEGA_MOE_BF16) +} + + +def check_deep_gemm_mega_moe_supported( + device: torch.device, hidden_size: int, intermediate_size: int +) -> None: + """Raise unless DeepGEMM MegaMoE can run on ``device`` with these shapes. + + Raises: + NotImplementedError: If the device architecture is unsupported. + ValueError: If the shapes are unsupported. + + """ + if not current_platform.is_device_capability_family( + 100, device_id=device.index or 0 + ): + raise NotImplementedError("DeepGEMM MegaMoE requires SM100 GPUs.") + if hidden_size % 128 != 0 or intermediate_size % 128 != 0: + raise ValueError( + "DeepGEMM MegaMoE requires hidden and intermediate sizes " + "to be multiples of 128." + ) + + +def _linear_weight_scale(linear: nn.Module) -> torch.Tensor | None: + scale = getattr(linear, "weight_scale", None) + return getattr(linear, "weight_scale_inv", None) if scale is None else scale + + +def _matches_weight_spec(linear: nn.Module, spec: MegaMoEWeightSpec) -> bool: + """Whether a checkpoint linear can be written into a ``spec`` weight.""" + weight, scale = linear.weight, _linear_weight_scale(linear) + if weight.dtype != spec.dtype or weight.dim() != 2: + return False + if spec.scale_group_k is None: + return scale is None + if scale is None or scale.dim() != 2: + return False + (rows, cols), (scale_rows, scale_cols) = weight.shape, scale.shape + scale_dtypes = {spec.scale_dtype} + if spec.scale_dtype == torch.float8_e8m0fnu: + scale_dtypes.add(torch.uint8) + return ( + scale.dtype in scale_dtypes + and rows % scale_rows == 0 + and cols % scale_cols == 0 + and (cols // scale_cols) % spec.scale_group_k == 0 + ) + + +class DeepGemmMegaMoEExperts(nn.Module): + """Routed experts, and optionally the shared expert, as one MegaMoE launch. + + Parameters are registered in the kernel layout and every checkpoint shard is + written straight into it, so loading (and reloading) needs no post-load + transform. Whether the shared expert runs inside the launch is decided here, + before any weight loads; a fused layer owns the shared-expert weights and + the caller drops its shared MLP. + + NVFP4 experts also take per-expert ``weight_scale_2`` (``w1`` and ``w3`` + must match: one L1 scale per expert) and ``w2`` ``input_scale``; the kernel + quantizes every L2 input with the max of the latter over all loaded + experts. Dispatch uses per-token scales, so ``w1``/``w3`` ``input_scale`` + are not loaded. + + Args: + num_experts: Physical experts across EP ranks, redundant included. + activation: Activation fused into the L1 epilogue. + weight_format: Routed-expert precision, a row of the feature matrix. + shared_experts: Shared MLP to fuse if it is unsharded and in one of + the ``weight_format.shared`` precisions; ``None`` never fuses. + + """ + + _symm_buffers: ClassVar[dict[tuple, Any]] = {} + _synchronized_ep_groups: ClassVar[set[tuple[int, int]]] = set() + + def __init__( self, - layer: nn.Module, + vllm_config: VllmConfig, *, + num_experts: int, num_local_experts: int, + experts_start_idx: int, + top_k: int, hidden_size: int, intermediate_size: int, - num_shared_experts: int, - ) -> None: - """Register the expert (and fused shared-expert) kernel-layout parameters: - ``w13_weight``, ``w13_weight_scale``, ``w2_weight``, ``w2_weight_scale`` - and, if ``num_shared_experts``, ``shared_l{1,2}_weight[_scale]``. Routed - parameters are indexed by local expert on dim 0, which EPLB moves.""" - raise NotImplementedError + activation: MegaMoEActivation, + weight_format: MegaMoEWeightFormat, + shared_experts: nn.Module | None = None, + num_logical_experts: int | None = None, + fast_math: bool = True, + prefix: str = "", + ): + super().__init__() + if activation.name not in weight_format.activations: + raise NotImplementedError( + f"DeepGEMM MegaMoE {weight_format.name} experts support " + f"{sorted(weight_format.activations)}, not {activation.name!r}." + ) + device = torch.get_default_device() + if device.type != current_platform.device_type: # meta under ipc_cache + device = torch.device( + current_platform.device_type, torch.accelerator.current_device_index() + ) + check_deep_gemm_mega_moe_supported(device, hidden_size, intermediate_size) - def load_weight( - self, - dst: torch.Tensor, - src: torch.Tensor, - shard_id: str, - is_scale: bool, - ) -> None: - """Write checkpoint shard ``src`` (``w1``/``w2``/``w3``) into ``dst``, one - expert's slice or a shared-expert parameter, in the kernel layout.""" - raise NotImplementedError + self.prefix = prefix + self.capture_fn: Callable[[torch.Tensor], None] | None = None + self.num_experts = num_experts + self.num_local_experts = num_local_experts + self.experts_start_idx = experts_start_idx + self.experts_end_idx = experts_start_idx + num_local_experts + self.num_logical_experts = ( + num_logical_experts if num_logical_experts is not None else num_experts + ) + self.top_k = top_k + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens + self.weight_format = weight_format + self.activation = activation + self.fast_math = fast_math + self._kernel_kwargs = activation.kernel_kwargs() + self.eplb_state = EplbLayerState() + # L2 input scale of every logical expert; the NVFP4 kernel takes the max. + self._w2_input_scales = torch.zeros(self.num_logical_experts, device="cpu") + self.l2_activation_scale = 1.0 + # The first of each expert's w1/w3 weight_scale_2 until the other loads. + self._pending_l1_scales: dict[int, tuple[str, float]] = {} - def supports_shared_experts(self, shared_experts: nn.Module) -> bool: - """Whether ``shared_experts`` can be fused into the MegaMoE kernel. + self.num_shared_experts, self.shared_spec = self._fusable_shared_experts( + shared_experts + ) + self._create_weights() + for name, param in self.named_parameters(recurse=False): + shared = name.startswith("shared_") + loader = self.shared_weight_loader if shared else self.weight_loader + set_weight_attrs(param, {"weight_loader": loader}) + if not shared and name.endswith("weight_scale"): + param.quant_method = "block" - Decided at construction from the installed kernel API and the shared - MLP's parameters, so the fused layout is fixed before any weight loads. + # Register in the static forward context so the custom-op wrapper + # can look up this module by name from within a torch.compile graph. + compilation_config = vllm_config.compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + @property + def has_fused_shared_experts(self) -> bool: + return self.num_shared_experts > 0 + + def _fusable_shared_experts( + self, shared_experts: nn.Module | None + ) -> tuple[int, MegaMoEWeightSpec | None]: + """How many shared experts the launch fuses, and in which precision.""" + if shared_experts is None: + return 0, None + gate_up, down = shared_experts.gate_up_proj, shared_experts.down_proj + hidden, inter = self.hidden_size, self.intermediate_size + shared_size = down.weight.shape[-1] + spec = next( + ( + spec + for spec in self.weight_format.shared + if all(_matches_weight_spec(lin, spec) for lin in (gate_up, down)) + ), + None, + ) + if ( + tuple(gate_up.weight.shape) != (2 * shared_size, hidden) + or tuple(down.weight.shape) != (hidden, shared_size) + or shared_size % inter + ): + reason = ( + f"needs an unsharded shared MLP over hidden size {hidden} with a " + f"multiple of {inter} intermediate; got {tuple(gate_up.weight.shape)}" + f" and {tuple(down.weight.shape)}" + ) + elif spec is None: + scale = _linear_weight_scale(gate_up) + supported = ", ".join( + f"{s.dtype}{' with 1x32 UE8M0 scales' if s.scale_group_k else ''}" + for s in self.weight_format.shared + ) + reason = ( + f"{self.weight_format.name} experts fuse {supported} shared " + f"experts; got weight {gate_up.weight.dtype} and scale " + f"{None if scale is None else (tuple(scale.shape), scale.dtype)}" + ) + elif not self._kernel_fuses_shared_experts(spec): + reason = ( + "the installed DeepGEMM Python API is older than the vLLM source. " + "Rebuild the vendored _deep_gemm_C extension to enable it" + ) + else: + return shared_size // inter, spec + logger.warning_once("Disabling MegaMoE shared-expert fusion: %s.", reason) + return 0, None + + def _kernel_fuses_shared_experts(self, spec: MegaMoEWeightSpec) -> bool: + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + fmt = self.weight_format + buffer_api = "get_symm_buffer_for_mega_moe" + block_m_api = "get_block_m_for_mega_moe" + if fmt.uses_nvfp4_api: + buffer_api = "NVFP4SymmBuffer" + block_m_api = "get_block_m_for_nvfp4_mega_moe" + try: + api = { + *signature(getattr(deep_gemm, buffer_api)).parameters, + *signature(getattr(deep_gemm, fmt.kernel)).parameters, + } + except (AttributeError, TypeError, ValueError): + return False + # Scaled shared inputs are laid out by the runtime BLOCK_M. + return { + "num_shared_experts", + "shared_l1_weights", + "shared_l2_weights", + } <= api and (spec.scale_group_k is None or hasattr(deep_gemm, block_m_api)) + + def _create_weights(self) -> None: + """Register the routed and fused shared-expert parameters. + + Routed parameters are indexed by local expert on dim 0, which EPLB moves. """ - return False + e, h, i = self.num_local_experts, self.hidden_size, self.intermediate_size + fmt = self.weight_format + self._register_weight("w13", (e, 2 * i, h), fmt.routed) + self._register_weight("w2", (e, h, i), fmt.routed) + if fmt.routed.has_global_scale: + # w13_weight_scale_2 is the kernel's L1 alpha (per-token dispatch). + for name in ("w13_weight_scale_2", "w2_weight_scale_2", "w2_input_scale"): + scale = torch.zeros(e, dtype=torch.float32) + self.register_parameter(name, nn.Parameter(scale, requires_grad=False)) + # w2 global weight scale times the L2 activation scale, per launch. + l2_alpha = torch.zeros(e, dtype=torch.float32) + self.register_buffer("_l2_alpha", l2_alpha, persistent=False) + if self.shared_spec is not None: + # The shared expert has the layout of one routed expert slice. + s = i * self.num_shared_experts + self._register_weight("shared_l1", (2 * s, h), self.shared_spec) + self._register_weight("shared_l2", (h, s), self.shared_spec) - def run_mega_moe( - self, - *, - y: torch.Tensor, - l1_weights: tuple[torch.Tensor, torch.Tensor], - l2_weights: tuple[torch.Tensor, torch.Tensor], - symm_buffer, - activation_clamp: float | None, - fast_math: bool, - activation: str | None = None, - activation_alpha: float | None = None, - activation_beta: float | None = None, - shared_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None, - shared_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None, - ) -> None: - raise NotImplementedError - - def supports_bf16_mega_gate(self) -> bool: - """Whether this backend can run the DeepGEMM bf16 MegaMoE gate kernel. - - Defaults to ``False`` (the symbol is not guaranteed to exist in every - deep_gemm build). Backends that target a deep_gemm build exposing the - ``bf16_mega_gate`` symbol override this and :meth:`bf16_mega_gate`. - """ - return False - - def bf16_mega_gate( - self, - *, - x: torch.Tensor, - weight: torch.Tensor, - num_topk: int, - scoring_func: str, - routed_scaling_factor: float, - ep_rank: int, - bias: torch.Tensor | None = None, - image_bias: torch.Tensor | None = None, - image_token_mask: torch.Tensor | None = None, - fix_routing_mask: torch.Tensor | None = None, - unmapped_topk_idx: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - raise NotImplementedError - - -class DeepGemmSm100MegaMoEBackend(DeepGemmMegaMoEBackend): - """Default MegaMoE backend targeting SM100 via deep_gemm fp8_fp4_mega_moe.""" - - mma_type = "fp8xfp4" - hidden_quant = kMxfp8Dynamic - - def create_weights( - self, - layer: nn.Module, - *, - num_local_experts: int, - hidden_size: int, - intermediate_size: int, - num_shared_experts: int, + def _register_weight( + self, name: str, shape: tuple[int, ...], spec: MegaMoEWeightSpec ) -> None: def register(name: str, data: torch.Tensor) -> None: - layer.register_parameter(name, nn.Parameter(data, requires_grad=False)) + self.register_parameter(name, nn.Parameter(data, requires_grad=False)) - # FP4 weights two per byte; UE8M0 scales MN-major, four per int32 along K. - e, h, i = num_local_experts, hidden_size, intermediate_size - register("w13_weight", torch.zeros(e, 2 * i, h // 2, dtype=torch.int8)) + *outer, n, k = shape register( - "w13_weight_scale", - torch.zeros(e, h // 128, 2 * i, dtype=torch.int32).transpose(1, 2), - ) - register("w2_weight", torch.zeros(e, h, i // 2, dtype=torch.int8)) - register( - "w2_weight_scale", - torch.zeros(e, i // 128, h, dtype=torch.int32).transpose(1, 2), - ) - if num_shared_experts == 0: - return - # The shared expert is FP8 with the layout of one routed expert slice. - s, fp8 = i * num_shared_experts, torch.float8_e4m3fn - register("shared_l1_weight", torch.zeros(2 * s, h, dtype=fp8)) - register( - "shared_l1_weight_scale", - torch.zeros(h // 128, 2 * s, dtype=torch.int32).transpose(0, 1), - ) - register("shared_l2_weight", torch.zeros(h, s, dtype=fp8)) - register( - "shared_l2_weight_scale", - torch.zeros(s // 128, h, dtype=torch.int32).transpose(0, 1), + f"{name}_weight", + torch.zeros(*outer, n, k // spec.values_per_element, dtype=spec.dtype), ) + if spec.scale_group_k is not None: + # Block scales, MN-major, four per int32 along K. + scale_k = k // (4 * spec.scale_group_k) + scale = torch.zeros(*outer, scale_k, n, dtype=torch.int32) + register(f"{name}_weight_scale", scale.transpose(-1, -2)) - def load_weight( + def _kernel_weights( + self, name: str + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + weight = getattr(self, f"{name}_weight") + scale = getattr(self, f"{name}_weight_scale", None) + return weight if scale is None else (weight, scale) + + @property + def l1_weights(self) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + return self._kernel_weights("w13") + + @property + def l2_weights(self) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + return self._kernel_weights("w2") + + def _map_global_expert_id(self, expert_id: int) -> list[int]: + """Local slots that logical expert ``expert_id`` occupies on this rank.""" + physical_ids: list[int] = [] + for p in range(self.experts_start_idx, self.experts_end_idx): + if p % self.num_logical_experts == expert_id: + physical_ids.append(p - self.experts_start_idx) + return physical_ids + + def weight_loader( + self, + param: nn.Parameter, + loaded_weight: torch.Tensor, + weight_name: str, + shard_id: str, + expert_id: int, + return_success: bool = False, + ) -> bool | None: + if not loaded_weight.is_meta: + self._record_global_scale(weight_name, loaded_weight, shard_id, expert_id) + local_expert_ids = self._map_global_expert_id(expert_id) + param_prefix = "w2_" if shard_id == "w2" else "w13_" + if not local_expert_ids or param_prefix not in weight_name: + return False if return_success else None + is_global_scale = weight_name.endswith(("weight_scale_2", "input_scale")) + is_scale = weight_name.endswith("weight_scale") + for local_expert_id in local_expert_ids: + if is_global_scale: + param.data[local_expert_id].copy_(loaded_weight.reshape(())) + else: + self._load_weight( + param.data[local_expert_id], loaded_weight, shard_id, is_scale + ) + return True if return_success else None + + def _record_global_scale( + self, weight_name: str, scale: torch.Tensor, shard_id: str, expert_id: int + ) -> None: + """Track NVFP4 global scales the kernel takes one of per launch or expert. + + Every rank sees the scales of all experts, so ranks agree and EPLB can + move experts. + """ + if weight_name.endswith("w2_input_scale"): + self._w2_input_scales[expert_id] = scale.max() + self.l2_activation_scale = float(self._w2_input_scales.max()) or 1.0 + elif weight_name.endswith("w13_weight_scale_2"): + # Compare w1 and w3 of the same load, not the live (maybe stale) slot. + prev = self._pending_l1_scales.pop(expert_id, None) + value = float(scale) + if prev is None or prev[0] == shard_id: + self._pending_l1_scales[expert_id] = (shard_id, value) + elif not math.isclose(value, prev[1], rel_tol=1e-5): + raise ValueError( + f"MegaMoE NVFP4 expert {expert_id}: {prev[0]} weight_scale_2 " + f"{prev[1]} != {shard_id} weight_scale_2 {value}; the kernel " + "takes one L1 scale per expert." + ) + + def load_shared_expert_weight(self, name: str, loaded_weight: torch.Tensor) -> str: + """Load one shared-expert checkpoint tensor; return the parameter name.""" + if ".w2." in name or ".down_proj." in name: + param_name = "shared_l2_weight" + elif ".w1." in name or ".w3." in name: + param_name = "shared_l1_weight" + else: + raise ValueError(f"Unexpected shared-expert tensor: {name}") + if "weight_scale" in name: + param_name += "_scale" + param = getattr(self, param_name) + param.weight_loader(param, loaded_weight, name) + return f"{name.partition('shared_experts.')[0]}experts.{param_name}" + + def shared_weight_loader( + self, param: nn.Parameter, loaded_weight: torch.Tensor, weight_name: str + ) -> None: + shard_id = ( + "w1" if ".w1." in weight_name else "w3" if ".w3." in weight_name else "w2" + ) + is_scale = "weight_scale" in weight_name + self._load_weight(param.data, loaded_weight, shard_id, is_scale) + + def _load_weight( self, dst: torch.Tensor, src: torch.Tensor, shard_id: str, is_scale: bool, ) -> None: + """Write checkpoint shard ``src`` into ``dst`` in the kernel layout. + + ``dst`` is one routed expert's slice or a shared-expert parameter. + """ if shard_id not in ("w1", "w2", "w3"): raise ValueError(f"Unsupported expert shard id: {shard_id}") rows = dst.shape[0] // (1 if shard_id == "w2" else 2) @@ -182,7 +571,7 @@ class DeepGemmSm100MegaMoEBackend(DeepGemmMegaMoEBackend): @staticmethod def _pack_scale(sf: torch.Tensor, rows: int, cols: int) -> torch.Tensor: - """Expand block UE8M0 scales to one per 1x32 and pack four along K per int32.""" + """Expand block scales to the kernel granularity, four per int32 along K.""" sf = sf.view(torch.uint8) if rows % sf.shape[0] or cols % sf.shape[1]: raise ValueError( @@ -214,163 +603,255 @@ class DeepGemmSm100MegaMoEBackend(DeepGemmMegaMoEBackend): dst = dst.view(-1, 2, 2, 8, 4, k)[:, :, half].permute(0, 3, 1, 2, 4) return dst, src.reshape(-1, 4, 2, 8, k) - def supports_shared_experts(self, shared_experts: nn.Module) -> bool: - for linear in (shared_experts.gate_up_proj, shared_experts.down_proj): - weight = linear.weight - scale = getattr(linear, "weight_scale", None) - if scale is None: - scale = getattr(linear, "weight_scale_inv", None) - if not self._tiles_1x32(weight, scale): - logger.warning_once( - "Disabling native MegaMoE shared-expert fusion: needs FP8 " - "weights with UE8M0 block scales that tile to 1x32; got " - "weight %s/%s and scale %s/%s.", - tuple(weight.shape), - weight.dtype, - None if scale is None else tuple(scale.shape), - None if scale is None else scale.dtype, - ) - return False - + def get_symm_buffer(self): from vllm.utils.deep_gemm import _import_deep_gemm deep_gemm = _import_deep_gemm() - try: - api = { - *signature(deep_gemm.get_symm_buffer_for_mega_moe).parameters, - *signature(deep_gemm.fp8_fp4_mega_moe).parameters, - } - except (AttributeError, TypeError, ValueError): - api = set() - if not ( - hasattr(deep_gemm, "get_block_m_for_mega_moe") - and {"num_shared_experts", "shared_l1_weights", "shared_l2_weights"} <= api - ): - logger.warning_once( - "Disabling native MegaMoE shared-expert fusion because the " - "installed DeepGEMM Python API is older than the vLLM " - "source. Rebuild the vendored _deep_gemm_C extension to enable it.", + group = get_ep_group().device_group + key = ( + id(group), + torch.accelerator.current_device_index(), + self.num_experts, + self.max_num_tokens, + self.top_k, + self.hidden_size, + self.intermediate_size, + self.num_shared_experts, + None if self.shared_spec is None else self.shared_spec.dtype, + self.weight_format.mma_type, + self.activation.name, + ) + symm_buffer = self._symm_buffers.get(key) + if symm_buffer is None: + args = ( + group, + self.num_experts, + self.max_num_tokens, + self.top_k, + self.hidden_size, + self.intermediate_size, ) - return False - return True + if self.weight_format.uses_nvfp4_api: + symm_buffer = deep_gemm.NVFP4SymmBuffer( + *args, + num_shared_experts=self.num_shared_experts, + shared_dtype=torch.float8_e4m3fn + if self.shared_spec is None + else self.shared_spec.dtype, + ) + else: + symm_buffer = deep_gemm.get_symm_buffer_for_mega_moe( + *args, + num_shared_experts=self.num_shared_experts, + mma_type=self.weight_format.mma_type, + activation=self.activation.name, + ) + self._symm_buffers[key] = symm_buffer + return symm_buffer - @staticmethod - def _tiles_1x32(weight: torch.Tensor, scale: torch.Tensor | None) -> bool: - if scale is None or weight.dim() != 2 or scale.dim() != 2: - return False - (rows, cols), (scale_rows, scale_cols) = weight.shape, scale.shape - return ( - weight.dtype == torch.float8_e4m3fn - and scale.dtype in (torch.uint8, torch.float8_e8m0fnu) - and rows % scale_rows == 0 - and cols % scale_cols == 0 - and (cols // scale_cols) % 32 == 0 + def get_block_m(self, num_tokens: int) -> int: + """BLOCK_M the MegaMoE scheduler runs ``num_tokens`` local tokens with.""" + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + args = ( + get_ep_group().world_size, + self.num_experts, + self.get_symm_buffer().num_max_tokens_per_rank, + num_tokens, + self.top_k, + ) + if self.weight_format.uses_nvfp4_api: + return deep_gemm.get_block_m_for_nvfp4_mega_moe(*args) + return deep_gemm.get_block_m_for_mega_moe(*args, self.weight_format.mma_type) + + def set_eplb_state( + self, + moe_layer_idx: int, + expert_load_view: torch.Tensor, + logical_to_physical_map: torch.Tensor, + logical_replica_count: torch.Tensor, + ) -> None: + self.eplb_state.set_layer_state( + moe_layer_idx, + expert_load_view, + logical_to_physical_map, + logical_replica_count, ) - def run_mega_moe( + def get_expert_weights(self) -> list[torch.Tensor]: + def _to_eplb_view(name: str, t: torch.Tensor) -> torch.Tensor: + """Return a (num_local_experts, -1) view with contiguous memory layout.""" + assert t.shape[0] == self.num_local_experts + if t.is_contiguous(): + return t.view(self.num_local_experts, -1) + elif t.dim() == 3 and t.stride(1) == 1 and t.stride(2) == t.shape[1]: + # scales have shape (E, M, N) with memory layout (E, N, M) + back = torch.transpose(t, 1, 2) + assert back.is_contiguous() + return back.view(self.num_local_experts, -1) + + raise AssertionError( + f"MegaMoE EPLB {name}: non-contiguous expert tensor with " + f"unexpected layout shape={tuple(t.shape)} " + f"stride={tuple(t.stride())} dtype={t.dtype}" + ) + + return [ + _to_eplb_view(name, param) + for name, param in self.named_parameters(recurse=False) + if not name.startswith("shared_") + ] + + def update_expert_map(self) -> None: + pass + + @property + def layer_id(self) -> int: + return extract_layer_index(self.prefix) + + def _synchronize_first_launch(self) -> None: + ep_group = get_ep_group() + key = (id(ep_group.cpu_group), torch.accelerator.current_device_index()) + if key in self._synchronized_ep_groups: + return + torch.accelerator.synchronize() + torch.distributed.barrier(group=ep_group.cpu_group) + self._synchronized_ep_groups.add(key) + + def forward( self, - *, - y: torch.Tensor, - l1_weights: tuple[torch.Tensor, torch.Tensor], - l2_weights: tuple[torch.Tensor, torch.Tensor], - symm_buffer, - activation_clamp: float | None, - fast_math: bool, - activation: str | None = None, - activation_alpha: float | None = None, - activation_beta: float | None = None, - shared_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None, - shared_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None, - ) -> None: + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_experts: nn.Module | None = None, + ) -> torch.Tensor: + """Run the routed and shared experts on this rank's tokens. + + Args: + hidden_states: BF16 ``[num_tokens, hidden_size]`` activations. + topk_weights: ``[num_tokens, top_k]`` routing weights. + topk_ids: ``[num_tokens, top_k]`` logical expert ids. + shared_experts: Unfused shared MLP, run on ``hidden_states`` and + added to the output. Must be ``None`` when this layer fused the + shared expert at construction. + + Returns: + BF16 ``[num_tokens, hidden_size]`` routed plus shared expert output. + + """ + if shared_experts is not None and self.has_fused_shared_experts: + raise ValueError( + f"{self.prefix} fuses the shared expert into MegaMoE; " + "pass shared_experts=None." + ) + num_tokens = hidden_states.shape[0] + if num_tokens > self.max_num_tokens: + raise ValueError( + f"MegaMoE got {num_tokens} tokens, but the symmetric buffer was " + f"sized for {self.max_num_tokens}." + ) + self._synchronize_first_launch() + + is_padding = None + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): + is_padding = get_forward_context().is_padding + if is_padding is not None: + is_padding = is_padding[:num_tokens] + + if self.capture_fn is not None: + self.capture_fn(topk_ids) + + # EPLB: map logical expert IDs to physical replicas and record load. + eplb_state = self.eplb_state + if eplb_state.logical_to_physical_map is not None: + assert eplb_state.expert_load_view is not None + assert eplb_state.logical_replica_count is not None + assert eplb_state.should_record_tensor is not None + if is_padding is not None: + topk_ids = torch.where(is_padding.unsqueeze(1), -1, topk_ids) + topk_ids = eplb_map_to_physical_and_record( + topk_ids=topk_ids, + expert_load_view=eplb_state.expert_load_view, + logical_to_physical_map=eplb_state.logical_to_physical_map, + logical_replica_count=eplb_state.logical_replica_count, + record_enabled=eplb_state.should_record_tensor, + num_unpadded_tokens=eplb_state.num_unpadded_tokens_tensors[ + dbo_current_ubatch_id() + ] + if eplb_state.num_unpadded_tokens_tensors is not None + else None, + ) + from vllm.utils.deep_gemm import _import_deep_gemm deep_gemm = _import_deep_gemm() - kwargs: dict[str, Any] = { - name: value - for name, value in ( - ("activation", activation), - ("activation_alpha", activation_alpha), - ("activation_beta", activation_beta), + fmt = self.weight_format + symm_buffer = self.get_symm_buffer() + shared_x_sf = None + shared_block_m = None + if self.shared_spec is not None and self.shared_spec.scale_group_k: + shared_x_sf = symm_buffer.shared_l1_acts_sf + shared_block_m = self.get_block_m(num_tokens) + kwargs = dict(self._kernel_kwargs) + if fmt.uses_nvfp4_api: + l2_activation_scale = self.l2_activation_scale + prepare_megamoe_nvfp4_inputs( + hidden_states, + topk_weights, + topk_ids, + symm_buffer.x[:num_tokens], + symm_buffer.x_sf[:num_tokens], + symm_buffer.x_scales[:num_tokens], + symm_buffer.topk_idx[:num_tokens], + symm_buffer.topk_weights[:num_tokens], + is_padding=is_padding, + shared_x=symm_buffer.shared_l1_acts[:num_tokens] + if self.has_fused_shared_experts + else None, + shared_x_sf=shared_x_sf, + shared_block_m=shared_block_m, + w2_global_scale=self.w2_weight_scale_2, + l2_alpha=self._l2_alpha, + l2_activation_scale=l2_activation_scale, ) - if value is not None - } - if shared_l1_weights is not None and shared_l2_weights is not None: - kwargs["shared_l1_weights"] = shared_l1_weights - kwargs["shared_l2_weights"] = shared_l2_weights - deep_gemm.fp8_fp4_mega_moe( + kwargs.update( + l1_alpha=self.w13_weight_scale_2, + l2_alpha=self._l2_alpha, + l2_activation_scale=l2_activation_scale, + use_x_scales=True, + ) + else: + quantized = fmt.hidden_quant is not None + prepare_megamoe_inputs( + hidden_states, + topk_weights, + topk_ids, + symm_buffer.x[:num_tokens], + symm_buffer.x_sf[:num_tokens] if quantized else None, + symm_buffer.topk_idx[:num_tokens], + symm_buffer.topk_weights[:num_tokens], + is_padding=is_padding, + shared_x_sf=shared_x_sf, + shared_block_m=shared_block_m, + hidden_quant=fmt.hidden_quant, + ) + + if self.has_fused_shared_experts: + kwargs["shared_l1_weights"] = self._kernel_weights("shared_l1") + kwargs["shared_l2_weights"] = self._kernel_weights("shared_l2") + y = torch.empty_like(hidden_states, dtype=torch.bfloat16) + getattr(deep_gemm, fmt.kernel)( y, - l1_weights, - l2_weights, + self.l1_weights, + self.l2_weights, symm_buffer, - activation_clamp=activation_clamp, - fast_math=fast_math, + fast_math=self.fast_math, **kwargs, ) - - def supports_bf16_mega_gate(self) -> bool: - """SM100 deep_gemm build is expected to provide the bf16 gate kernel.""" - return True - - def bf16_mega_gate( - self, - *, - x: torch.Tensor, - weight: torch.Tensor, - num_topk: int, - scoring_func: str, - routed_scaling_factor: float, - ep_rank: int, - bias: torch.Tensor | None = None, - image_bias: torch.Tensor | None = None, - image_token_mask: torch.Tensor | None = None, - fix_routing_mask: torch.Tensor | None = None, - unmapped_topk_idx: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - from vllm.utils.deep_gemm import bf16_mega_gate - - return bf16_mega_gate( - x, - weight, - num_topk, - scoring_func=scoring_func, - routed_scaling_factor=routed_scaling_factor, - ep_rank=ep_rank, - bias=bias, - image_bias=image_bias, - image_token_mask=image_token_mask, - fix_routing_mask=fix_routing_mask, - unmapped_topk_idx=unmapped_topk_idx, - ) + if shared_experts is not None: + y += shared_experts(hidden_states) + return y -def get_deep_gemm_mega_moe_backend( - device: torch.device, - hidden_size: int, - intermediate_size: int, -) -> DeepGemmMegaMoEBackend: - """Return the DeepGEMM MegaMoE backend for ``device``. - - A successful return implies the backend can run on this device with the - given ``hidden_size`` / ``intermediate_size``. All runtime capability - checks (arch + shape divisibility) live here; there is no separate check - on the returned backend. - - Currently only the SM100 deep_gemm backend is registered here. SM90 and - other backends are expected to be provided by new backends or plugins - that override this factory. - - Raises: - NotImplementedError: If the device architecture is unsupported. - ValueError: If the shapes are unsupported. - - """ - if not current_platform.is_device_capability_family( - 100, device_id=device.index or 0 - ): - raise NotImplementedError("DeepGEMM MegaMoE requires SM100 GPUs.") - if hidden_size % 128 != 0 or intermediate_size % 128 != 0: - raise ValueError( - "DeepGEMM MegaMoE requires hidden and intermediate sizes " - "to be multiples of 128." - ) - return DeepGemmSm100MegaMoEBackend() +DeepGemmMegaMoEExperts.weight_loader.supports_moe_loading = True # type: ignore[attr-defined] diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py index 5423a93b5f..4706fb50b7 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py @@ -16,6 +16,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kMxfp4Static, ) +from vllm.triton_utils import triton __all__ = [ "AiterW4A16ExpertsMonolithic", @@ -52,10 +53,10 @@ def _aiter_w4a16_silu_via_a8w4( from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( - should_use_cdna4_mx_scale_swizzle, + get_aiter_mx_scale_swizzle, ) - swz = "CDNA4_SCALE" if should_use_cdna4_mx_scale_swizzle() else None + swz = get_aiter_mx_scale_swizzle() quant_dtype = torch.float8_e4m3fn g1_gammas = gammas if apply_router_weight_on_input else None @@ -134,6 +135,8 @@ def aiter_triton_kernel_w4a16_moe_forward( e_score_correction_bias: torch.Tensor | None = None, routed_scaling_factor: float | None = None, score_mode: str | None = None, + input_ids: torch.Tensor | None = None, + hash_indices_table: torch.Tensor | None = None, ): assert quant_config is not None and rocm_aiter_ops.is_enabled() from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( @@ -154,17 +157,37 @@ def aiter_triton_kernel_w4a16_moe_forward( _routing_mod.is_tdm_avail = lambda: False aiter_routing = _routing_mod.routing - if score_mode is not None: + scale = routed_scaling_factor if routed_scaling_factor is not None else 1.0 + if hash_indices_table is not None: + assert score_mode == "sqrtsoftplus" and input_ids is not None + # Same block_m heuristic as aiter's routing(). + n_tokens, n_experts = gating_output.shape + block_m = max( + 16, min(triton.next_power_of_2(max(1, n_tokens * topk // n_experts)), 128) + ) + routing_data, gather_idx, scatter_idx = _routing_mod.routing_from_hash( + gating_output, + hash_indices_table, + input_ids, + topk, + block_m, + score_mode=score_mode, + renorm=renormalize, + routed_scaling_factor=scale, + ) + elif score_mode is not None: use_grouped_topk = num_expert_group is not None and num_expert_group > 1 routing_data, gather_idx, scatter_idx = aiter_routing( gating_output, topk, score_mode=score_mode, - bias=e_score_correction_bias, - renorm=renormalize, - routed_scaling_factor=( - routed_scaling_factor if routed_scaling_factor is not None else 1.0 + bias=( + e_score_correction_bias.float() + if e_score_correction_bias is not None + else None ), + renorm=renormalize, + routed_scaling_factor=scale, use_grouped_topk=use_grouped_topk, num_expert_group=num_expert_group, topk_group=topk_group, @@ -348,6 +371,10 @@ class AiterW4A16ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): RoutingMethodType.DeepSeekV3, ] + @staticmethod + def _supports_hash_routing() -> bool: + return True + @staticmethod def _supports_router_logits_dtype( router_logits_dtype: torch.dtype | None, @@ -375,6 +402,8 @@ class AiterW4A16ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): routed_scaling_factor: float | None = None, topk_group: int | None = None, routing_replay_out: torch.Tensor | None = None, + input_ids: torch.Tensor | None = None, + hash_indices_table: torch.Tensor | None = None, ) -> torch.Tensor: routing_method = self.moe_config.routing_method if routing_method == RoutingMethodType.DeepseekV4: @@ -411,4 +440,6 @@ class AiterW4A16ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): e_score_correction_bias=e_score_correction_bias, routed_scaling_factor=routed_scaling_factor, score_mode=score_mode, + input_ids=input_ids, + hash_indices_table=hash_indices_table, ) diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a4_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a4_moe.py new file mode 100644 index 0000000000..a40d8184fb --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a4_moe.py @@ -0,0 +1,373 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm._aiter_ops import rocm_aiter_ops +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, + RoutingMethodType, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kMxfp4Dynamic, + kMxfp4Static, +) + +__all__ = [ + "AiterW4A4ExpertsMonolithic", + "aiter_triton_kernel_w4a4_moe_forward", +] + + +def _aiter_raw(t): + if t is None or isinstance(t, torch.Tensor): + return t + return t.storage.data if hasattr(t, "storage") else t + + +def aiter_triton_kernel_w4a4_moe_forward( + hidden_states: torch.Tensor, + w1, + w2, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + activation: MoEActivation = MoEActivation.SWIGLUOAI, + quant_config: FusedMoEQuantConfig | None = None, + apply_router_weight_on_input: bool = False, + global_num_experts: int = -1, + expert_map: torch.Tensor | None = None, + unpadded_N_w1=None, + unpadded_K_w1=None, + unpadded_N_w2=None, + unpadded_K_w2=None, + num_expert_group: int | None = None, + topk_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + score_mode: str | None = None, + input_ids: torch.Tensor | None = None, + hash_indices_table: torch.Tensor | None = None, +): + assert quant_config is not None and rocm_aiter_ops.is_enabled() + + from aiter.ops.triton.moe.moe_op_gemm_a4w4 import moe_gemm_a4w4, mxfp4_quant + + try: + from aiter.ops.triton.moe.moe_routing import routing as _routing_mod + except ImportError: + from aiter.ops.triton.moe_routing import routing as _routing_mod + + from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + get_aiter_mx_scale_swizzle, + ) + from vllm.platforms.rocm import on_gfx1250 + + if on_gfx1250(): + _routing_mod.is_tdm_avail = lambda: False + aiter_routing = _routing_mod.routing + + gating_output = torch.nan_to_num(gating_output, nan=0.0, posinf=0.0, neginf=0.0) + + if hash_indices_table is not None: + assert input_ids is not None, "hash routing requires input_ids" + n_tokens, n_expts_tot = gating_output.shape + tokens_per_expt = max(1, n_tokens * topk // n_expts_tot) + block_m = max(16, min(1 << (tokens_per_expt - 1).bit_length(), 128)) + routing_data, gather_idx, scatter_idx = _routing_mod.routing_from_hash( + gating_output, + hash_indices_table, + input_ids.to(hash_indices_table.dtype), + topk, + block_m, + score_mode=score_mode or "sqrtsoftplus", + renorm=renormalize, + routed_scaling_factor=( + routed_scaling_factor if routed_scaling_factor is not None else 1.0 + ), + ) + elif score_mode is not None: + use_grouped_topk = num_expert_group is not None and num_expert_group > 1 + routing_data, gather_idx, scatter_idx = aiter_routing( + gating_output, + topk, + score_mode=score_mode, + bias=( + e_score_correction_bias.float() + if e_score_correction_bias is not None + else None + ), + renorm=renormalize, + routed_scaling_factor=( + routed_scaling_factor if routed_scaling_factor is not None else 1.0 + ), + use_grouped_topk=use_grouped_topk, + num_expert_group=num_expert_group, + topk_group=topk_group, + ) + else: + routing_data, gather_idx, scatter_idx = aiter_routing( + gating_output, topk, sm_first=not renormalize + ) + + if on_gfx1250(): + gather_src = gather_idx.to(torch.long) // topk + hidden_states = hidden_states[gather_src] + gather_idx = None + + w1_data = _aiter_raw(w1) + w2_data = _aiter_raw(w2) + w1_scale = _aiter_raw(quant_config.w1_scale) + w2_scale = _aiter_raw(quant_config.w2_scale) + + gammas = routing_data.gate_scal if routing_data else None + + swiglu_limit = ( + quant_config.gemm1_clamp_limit + if quant_config.gemm1_clamp_limit is not None + else 7.0 + ) + # SWIGLUOAI needs the alpha/residual swiglu fused into the GEMM1 epilogue, + # which reads gate/up interleaved along N (see the oracle's weight prep). + fused_swiglu = activation in ( + MoEActivation.SWIGLUOAI, + MoEActivation.SWIGLUOAI_UNINTERLEAVE, + ) + swiglu_kwargs: dict[str, object] = {} + if fused_swiglu: + assert quant_config.gemm1_beta in (None, 1.0), ( + "aiter's fused swiglu hardcodes the residual to (up + 1); " + f"gemm1_beta={quant_config.gemm1_beta} cannot be expressed" + ) + swiglu_kwargs = { + "alpha": ( + quant_config.gemm1_alpha + if quant_config.gemm1_alpha is not None + else 1.0 + ), + "limit": swiglu_limit, + "swiglu_add_residual": quant_config.gemm1_beta == 1.0, + } + swizzle_mx_scale = get_aiter_mx_scale_swizzle() + + x_q, x_scale = mxfp4_quant(hidden_states.to(torch.bfloat16)) + + # GEMM1: gate+up projection. + raw_intermediate = moe_gemm_a4w4( + x_q, + w1_data, + x_scale, + w1_scale, + bias=quant_config.w1_bias, + routing_data=routing_data, + gather_indx=gather_idx, + gammas=gammas if apply_router_weight_on_input else None, + swizzle_mx_scale=swizzle_mx_scale, + apply_swiglu=fused_swiglu, + **swiglu_kwargs, + ) + + if fused_swiglu: + # The swiglu epilogue already halved N. + intermediate = ( + raw_intermediate + if unpadded_N_w1 is None + else raw_intermediate[:, : unpadded_N_w1 // 2] + ) + else: + if unpadded_N_w1 is not None: + raw_intermediate = raw_intermediate[:, :unpadded_N_w1] + + # SiLU(gate) * up on the concatenated [gate | up] halves + from aiter.ops.triton.fusions.fused_clamp_act_mul import fused_clamp_act_mul + + half_n = raw_intermediate.shape[-1] // 2 + intermediate = torch.empty( + raw_intermediate.shape[0], + half_n, + dtype=raw_intermediate.dtype, + device=raw_intermediate.device, + ) + fused_clamp_act_mul( + raw_intermediate, + out=intermediate, + swiglu_limit=swiglu_limit, + activation="silu", + dtype_quant=None, + ) + + # GEMM2: down projection with scatter-reduce + mid_q, mid_scale = mxfp4_quant(intermediate.to(torch.bfloat16)) + + out = moe_gemm_a4w4( + mid_q, + w2_data, + mid_scale, + w2_scale, + bias=quant_config.w2_bias, + routing_data=routing_data, + scatter_indx=scatter_idx, + gammas=None if apply_router_weight_on_input else gammas, + swizzle_mx_scale=swizzle_mx_scale, + apply_swiglu=False, + ) + + return out + + +class AiterW4A4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): + """Monolithic MXFP4 W4A4 expert using AITER Triton/Gluon kernels (gfx1250). + + Uses moe_gemm_a4w4 (auto-selects gluon backend on gfx1250) with dynamic + MXFP4 activation quantization via mxfp4_quant. + """ + + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + ): + super().__init__(moe_config, quant_config) + self.topk = moe_config.experts_per_token + self.renormalize = moe_config.routing_method in ( + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + RoutingMethodType.DeepseekV4, + RoutingMethodType.DeepSeekV3, + RoutingMethodType.MiniMax2, + ) + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @staticmethod + def _supports_hash_routing() -> bool: + return True + + @staticmethod + def _supports_current_device() -> bool: + if not rocm_aiter_ops.is_enabled(): + return False + from vllm.platforms.rocm import on_gfx1250 + + return on_gfx1250() + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return (weight_key, activation_key) == (kMxfp4Static, kMxfp4Dynamic) + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation in ( + MoEActivation.SWIGLUOAI, + MoEActivation.SWIGLUOAI_UNINTERLEAVE, + MoEActivation.SILU, + ) + + @staticmethod + def _supports_parallel_config( + moe_parallel_config: FusedMoEParallelConfig, + ) -> bool: + return ( + not moe_parallel_config.use_all2all_kernels + and not moe_parallel_config.enable_eplb + and moe_parallel_config.dp_size <= 1 + ) + + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return routing_method in [ + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + RoutingMethodType.DeepseekV4, + RoutingMethodType.DeepSeekV3, + RoutingMethodType.MiniMax2, + ] + + @staticmethod + def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, + ) -> bool: + return True + + @property + def expects_unquantized_inputs(self) -> bool: + return True + + def apply( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + apply_router_weight_on_input: bool, + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, + routing_replay_out: torch.Tensor | None = None, + input_ids: torch.Tensor | None = None, + hash_indices_table: torch.Tensor | None = None, + ) -> torch.Tensor: + assert self.moe_config.intermediate_size_per_partition_unpadded is not None + assert self.moe_config.hidden_dim_unpadded is not None + routing_method = self.moe_config.routing_method + if routing_method == RoutingMethodType.DeepseekV4: + score_mode = "sqrtsoftplus" + elif routing_method in ( + RoutingMethodType.DeepSeekV3, + RoutingMethodType.MiniMax2, + ): + # Sigmoid scores + routing bias; grouped for DeepSeek-R1/V3, + # ungrouped for MiniMax-M3. + score_mode = "sigmoid" + else: + score_mode = None + return aiter_triton_kernel_w4a4_moe_forward( + hidden_states=hidden_states, + w1=w1, + w2=w2, + gating_output=router_logits, + topk=self.topk, + renormalize=self.renormalize, + activation=activation, + global_num_experts=global_num_experts, + expert_map=expert_map, + quant_config=self.quant_config, + apply_router_weight_on_input=apply_router_weight_on_input, + unpadded_N_w1=( + self.moe_config.intermediate_size_per_partition_unpadded * 2 + ), + unpadded_K_w1=self.moe_config.hidden_dim_unpadded, + unpadded_N_w2=self.moe_config.hidden_dim_unpadded, + unpadded_K_w2=(self.moe_config.intermediate_size_per_partition_unpadded), + num_expert_group=num_expert_group, + topk_group=topk_group, + e_score_correction_bias=e_score_correction_bias, + routed_scaling_factor=routed_scaling_factor, + score_mode=score_mode, + input_ids=input_ids, + hash_indices_table=hash_indices_table, + ) diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py index bc1a4af6c0..f7d2be94aa 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py @@ -61,12 +61,6 @@ def aiter_triton_kernel_w4a8_moe_forward( gating_output, topk, sm_first=not renormalize ) - # gfx1250: aiter's in-kernel gather is numerically broken - if on_gfx1250(): - gather_src = gather_idx.to(torch.long) // topk - hidden_states = hidden_states[gather_src] - gather_idx = None - return triton_kernel_fused_mxfp4_w4a8_experts( None, hidden_states, @@ -130,11 +124,11 @@ def triton_kernel_fused_mxfp4_w4a8_experts( from aiter.ops.triton.quant_moe import downcast_to_static_fp8 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( - should_use_cdna4_mx_scale_swizzle, + get_aiter_mx_scale_swizzle, weight_mx_scale, ) - _swizzle_mx_scale = "CDNA4_SCALE" if should_use_cdna4_mx_scale_swizzle() else None + _swizzle_mx_scale = get_aiter_mx_scale_swizzle() assert quant_config.w1_precision is not None, ( "w1_precision in quant config can't be None" diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 839eba5d78..a57bb3472b 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -436,6 +436,7 @@ def FusedMoEFactory( # TODO get from router? needs to be truncated? e_score_correction_bias=e_score_correction_bias, apply_router_weight_on_input=apply_router_weight_on_input, + hash_indices_table=hash_indices_table, **routed_experts_args if routed_experts_args is not None else {}, ) diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 785369a344..fccbd1a9d8 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -5,7 +5,7 @@ from collections.abc import Callable from dataclasses import dataclass from enum import Enum from math import prod -from typing import final +from typing import Any, final import torch @@ -579,7 +579,11 @@ class FusedMoEExperts(ABC): return False, _make_reason( f"parallel config {moe_config.moe_parallel_config}" ) - elif moe_config.has_hash_routing and cls.is_monolithic(): + elif ( + moe_config.has_hash_routing + and cls.is_monolithic() + and not cls._supports_hash_routing() + ): return False, _make_reason("hash routing") elif moe_config.require_decomposed_backend and cls.is_monolithic(): return False, _make_reason("routes adjusted after routing") @@ -659,6 +663,15 @@ class FusedMoEExperts(ABC): """ return True + @staticmethod + def _supports_hash_routing() -> bool: + """Whether a monolithic kernel implements hash routing itself. + + Kernels that return True receive `input_ids` and `hash_indices_table` + in `apply()` for layers with a hash-routing table. + """ + return False + @staticmethod def _supports_router_logits_dtype( router_logits_dtype: torch.dtype | None, @@ -1584,11 +1597,20 @@ class FusedMoEKernelMonolithicImpl: topk_group: int | None = None, *, routing_sink: RoutedExpertsSink | None, + input_ids: torch.Tensor | None = None, + hash_indices_table: torch.Tensor | None = None, ) -> torch.Tensor | UnfinalizedMoEOutput: """Same as forward(), except uses router_logits as opposed to the topk_ids and topk_weights. This is used for kernels that have fused router + experts (e.g. FLASHINFER_TRTLLM). """ + # Only experts with _supports_hash_routing() accept these. + hash_kwargs: dict[str, Any] = {} + if hash_indices_table is not None: + hash_kwargs = { + "input_ids": input_ids, + "hash_indices_table": hash_indices_table, + } a1q, a1q_scale, router_logits = self.prepare_finalize.prepare( hidden_states, router_logits=router_logits, @@ -1605,6 +1627,11 @@ class FusedMoEKernelMonolithicImpl: f"kernel {type(self.fused_experts).__name__}." ) routing_replay_out = routing_sink.buffer[: len(a1q)] + hash_kwargs = ( + {"input_ids": input_ids, "hash_indices_table": hash_indices_table} + if hash_indices_table is not None + else {} + ) fused_out = self.fused_experts.apply( hidden_states=a1q, w1=w1, @@ -1621,6 +1648,7 @@ class FusedMoEKernelMonolithicImpl: routed_scaling_factor=routed_scaling_factor, topk_group=topk_group, routing_replay_out=routing_replay_out, + **hash_kwargs, ) if routing_sink is not None: routing_sink.capture_fn(routing_replay_out) @@ -1733,6 +1761,8 @@ class FusedMoEKernel: topk_group: int | None = None, *, routing_sink: RoutedExpertsSink | None, + input_ids: torch.Tensor | None = None, + hash_indices_table: torch.Tensor | None = None, ) -> torch.Tensor | UnfinalizedMoEOutput: assert isinstance(self.impl, FusedMoEKernelMonolithicImpl) return self.impl.apply( @@ -1749,6 +1779,8 @@ class FusedMoEKernel: routed_scaling_factor=routed_scaling_factor, topk_group=topk_group, routing_sink=routing_sink, + input_ids=input_ids, + hash_indices_table=hash_indices_table, ) def apply( diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index aef5c5b5b5..07e8815fc7 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -841,6 +841,11 @@ def make_fp8_moe_quant_config( block_shape=block_shape, per_act_token_quant=per_act_token_quant, per_out_ch_quant=per_out_ch_quant, + # TRTLLM fp8-block consumes activation scales as [K/128, M]; emit them + # transposed from the quant kernel instead of copying afterwards. + is_scale_transposed=( + fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM and block_shape == [128, 128] + ), gemm1_alpha=gemm1_alpha, gemm1_beta=gemm1_beta, gemm1_clamp_limit=swiglu_limit, diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 3abc0cbede..ae79298eea 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -152,7 +152,9 @@ class Mxfp4MoeBackend(Enum): # W4A16: aiter Triton moe_gemm_a16w4 kernel (gfx942/gfx950/gfx1250) AITER_TRITON_MXFP4_BF16 = "AITER_TRITON_MXFP4_BF16" AITER_MXFP4_FP8 = "AITER_MXFP4_FP8" # W4A8: triton kernel - AITER_MXFP4_MXFP4 = "AITER_MXFP4_MXFP4" # W4A4: CK kernel + AITER_MXFP4_MXFP4 = "AITER_MXFP4_MXFP4" # W4A4: CK kernel (gfx950) + # W4A4: aiter Triton/Gluon moe_gemm_a4w4 kernel (gfx1250) + AITER_TRITON_MXFP4_MXFP4 = "AITER_TRITON_MXFP4_MXFP4" # Triton TRITON = "TRITON" TRITON_UNFUSED = "TRITON_UNFUSED" @@ -310,6 +312,13 @@ def backend_to_kernel_cls( return [AiterMxfp4Experts] + elif backend == Mxfp4MoeBackend.AITER_TRITON_MXFP4_MXFP4: + from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a4_moe import ( + AiterW4A4ExpertsMonolithic, + ) + + return [AiterW4A4ExpertsMonolithic] + elif backend == Mxfp4MoeBackend.XPU: from vllm.model_executor.layers.fused_moe.experts.xpu_moe import XPUExpertsMxFp4 @@ -370,9 +379,11 @@ def map_mxfp4_backend(runner_backend: MoEBackend) -> list[Mxfp4MoeBackend]: Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16, Mxfp4MoeBackend.AITER_MXFP4_FP8, Mxfp4MoeBackend.AITER_MXFP4_MXFP4, + Mxfp4MoeBackend.AITER_TRITON_MXFP4_MXFP4, ], "aiter_triton_mxfp4_bf16": [Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16], "aiter_mxfp4_fp8": [Mxfp4MoeBackend.AITER_MXFP4_FP8], + "aiter_triton_mxfp4_mxfp4": [Mxfp4MoeBackend.AITER_TRITON_MXFP4_MXFP4], "aiter_mxfp4_mxfp4": [Mxfp4MoeBackend.AITER_MXFP4_MXFP4], "xpu": [Mxfp4MoeBackend.XPU], "cpu": [Mxfp4MoeBackend.CPU], @@ -464,6 +475,7 @@ def _backend_activation_key(backend: Mxfp4MoeBackend) -> QuantKey | None: return kFp8StaticTensorSym if backend in ( Mxfp4MoeBackend.AITER_MXFP4_MXFP4, + Mxfp4MoeBackend.AITER_TRITON_MXFP4_MXFP4, Mxfp4MoeBackend.CUTLASS_MXFP4_MXFP4, ): return kMxfp4Dynamic @@ -896,6 +908,18 @@ def mxfp4_round_up_hidden_size_and_intermediate_size( return hidden_size, intermediate_size +def _interleave_gate_up_rows(t: torch.Tensor) -> torch.Tensor: + """Repack [E, 2n, X] from [all gates | all ups] to [g0, u0, g1, u1, ...].""" + num_experts, two_n, inner = t.shape + assert two_n % 2 == 0, f"expected an even gate/up row count, got {two_n}" + return ( + t.view(num_experts, 2, two_n // 2, inner) + .transpose(1, 2) + .contiguous() + .view(num_experts, two_n, inner) + ) + + def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( mxfp4_backend: Mxfp4MoeBackend, layer: torch.nn.Module, @@ -1192,9 +1216,56 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( w2_bias, ) - elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4: - from vllm._aiter_ops import rocm_aiter_ops + elif mxfp4_backend == Mxfp4MoeBackend.AITER_TRITON_MXFP4_MXFP4: + if w13_bias is not None: + w13_bias = w13_bias.data.to(torch.float32) + if w2_bias is not None: + w2_bias = w2_bias.data.to(torch.float32) + # Triton moe_gemm_a4w4: uint8 column-major weights. + w13_data = w13_weight.data.view(torch.uint8) + w2_data = w2_weight.data.view(torch.uint8) + w13_scale_data = w13_weight_scale.data.view(torch.uint8) + + # SILU applies the activation outside the GEMM and keeps the + # checkpoint's packed halves; the SWIGLUOAI epilogue needs them + # interleaved. + if getattr(layer, "activation", None) in ( + MoEActivation.SWIGLUOAI, + MoEActivation.SWIGLUOAI_UNINTERLEAVE, + ): + w13_data = _interleave_gate_up_rows(w13_data) + w13_scale_data = _interleave_gate_up_rows(w13_scale_data) + + # Transpose to [E, K, N] — the view has stride(-2)==1 (column-major), + # which is what moe_gemm_a4w4 asserts. + w13_data = w13_data.transpose(1, 2) + w2_data = w2_data.transpose(1, 2) + + # Scales are [E, N, K_scale]; both kernels index them as + # [E, K_scale, N] with K innermost, so transpose(1, 2) + # Keep them uint8: Triton moe_gemm_a4w4 takes e8m0 scales + w13_scale = w13_scale_data.transpose(1, 2) + w2_scale = w2_weight_scale.data.view(torch.uint8).transpose(1, 2) + + from aiter.ops.triton.utils.shuffle import shuffle_scale_moe + + w13_scale = shuffle_scale_moe( + w13_scale, + arch="gfx1250", + preshuffle_factor=32, + scale_kwidth=4, + ) + w2_scale = shuffle_scale_moe( + w2_scale, + arch="gfx1250", + preshuffle_factor=32, + scale_kwidth=4, + ) + + return (w13_data, w2_data, w13_scale, w2_scale, w13_bias, w2_bias) + + elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4: if w13_bias is not None: w13_bias = w13_bias.data.to(torch.float32) if w2_bias is not None: @@ -1203,6 +1274,8 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( # e8m0_shuffle on weight scales (GFX950 swizzle layout) from aiter.utility.fp4_utils import e8m0_shuffle + from vllm._aiter_ops import rocm_aiter_ops + s0, s1, _ = w13_weight_scale.shape w13_weight_scale.data = e8m0_shuffle(w13_weight_scale.view(s0 * s1, -1)).view( s0, s1, -1 @@ -1935,6 +2008,7 @@ def convert_weight_to_mxfp4_moe_kernel_format( ) elif mxfp4_backend in ( Mxfp4MoeBackend.AITER_MXFP4_MXFP4, + Mxfp4MoeBackend.AITER_TRITON_MXFP4_MXFP4, Mxfp4MoeBackend.AITER_MXFP4_FP8, ): return convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( @@ -2088,6 +2162,7 @@ def make_mxfp4_moe_quant_config( ) elif mxfp4_backend in ( Mxfp4MoeBackend.AITER_MXFP4_MXFP4, + Mxfp4MoeBackend.AITER_TRITON_MXFP4_MXFP4, Mxfp4MoeBackend.CUTLASS_MXFP4_MXFP4, ): return ocp_mx_moe_quant_config( diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py index 660807fea3..fe5d518cf3 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py @@ -31,6 +31,7 @@ def _quantize_input( per_act_token_quant=quant_config.per_act_token_quant, block_shape=quant_config.block_shape, is_scale_swizzled=quant_config.is_scale_swizzled, + is_scale_transposed=quant_config.is_scale_transposed, mx_alignment=quant_config.mx_alignment, ) diff --git a/vllm/model_executor/layers/fused_moe/prepare_megamoe.py b/vllm/model_executor/layers/fused_moe/prepare_megamoe.py new file mode 100644 index 0000000000..5e8d8cde9a --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/prepare_megamoe.py @@ -0,0 +1,682 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Triton input-staging kernels for DeepGEMM MegaMoE. + +Quantize hidden states to fp8 with E8M0 group scales, to NVFP4 with per-token +global scales, or copy them as BF16, and repack the routing top-k tensors into +the int64/float32 layout that the DeepGEMM MegaMoE kernels consume. +""" + +import torch + +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + FP8_DTYPE, + QuantKey, + kMxfp8Dynamic, +) +from vllm.triton_utils import tl, triton + +_E8M0_SCALE_DTYPES = (torch.uint8, torch.float8_e8m0fnu) +# Inverse of the NVFP4 range: E2M1 max (6) times E4M3 max (448). +_INV_NVFP4_MAX = tl.constexpr(1.0 / 2688.0) + + +@triton.jit +def _quantize_fp8_groups( + hidden, + BLOCK_M: tl.constexpr, + BLOCK_K: tl.constexpr, + GROUP_K: tl.constexpr, + USE_UE8M0: tl.constexpr, +): + """Return fp8 values, fp32 group scales and, for UE8M0, the scale exponents.""" + num_groups: tl.constexpr = BLOCK_K // GROUP_K + hidden_groups = tl.reshape(tl.abs(hidden), [BLOCK_M, num_groups, GROUP_K]) + amax = tl.max(hidden_groups, axis=2) + amax = tl.maximum(amax, 1.0e-4) + + scale = amax / 448.0 + if USE_UE8M0: + scale_bits = scale.to(tl.uint32, bitcast=True) + scale_exp = ((scale_bits >> 23) & 0xFF) + ((scale_bits & 0x7FFFFF) != 0).to( + tl.uint32 + ) + scale_exp = tl.minimum(tl.maximum(scale_exp, 1), 254) + quant_scale = (scale_exp << 23).to(tl.float32, bitcast=True) + else: + scale_exp = scale.to(tl.uint32, bitcast=True) + quant_scale = scale + + hidden_groups = tl.reshape(hidden, [BLOCK_M, num_groups, GROUP_K]) + scaled = hidden_groups * (1.0 / quant_scale)[:, :, None] + scaled = tl.reshape(scaled, [BLOCK_M, BLOCK_K]) + return scaled.to(tl.float8e4nv), quant_scale, scale_exp + + +@triton.jit +def _pack_bytes(values, NUM_BYTES: tl.constexpr): + """Pack the last dim of ``[..., NUM_BYTES]`` byte values into int32 words.""" + offsets = tl.arange(0, NUM_BYTES) + return tl.sum(values.to(tl.uint32) << (offsets * 8), axis=-1).to(tl.int32) + + +@triton.jit +def _shared_sf_row(token_id, SHARED_BLOCK_M: tl.constexpr): + """Row of ``token_id`` in the shared-expert MN-major UTCCP scale pages.""" + m_block_id = token_id // SHARED_BLOCK_M + m_in_block = token_id % SHARED_BLOCK_M + aligned_block_m: tl.constexpr = triton.cdiv(SHARED_BLOCK_M, 128) * 128 + transposed_m = ( + (m_in_block // 128) * 128 + (m_in_block % 32) * 4 + (m_in_block % 128) // 32 + ) + return m_block_id * aligned_block_m + transposed_m + + +@triton.jit +def _store_topk( + token_id, + token_mask, + topk_ids, + topk_weights, + is_padding, + topk_idx_out, + topk_weights_out, + topk_ids_stride_m: tl.constexpr, + topk_ids_stride_k: tl.constexpr, + topk_weights_stride_m: tl.constexpr, + topk_weights_stride_k: tl.constexpr, + is_padding_stride_m: tl.constexpr, + topk_idx_stride_m: tl.constexpr, + topk_idx_stride_k: tl.constexpr, + topk_weights_out_stride_m: tl.constexpr, + topk_weights_out_stride_k: tl.constexpr, + top_k: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_TOPK: tl.constexpr, +) -> None: + """Repack top-k routing as int64 ids and fp32 weights; padding is unrouted.""" + topk_offsets = tl.arange(0, BLOCK_TOPK) + topk_mask = token_mask[:, None] & (topk_offsets[None, :] < top_k) + token_is_padding = tl.full((BLOCK_M,), False, tl.int1) + if is_padding is not None: + token_is_padding = tl.load( + is_padding + token_id * is_padding_stride_m, + mask=token_mask, + other=True, + ) + + ids = tl.load( + topk_ids + + token_id[:, None] * topk_ids_stride_m + + topk_offsets[None, :] * topk_ids_stride_k, + mask=topk_mask, + other=0, + ).to(tl.int64) + ids = tl.where(token_is_padding[:, None], -1, ids) + tl.store( + topk_idx_out + + token_id[:, None] * topk_idx_stride_m + + topk_offsets[None, :] * topk_idx_stride_k, + ids, + mask=topk_mask, + ) + + weights = tl.load( + topk_weights + + token_id[:, None] * topk_weights_stride_m + + topk_offsets[None, :] * topk_weights_stride_k, + mask=topk_mask, + other=0.0, + ) + weights = tl.where(token_is_padding[:, None], 0.0, weights) + tl.store( + topk_weights_out + + token_id[:, None] * topk_weights_out_stride_m + + topk_offsets[None, :] * topk_weights_out_stride_k, + weights, + mask=topk_mask, + ) + + +@triton.jit(do_not_specialize=["num_tokens"]) +def _prepare_megamoe_inputs_kernel( + hidden_states, + x_out, + x_sf, + shared_x_sf, + topk_ids, + topk_weights, + is_padding, + topk_idx_out, + topk_weights_out, + hidden_stride_m: tl.constexpr, + hidden_stride_k: tl.constexpr, + x_stride_m: tl.constexpr, + x_stride_k: tl.constexpr, + x_sf_stride_m: tl.constexpr, + x_sf_stride_k: tl.constexpr, + shared_x_sf_stride_m: tl.constexpr, + shared_x_sf_stride_k: tl.constexpr, + topk_ids_stride_m: tl.constexpr, + topk_ids_stride_k: tl.constexpr, + topk_weights_stride_m: tl.constexpr, + topk_weights_stride_k: tl.constexpr, + is_padding_stride_m: tl.constexpr, + topk_idx_stride_m: tl.constexpr, + topk_idx_stride_k: tl.constexpr, + topk_weights_out_stride_m: tl.constexpr, + topk_weights_out_stride_k: tl.constexpr, + num_tokens, + hidden_size: tl.constexpr, + top_k: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_K: tl.constexpr, + GROUP_K: tl.constexpr, + BLOCK_TOPK: tl.constexpr, + SHARED_BLOCK_M: tl.constexpr, + QUANTIZE: tl.constexpr, + USE_UE8M0: tl.constexpr, +) -> None: + token_id = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) + token_mask = token_id < num_tokens + k_block_id = tl.program_id(1) + + k_offsets = k_block_id * BLOCK_K + tl.arange(0, BLOCK_K) + k_mask = token_mask[:, None] & (k_offsets[None, :] < hidden_size) + hidden = tl.load( + hidden_states + + token_id[:, None] * hidden_stride_m + + k_offsets[None, :] * hidden_stride_k, + mask=k_mask, + other=0.0, + ).to(tl.float32) + + if not QUANTIZE: + tl.store( + x_out + token_id[:, None] * x_stride_m + k_offsets[None, :] * x_stride_k, + hidden.to(x_out.dtype.element_ty), + mask=k_mask, + ) + else: + num_groups: tl.constexpr = BLOCK_K // GROUP_K + fp8, quant_scale, scale_exp = _quantize_fp8_groups( + hidden, BLOCK_M, BLOCK_K, GROUP_K, USE_UE8M0 + ) + tl.store( + x_out + token_id[:, None] * x_stride_m + k_offsets[None, :] * x_stride_k, + fp8, + mask=k_mask, + ) + + if USE_UE8M0: + packed_scale = _pack_bytes(scale_exp, num_groups) + tl.store( + x_sf + token_id * x_sf_stride_m + k_block_id * x_sf_stride_k, + packed_scale, + mask=token_mask, + ) + else: + sf_offset = tl.arange(0, num_groups) + tl.store( + x_sf + + token_id[:, None] * x_sf_stride_m + + k_block_id * x_sf_stride_k + + sf_offset[None, :], + quant_scale, + mask=token_mask[:, None], + ) + + # DeepGEMM's SM100 shared-expert TMA loads require the activation scales + # in an MN-major layout whose row permutation depends on the MegaMoE + # scheduler's runtime BLOCK_M. Write that view while the packed UE8M0 scale + # is already resident, avoiding another kernel and temporary tensor. + if USE_UE8M0 and shared_x_sf is not None: + shared_row = _shared_sf_row(token_id, SHARED_BLOCK_M) + tl.store( + shared_x_sf + + shared_row * shared_x_sf_stride_m + + k_block_id * shared_x_sf_stride_k, + packed_scale, + mask=token_mask, + ) + + if k_block_id == 0: + _store_topk( + token_id, + token_mask, + topk_ids, + topk_weights, + is_padding, + topk_idx_out, + topk_weights_out, + topk_ids_stride_m, + topk_ids_stride_k, + topk_weights_stride_m, + topk_weights_stride_k, + is_padding_stride_m, + topk_idx_stride_m, + topk_idx_stride_k, + topk_weights_out_stride_m, + topk_weights_out_stride_k, + top_k, + BLOCK_M, + BLOCK_TOPK, + ) + + +@triton.jit(do_not_specialize=["num_tokens", "l2_activation_scale"]) +def _prepare_megamoe_nvfp4_inputs_kernel( + hidden_states, + x_out, + x_sf, + x_scales, + shared_x, + shared_x_sf, + topk_ids, + topk_weights, + is_padding, + topk_idx_out, + topk_weights_out, + w2_global_scale, + l2_alpha, + hidden_stride_m: tl.constexpr, + hidden_stride_k: tl.constexpr, + x_stride_m: tl.constexpr, + x_sf_stride_m: tl.constexpr, + shared_x_stride_m: tl.constexpr, + shared_x_sf_stride_m: tl.constexpr, + shared_x_sf_stride_k: tl.constexpr, + topk_ids_stride_m: tl.constexpr, + topk_ids_stride_k: tl.constexpr, + topk_weights_stride_m: tl.constexpr, + topk_weights_stride_k: tl.constexpr, + is_padding_stride_m: tl.constexpr, + topk_idx_stride_m: tl.constexpr, + topk_idx_stride_k: tl.constexpr, + topk_weights_out_stride_m: tl.constexpr, + topk_weights_out_stride_k: tl.constexpr, + num_tokens, + l2_activation_scale, + hidden_size: tl.constexpr, + top_k: tl.constexpr, + num_local_experts: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_TOPK: tl.constexpr, + BLOCK_E: tl.constexpr, + SHARED_BLOCK_M: tl.constexpr, + SHARED_MXFP8: tl.constexpr, +) -> None: + if l2_alpha is not None: + experts = tl.arange(0, BLOCK_E) + expert_mask = (experts < num_local_experts) & (tl.program_id(0) == 0) + alpha = tl.load(w2_global_scale + experts, mask=expert_mask) + tl.store(l2_alpha + experts, alpha * l2_activation_scale, mask=expert_mask) + + token_id = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) + token_mask = token_id < num_tokens + k_offsets = tl.arange(0, BLOCK_K) + hidden_ptrs = hidden_states + token_id[:, None] * hidden_stride_m + + amax = tl.zeros((BLOCK_M,), tl.float32) + for k_block_id in range(hidden_size // BLOCK_K): + hidden = tl.load( + hidden_ptrs + (k_block_id * BLOCK_K + k_offsets[None, :]) * hidden_stride_k, + mask=token_mask[:, None], + other=0.0, + ).to(tl.float32) + amax = tl.maximum(amax, tl.max(tl.abs(hidden), axis=1)) + x_scale = tl.maximum(amax, 1.0e-12) * _INV_NVFP4_MAX + tl.store(x_scales + token_id, x_scale, mask=token_mask) + + for k_block_id in range(hidden_size // BLOCK_K): + k_ids = k_block_id * BLOCK_K + k_offsets + hidden = tl.load( + hidden_ptrs + k_ids[None, :] * hidden_stride_k, + mask=token_mask[:, None], + other=0.0, + ) + + # E2M1 values with one E4M3 scale per 16, relative to the token scale. + values = tl.reshape( + tl.div_rn(hidden.to(tl.float32), x_scale[:, None]), + [BLOCK_M, BLOCK_K // 16, 16], + ) + block_scale = tl.max(tl.abs(values), axis=2) * (1.0 / 6.0) + block_scale = tl.minimum(tl.maximum(block_scale, 2.0**-9), 448.0) + block_scale = block_scale.to(tl.float8e4nv) + values = tl.div_rn(values, block_scale.to(tl.float32)[:, :, None]) + magnitude = tl.abs(values) + # Round to nearest even on the grid {0, 0.5, 1, 1.5, 2, 3, 4, 6}. + code = ( + (magnitude > 0.25).to(tl.int32) + + (magnitude >= 0.75).to(tl.int32) + + (magnitude > 1.25).to(tl.int32) + + (magnitude >= 1.75).to(tl.int32) + + (magnitude > 2.5).to(tl.int32) + + (magnitude >= 3.5).to(tl.int32) + + (magnitude > 5.0).to(tl.int32) + ) + code = tl.where((values < 0) & (code != 0), code | 8, code) + low, high = tl.split(tl.reshape(code, [BLOCK_M, BLOCK_K // 2, 2])) + packed_offsets = k_block_id * (BLOCK_K // 2) + tl.arange(0, BLOCK_K // 2) + tl.store( + x_out + token_id[:, None] * x_stride_m + packed_offsets[None, :], + (low | (high << 4)).to(tl.uint8), + mask=token_mask[:, None], + ) + sf_offsets = k_block_id * (BLOCK_K // 64) + tl.arange(0, BLOCK_K // 64) + sf_bytes = tl.reshape( + block_scale.to(tl.uint8, bitcast=True), [BLOCK_M, BLOCK_K // 64, 4] + ) + tl.store( + x_sf + token_id[:, None] * x_sf_stride_m + sf_offsets[None, :], + _pack_bytes(sf_bytes, 4), + mask=token_mask[:, None], + ) + + if shared_x is not None: + shared_ptrs = shared_x + token_id[:, None] * shared_x_stride_m + if SHARED_MXFP8: + fp8, _, scale_exp = _quantize_fp8_groups( + hidden.to(tl.float32), BLOCK_M, BLOCK_K, 32, True + ) + tl.store(shared_ptrs + k_ids[None, :], fp8, mask=token_mask[:, None]) + tl.store( + shared_x_sf + + _shared_sf_row(token_id, SHARED_BLOCK_M) * shared_x_sf_stride_m + + k_block_id * shared_x_sf_stride_k, + _pack_bytes(scale_exp, BLOCK_K // 32), + mask=token_mask, + ) + else: + tl.store(shared_ptrs + k_ids[None, :], hidden, mask=token_mask[:, None]) + + _store_topk( + token_id, + token_mask, + topk_ids, + topk_weights, + is_padding, + topk_idx_out, + topk_weights_out, + topk_ids_stride_m, + topk_ids_stride_k, + topk_weights_stride_m, + topk_weights_stride_k, + is_padding_stride_m, + topk_idx_stride_m, + topk_idx_stride_k, + topk_weights_out_stride_m, + topk_weights_out_stride_k, + top_k, + BLOCK_M, + BLOCK_TOPK, + ) + + +def _resolve_hidden_quant(hidden_quant: QuantKey, block_k: int) -> tuple[int, bool]: + """Map a hidden-state ``QuantKey`` to the kernel's ``(GROUP_K, USE_UE8M0)``.""" + scale = hidden_quant.scale + group_shape = scale.group_shape + if ( + hidden_quant.dtype != FP8_DTYPE + or not hidden_quant.symmetric + or hidden_quant.scale2 is not None + or scale.static + or not group_shape.is_per_group() + or block_k % group_shape.col != 0 + or scale.dtype not in (*_E8M0_SCALE_DTYPES, torch.float32) + ): + raise ValueError( + "MegaMoE input staging requires symmetric dynamic fp8 " + f"per-group quantization with a group size dividing {block_k} and " + f"E8M0 or fp32 scales, got {hidden_quant}." + ) + return group_shape.col, scale.dtype in _E8M0_SCALE_DTYPES + + +def prepare_megamoe_inputs( + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + x: torch.Tensor, + x_sf: torch.Tensor | None, + topk_idx_out: torch.Tensor, + topk_weights_out: torch.Tensor, + is_padding: torch.Tensor | None = None, + shared_x_sf: torch.Tensor | None = None, + shared_block_m: int | None = None, + hidden_quant: QuantKey | None = kMxfp8Dynamic, +) -> None: + """Quantize hidden states and repack top-k routing for DeepGEMM MegaMoE. + + Args: + hidden_states: Input activations of shape ``[num_tokens, hidden_dim]``. + topk_weights: Router top-k weights of shape ``[num_tokens, top_k]``. + topk_ids: Router top-k expert ids of shape ``[num_tokens, top_k]``. + x: Output buffer for the dispatched hidden states. + x_sf: Output buffer for the hidden-state scale factors; ``None`` when + ``hidden_quant`` is ``None``. + topk_idx_out: Output buffer for the repacked top-k expert ids. + topk_weights_out: Output buffer for the repacked top-k weights. + is_padding: Optional per-token mask; padded tokens are not routed. + shared_x_sf: Optional output buffer for shared-expert scale factors. + Must be given together with ``shared_block_m``. + shared_block_m: Block M used to lay out ``shared_x_sf``. + hidden_quant: Hidden-state quantization scheme. E8M0 scales are packed + four per int32 into ``x_sf``; fp32 scales are stored directly. + ``None`` copies the hidden states into ``x`` unquantized. + + """ + num_tokens, hidden_size = hidden_states.shape + if num_tokens == 0: + return + if hidden_size % 128 != 0: + raise ValueError( + "MegaMoE input staging requires hidden_size to be a multiple of 128." + ) + block_k = 128 + quantize = hidden_quant is not None + if quantize != (x_sf is not None): + raise ValueError("MegaMoE input staging needs x_sf iff hidden_quant is set.") + group_k, use_ue8m0 = ( + (block_k, False) + if hidden_quant is None + else _resolve_hidden_quant(hidden_quant, block_k) + ) + top_k = topk_ids.shape[1] + if topk_weights.shape != topk_ids.shape: + raise ValueError( + "MegaMoE input staging requires topk_weights and " + "topk_ids to have the same shape." + ) + if (shared_x_sf is None) != (shared_block_m is None): + raise ValueError( + "MegaMoE shared input staging requires both shared_x_sf and shared_block_m." + ) + if shared_x_sf is not None and not use_ue8m0: + raise ValueError( + "MegaMoE shared input staging currently requires " + "UE8M0-packed hidden scales." + ) + if shared_x_sf is not None: + assert shared_block_m is not None + if shared_block_m <= 0: + raise ValueError("MegaMoE shared_block_m must be positive.") + expected_sf_k = hidden_size // 128 + if shared_x_sf.ndim != 2 or shared_x_sf.shape[1] != expected_sf_k: + raise ValueError( + "MegaMoE shared_x_sf must have shape " + f"(*, {expected_sf_k}), got {tuple(shared_x_sf.shape)}." + ) + aligned_block_m = triton.cdiv(shared_block_m, 128) * 128 + required_rows = triton.cdiv(num_tokens, shared_block_m) * aligned_block_m + if shared_x_sf.shape[0] < required_rows: + raise ValueError( + "MegaMoE shared_x_sf has insufficient rows: requires " + f"{required_rows}, got {shared_x_sf.shape[0]}." + ) + + # On GB200, eight-row tiles win from 64 tokens; keep smaller batches untiled. + block_m = 8 if num_tokens >= 64 else 1 + grid = (triton.cdiv(num_tokens, block_m), triton.cdiv(hidden_size, block_k)) + block_topk = triton.next_power_of_2(top_k) + padding_stride_m = is_padding.stride(0) if is_padding is not None else 0 + _prepare_megamoe_inputs_kernel[grid]( + hidden_states, + x, + x_sf, + shared_x_sf, + topk_ids, + topk_weights, + is_padding, + topk_idx_out, + topk_weights_out, + hidden_states.stride(0), + hidden_states.stride(1), + x.stride(0), + x.stride(1), + x_sf.stride(0) if x_sf is not None else 0, + x_sf.stride(1) if x_sf is not None else 0, + shared_x_sf.stride(0) if shared_x_sf is not None else 0, + shared_x_sf.stride(1) if shared_x_sf is not None else 0, + topk_ids.stride(0), + topk_ids.stride(1), + topk_weights.stride(0), + topk_weights.stride(1), + padding_stride_m, + topk_idx_out.stride(0), + topk_idx_out.stride(1), + topk_weights_out.stride(0), + topk_weights_out.stride(1), + num_tokens, + hidden_size, + top_k, + BLOCK_M=block_m, + BLOCK_K=block_k, + GROUP_K=group_k, + BLOCK_TOPK=block_topk, + SHARED_BLOCK_M=shared_block_m or 1, + QUANTIZE=quantize, + USE_UE8M0=use_ue8m0, + num_warps=4, + ) + + +def prepare_megamoe_nvfp4_inputs( + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + x: torch.Tensor, + x_sf: torch.Tensor, + x_scales: torch.Tensor, + topk_idx_out: torch.Tensor, + topk_weights_out: torch.Tensor, + is_padding: torch.Tensor | None = None, + shared_x: torch.Tensor | None = None, + shared_x_sf: torch.Tensor | None = None, + shared_block_m: int | None = None, + w2_global_scale: torch.Tensor | None = None, + l2_alpha: torch.Tensor | None = None, + l2_activation_scale: float = 1.0, +) -> None: + """Stage DeepGEMM NVFP4 MegaMoE inputs (``kNvfp4DynamicToken``). + + Each token is scaled by its own fp32 ``x_scales`` (row amax / 2688) and + quantized to E2M1 with one E4M3 scale per 16 values. + + Args: + hidden_states: BF16 activations of shape ``[num_tokens, hidden_dim]``. + topk_weights: Router top-k weights of shape ``[num_tokens, top_k]``. + topk_ids: Router top-k expert ids of shape ``[num_tokens, top_k]``. + x: Packed E2M1 output, ``[num_tokens, hidden_dim // 2]``. + x_sf: E4M3 scales packed four per int32, ``[num_tokens, hidden_dim // 64]``. + x_scales: Per-token fp32 dequantization scales, ``[num_tokens]``. + topk_idx_out: Output buffer for the repacked top-k expert ids. + topk_weights_out: Output buffer for the repacked top-k weights. + is_padding: Optional per-token mask; padded tokens are not routed. + shared_x: Optional shared-expert input, fp8 (MXFP8 with ``shared_x_sf``) + or a BF16 copy of the hidden states. + shared_x_sf: UE8M0 scales of an fp8 ``shared_x`` in the shared-expert + TMA layout for ``shared_block_m``. + shared_block_m: Block M used to lay out ``shared_x_sf``. + w2_global_scale: Per-local-expert L2 weight global scales. + l2_alpha: Output ``w2_global_scale * l2_activation_scale``, the L2 + dequantization scale of each local expert. Written even without + local tokens, as other ranks' tokens still reach these experts. + l2_activation_scale: Global scale the kernel quantizes L2 inputs with. + + """ + num_tokens, hidden_size = hidden_states.shape + if hidden_size % 128 != 0: + raise ValueError( + "MegaMoE input staging requires hidden_size to be a multiple of 128." + ) + if topk_weights.shape != topk_ids.shape: + raise ValueError( + "MegaMoE input staging requires topk_weights and " + "topk_ids to have the same shape." + ) + shared_mxfp8 = shared_x is not None and shared_x.dtype == FP8_DTYPE + if shared_mxfp8 != (shared_x_sf is not None) or (shared_x_sf is None) != ( + shared_block_m is None + ): + raise ValueError( + "MegaMoE NVFP4 staging needs shared_x_sf and shared_block_m exactly " + "for an fp8 shared_x." + ) + if (w2_global_scale is None) != (l2_alpha is None): + raise ValueError("MegaMoE NVFP4 staging needs w2_global_scale with l2_alpha.") + for name, tensor in (("x", x), ("x_sf", x_sf), ("shared_x", shared_x)): + if tensor is not None and tensor.stride(-1) != 1: + raise ValueError(f"MegaMoE NVFP4 staging requires a row-major {name}.") + + block_m = 8 if num_tokens >= 64 else 1 + top_k = topk_ids.shape[1] + num_local_experts = 0 if l2_alpha is None else l2_alpha.numel() + # One program even without tokens, to write l2_alpha. + grid = (max(triton.cdiv(num_tokens, block_m), 1),) + _prepare_megamoe_nvfp4_inputs_kernel[grid]( + hidden_states, + x.view(torch.uint8), + x_sf, + x_scales, + shared_x, + shared_x_sf, + topk_ids, + topk_weights, + is_padding, + topk_idx_out, + topk_weights_out, + w2_global_scale, + l2_alpha, + hidden_states.stride(0), + hidden_states.stride(1), + x.stride(0), + x_sf.stride(0), + shared_x.stride(0) if shared_x is not None else 0, + shared_x_sf.stride(0) if shared_x_sf is not None else 0, + shared_x_sf.stride(1) if shared_x_sf is not None else 0, + topk_ids.stride(0), + topk_ids.stride(1), + topk_weights.stride(0), + topk_weights.stride(1), + is_padding.stride(0) if is_padding is not None else 0, + topk_idx_out.stride(0), + topk_idx_out.stride(1), + topk_weights_out.stride(0), + topk_weights_out.stride(1), + num_tokens, + l2_activation_scale, + hidden_size, + top_k, + num_local_experts, + BLOCK_M=block_m, + BLOCK_K=128, + BLOCK_TOPK=triton.next_power_of_2(top_k), + BLOCK_E=triton.next_power_of_2(max(num_local_experts, 1)), + SHARED_BLOCK_M=shared_block_m or 1, + SHARED_MXFP8=shared_mxfp8, + num_warps=4, + ) diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 6fb49314fc..7b5fff06e2 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -79,6 +79,9 @@ class RoutedExperts(PluggableLayer): - Executing routed experts via quant_method.apply() """ + # Hash-routing table, owned by the gate. + hash_indices_table: torch.Tensor | None = None + def __init__( self, layer_name: str, @@ -107,6 +110,7 @@ class RoutedExperts(PluggableLayer): swiglu_beta: float | None = None, e_score_correction_bias: torch.Tensor | None = None, apply_router_weight_on_input: bool = False, + hash_indices_table: torch.Tensor | None = None, ): super().__init__() self.layer_name = layer_name @@ -152,6 +156,8 @@ class RoutedExperts(PluggableLayer): # Set by bind_routed_experts_capturer for monolithic kernels. self.routing_sink: RoutedExpertsSink | None = None self.apply_router_weight_on_input = apply_router_weight_on_input + # The gate owns this Parameter: keep it out of this module's parameters. + object.__setattr__(self, "hash_indices_table", hash_indices_table) # End random parameters self._loaded_expert_biases: set[str] = set() self._fused_shared_expert_quantizer = self._get_fused_shared_expert_quantizer() diff --git a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py b/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py index 5d111fa635..8ece4da679 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py @@ -14,7 +14,7 @@ import torch from vllm.config import VllmConfig from vllm.distributed.parallel_state import get_tp_group -from vllm.forward_context import get_forward_context +from vllm.forward_context import ForwardContext, get_forward_context from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig from vllm.platforms import current_platform @@ -85,103 +85,74 @@ class RoutedExpertsCapturer: self.tp_size = vllm_config.parallel_config.tensor_parallel_size def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None: - """Capture expert routing decisions for a specific layer. - - Under data parallelism, ``topk_ids`` may have four different batch - layouts depending on where the DP combine happens and whether - Expert Parallelism (EP) or Sequence Parallelism (SP) is active for the - MoE layer: - - ``n == total`` (naive dispatch): all DP ranks' tokens are - concatenated before routing; we slice out this rank's span - using the cumulative per-rank counts. - - ``n == token_num_per_dp`` (modular-kernel path): DP combine - happens inside ``quant_method.apply``; ``select_experts`` only - ever sees this rank's tokens, so we take the whole tensor. - - ``n == sum(dp_metadata.local_sizes)`` (naive DP+EP dispatch): - sequence-parallel shards from every DP rank are gathered through - the flattened EP group. The shard sizes include CUDA-graph / SP - padding, so we use them to locate this DP rank's unpadded rows. - - ``n == ceil(token_num_per_dp / tp_size)`` (SP + modular-kernel - path): tokens were split along dim=0 across the TP group by - ``_sequence_parallel_context`` - (``moe_runner_base.py:_sequence_parallel_context``), so each - TP rank only sees its shard. We all-gather along dim=0 to - reconstruct this DP rank's full routing tensor. SP pads with - ceil-div (see ``_compute_sp_num_tokens`` in - ``forward_context.py``), so the gathered tensor may contain a - few trailing padding rows which are trimmed by the downstream - ``[:token_num_per_dp]`` slice. - - Args: - layer_id: The layer index. - topk_ids: Tensor of shape (batch_size, num_routed_experts). - - """ - ctx = get_forward_context() - if ctx.dp_metadata is None: - local_topk_ids = topk_ids - else: - num_tokens_dp = ctx.dp_metadata.num_tokens_across_dp_cpu - num_local_tokens = int(num_tokens_dp[self.dp_rank].item()) - total = int(num_tokens_dp.sum().item()) - n = topk_ids.shape[0] - shard_sizes = getattr(ctx.dp_metadata, "local_sizes", None) - gathered_size = sum(shard_sizes) if shard_sizes is not None else -1 - - if n == total: - end = int(num_tokens_dp[: self.dp_rank + 1].sum().item()) - local_topk_ids = topk_ids[end - num_local_tokens : end] - elif n == num_local_tokens: - local_topk_ids = topk_ids - elif shard_sizes is not None and n == gathered_size: - # Naive DP+EP dispatch gathers the sequence-parallel shards in - # flattened DP-then-TP order. ``local_sizes`` is the exact - # all-gatherv layout, including padding. Select this DP rank's - # contiguous shard group, then trim its trailing padding. - num_dp_ranks = len(num_tokens_dp) - assert len(shard_sizes) % num_dp_ranks == 0 - shards_per_dp_rank = len(shard_sizes) // num_dp_ranks - first_shard = self.dp_rank * shards_per_dp_rank - start = sum(shard_sizes[:first_shard]) - local_topk_ids = topk_ids[start : start + num_local_tokens] - elif ( - self.tp_size > 1 - and n == (num_local_tokens + self.tp_size - 1) // self.tp_size - ): - # SP + modular-kernel path. All-gather across the TP - # group along dim=0 to reconstruct the full per-DP-rank - # tensor; keep only the first ``token_num_per_dp`` rows - # (trailing rows are SP ceil-div padding). The TP group - # is always initialized on real rollout workers, and - # every rank in the group reaches this branch in - # lockstep (bind is per-FusedMoEFactory layer, SP is a global - # condition), so a bare all_gather here will not - # deadlock -- let it raise if the precondition is - # violated rather than skip silently. - # - # ``topk_ids`` is already whatever the router produced - # (typically int32/int64, both supported by NCCL); the - # downstream buffer assignment narrows to the capture dtype. - local_topk_ids = get_tp_group().all_gather(topk_ids, dim=0)[ - :num_local_tokens - ] - else: - sp_expected = (num_local_tokens + self.tp_size - 1) // self.tp_size - raise AssertionError( - "RoutedExpertsCapturer: unexpected topk_ids batch " - f"dim {n} (expected {total}, {num_local_tokens}, " - f"{gathered_size}, or {sp_expected} for " - f"dp_rank={self.dp_rank}, tp_size={self.tp_size})" - ) - if not 0 <= layer_id < self.device_buffer.shape[1]: raise IndexError( f"routed-experts layer {layer_id} exceeds capture buffer " f"layer count {self.device_buffer.shape[1]}" ) + local_topk_ids = self._restore_local_topk_ids(topk_ids, get_forward_context()) self.device_buffer[: len(local_topk_ids), layer_id] = local_topk_ids + def _restore_local_topk_ids( + self, topk_ids: torch.Tensor, ctx: ForwardContext + ) -> torch.Tensor: + """Restore this DP rank's token rows, trimming trailing SP padding. + + Layout priority: cross-DP batch, local batch, gathered DP/TP shards, + then a local TP shard. Gathered rows are DP-major/TP-minor; preserve + this priority when layout sizes coincide. PCP is unsupported. + + Without DP metadata, use the pre-SP batch descriptor. Unknown counts + pass through; known counts must match a full, sharded or padded batch. + A full local batch avoids gathering, including zero/one-token batches + where only the publishing TP0 rank needs the real rows. + """ + n = topk_ids.shape[0] + dp_metadata = ctx.dp_metadata + if dp_metadata is None: + if ctx.batch_descriptor is None: + return topk_ids + num_local_tokens = ctx.batch_descriptor.num_tokens + else: + num_tokens_dp = dp_metadata.num_tokens_across_dp_cpu + num_local_tokens = int(num_tokens_dp[self.dp_rank].item()) + if n == int(num_tokens_dp.sum().item()): + end = int(num_tokens_dp[: self.dp_rank + 1].sum().item()) + return topk_ids[end - num_local_tokens : end] + + if n == num_local_tokens: + return topk_ids + + if dp_metadata is not None: + shard_sizes = dp_metadata.local_sizes + if shard_sizes is not None and n == sum(shard_sizes): + num_dp_ranks = len(num_tokens_dp) + if len(shard_sizes) % num_dp_ranks != 0: + raise ValueError( + "RoutedExpertsCapturer: gathered shard count " + f"{len(shard_sizes)} is not divisible by DP rank count " + f"{num_dp_ranks}" + ) + shards_per_dp_rank = len(shard_sizes) // num_dp_ranks + first_shard = self.dp_rank * shards_per_dp_rank + start = sum(shard_sizes[:first_shard]) + return topk_ids[start : start + num_local_tokens] + + shard_n = (num_local_tokens + self.tp_size - 1) // self.tp_size + if n == shard_n: + return get_tp_group().all_gather(topk_ids.contiguous(), dim=0)[ + :num_local_tokens + ] + if dp_metadata is None and n == shard_n * self.tp_size: + return topk_ids[:num_local_tokens] + + raise AssertionError( + f"RoutedExpertsCapturer: unexpected topk_ids batch dim {n} " + f"(local_tokens={num_local_tokens}, " + f"dp_rank={self.dp_rank}, tp_size={self.tp_size})" + ) + def snapshot_routing_data(self, num_tokens: int) -> torch.Tensor: """Return a stable snapshot of the current routing data.""" return self.device_buffer[:num_tokens].to(self.output_dtype) diff --git a/vllm/model_executor/layers/fused_moe/router/aiter_shared_routed_fused_moe_router.py b/vllm/model_executor/layers/fused_moe/router/aiter_shared_routed_fused_moe_router.py index 42d4b5d4ad..19e631659a 100644 --- a/vllm/model_executor/layers/fused_moe/router/aiter_shared_routed_fused_moe_router.py +++ b/vllm/model_executor/layers/fused_moe/router/aiter_shared_routed_fused_moe_router.py @@ -84,7 +84,14 @@ class AiterSharedRoutedFusedMoERouter(BaseRouter): total_topk_weights_slice = total_topk_weights[:M] topk_ids_slice = total_topk_ids[:M, :topk] - topk_func = dispatch_topk_softmax_func(use_rocm_aiter=True) + topk_func = dispatch_topk_softmax_func( + use_rocm_aiter=True, + topk_weights=total_topk_weights_slice, + topk_indices=topk_ids_slice, + gating_output=router_logits, + num_shared_experts=num_fse, + shared_expert_scoring_func="sigmoid", + ) topk_func( total_topk_weights_slice, topk_ids_slice, @@ -110,7 +117,10 @@ class AiterSharedRoutedFusedMoERouter(BaseRouter): ) topk_func = dispatch_topk_softmax_func( - use_rocm_aiter=rocm_aiter_ops.is_fused_moe_enabled() + use_rocm_aiter=rocm_aiter_ops.is_fused_moe_enabled(), + topk_weights=topk_weights, + topk_indices=topk_ids, + gating_output=routing_logits, ) topk_weights, topk_ids = topk_func( topk_weights, diff --git a/vllm/model_executor/layers/fused_moe/router/fused_topk_router.py b/vllm/model_executor/layers/fused_moe/router/fused_topk_router.py index 28f67b5c69..766940f271 100644 --- a/vllm/model_executor/layers/fused_moe/router/fused_topk_router.py +++ b/vllm/model_executor/layers/fused_moe/router/fused_topk_router.py @@ -61,10 +61,73 @@ def vllm_topk_sigmoid( return topk_weights, topk_indices +# At E=512 / top-10, AITER's topk_gating is faster than the legacy launcher +# through this many tokens; its one-row generic prefill path regresses above it. +_AITER_TOPK_GATING_MAX_TOKENS = 4096 + + +def _aiter_topk_gating_supported( + topk_weights: torch.Tensor, + topk_indices: torch.Tensor, + gating_output: torch.Tensor, + num_shared_experts: int, + shared_expert_scoring_func: str, +) -> bool: + """Whether AITER's ``topk_gating`` can serve this softmax top-k launch. + + It needs contiguous gating rows and at most ``_AITER_TOPK_GATING_MAX_TOKENS`` + tokens. Fused shared experts are scored with sigmoid after the routed + top-k: 1, 2, 4 or 8 of them, with ``topk_weights`` and ``topk_indices`` + sharing a row stride and ``topk_weights`` wide enough for the shared + columns. Shared expert ids stay as the caller prefilled them. + """ + if ( + not gating_output.is_contiguous() + or gating_output.shape[0] > _AITER_TOPK_GATING_MAX_TOKENS + ): + return False + if num_shared_experts == 0: + return not shared_expert_scoring_func + return ( + shared_expert_scoring_func == "sigmoid" + and num_shared_experts in (1, 2, 4, 8) + and topk_weights.stride(0) == topk_indices.stride(0) + and topk_weights.shape[-1] >= topk_indices.shape[-1] + num_shared_experts + ) + + def dispatch_topk_softmax_func( use_rocm_aiter: bool = False, + *, + topk_weights: torch.Tensor | None = None, + topk_indices: torch.Tensor | None = None, + gating_output: torch.Tensor | None = None, + num_shared_experts: int = 0, + shared_expert_scoring_func: str = "", ) -> Callable[..., tuple[torch.Tensor, ...]]: + """Pick the softmax top-k implementation for one launch. + + On ROCm with AITER, ``topk_gating`` is chosen on gfx942 and gfx950 when the + launch described by the tensors and shared-expert arguments is one it supports; + every other AITER launch uses the legacy ``topk_softmax``. The tensors are + optional: without them the launch cannot be judged and the legacy launcher + is used. + """ if use_rocm_aiter: + if ( + rocm_aiter_ops.is_topk_gating_enabled() + and topk_weights is not None + and topk_indices is not None + and gating_output is not None + and _aiter_topk_gating_supported( + topk_weights, + topk_indices, + gating_output, + num_shared_experts, + shared_expert_scoring_func, + ) + ): + return rocm_aiter_ops.topk_gating return rocm_aiter_ops.topk_softmax return vllm_topk_softmax @@ -104,7 +167,10 @@ def fused_topk( if scoring_func == "softmax": topk_func = dispatch_topk_softmax_func( - use_rocm_aiter=rocm_aiter_ops.is_fused_moe_enabled() + use_rocm_aiter=rocm_aiter_ops.is_fused_moe_enabled(), + topk_weights=topk_weights, + topk_indices=topk_ids, + gating_output=gating_output, ) topk_weights, topk_ids = topk_func( topk_weights, topk_ids, token_expert_indices, gating_output, renormalize diff --git a/vllm/model_executor/layers/fused_moe/router/gate_linear.py b/vllm/model_executor/layers/fused_moe/router/gate_linear.py index 8ecd244b61..c6b1da890b 100644 --- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py +++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py @@ -24,14 +24,15 @@ class GateLinear(ReplicatedLinear): 1. cuteDSL ll_bf16_gemm (SM90+, M<=16, bf16 in, fp32 out, any shape w/ K divisible by 8) - 2. fp32 specialized kernel (SM90+ or gfx950, bf16/fp32 in, fp32 out, - M<=32, model-specific shapes) + 2. fp32 specialized kernel (bf16/fp32 in, fp32 out, M<=32; SM90+ + model-specific shapes, gfx950 any shape) 3. cuteDSL ll_fp32w_gemm (SM90+, M<=32, bf16/fp32 in, fp32 out, any shape) 4. bf16x3 CuteDSL kernel (SM100, bf16 in, fp32 weight) 5. ROCm bf16x3 router GEMM (gfx950, bf16 in, fp32 weight, fp32 out, M>=2048) - 6. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype) + 6. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype); gfx950 + routes low M through the tier 2 kernel 7. F.linear via ReplicatedLinear (ultimate fallback) The ``out_dtype`` attribute is mutable and can be set after init @@ -73,16 +74,6 @@ class GateLinear(ReplicatedLinear): from vllm.platforms.rocm import on_gfx950 is_gfx950 = on_gfx950() - is_rocm_fp32_shape = False - if is_gfx950: - from vllm.model_executor.layers.fused_moe.router.rocm_fp32_router_gemm import ( # noqa: E501 - ROCM_FP32_ROUTER_GEMM_SUPPORTED_SHAPES, - ) - - is_rocm_fp32_shape = ( - input_size, - output_size, - ) in ROCM_FP32_ROUTER_GEMM_SUPPORTED_SHAPES batch_invariant = envs.VLLM_BATCH_INVARIANT can_use_specialized_kernels = ( current_platform.is_cuda() @@ -107,6 +98,7 @@ class GateLinear(ReplicatedLinear): return_bias=return_bias, ) self.out_dtype = out_dtype + self.is_gfx950 = is_gfx950 # A quantized gate exposes no plain ``weight``, so every specialized # tier below is disabled and forward falls back to ReplicatedLinear. @@ -127,7 +119,7 @@ class GateLinear(ReplicatedLinear): and (is_hopper or is_blackwell) and (input_size, output_size) in self.FP32_SUPPORTED_SHAPES ) - or (is_gfx950 and is_rocm_fp32_shape) + or is_gfx950 ) ) self.allow_bf16x3_router_gemm = ( @@ -349,7 +341,10 @@ class GateLinear(ReplicatedLinear): # Tier 6: cuBLAS bf16→fp32 if self.allow_cublas_router_gemm and x.dtype == torch.bfloat16: - output = torch.mm(x, self.weight.T, out_dtype=torch.float32) + if self.is_gfx950: + output = torch.ops.vllm.fp32_router_gemm_dispatch(x, self.weight, False) + else: + output = torch.mm(x, self.weight.T, out_dtype=torch.float32) return self._return(output) # Tier 7: F.linear (ReplicatedLinear) @@ -465,17 +460,20 @@ def fp32_router_gemm_dispatch_impl( This must be wrapped in a custom op because our torch.compile integration does not support runtime dispatching on num_tokens. """ - if x.shape[0] <= _FP32_ROUTER_GEMM_MAX_TOKENS: - if current_platform.is_rocm(): - from vllm.model_executor.layers.fused_moe.router.rocm_fp32_router_gemm import ( # noqa: E501 - can_use_rocm_fp32_router_gemm, - rocm_fp32_router_gemm, - ) + if current_platform.is_rocm(): + from vllm.model_executor.layers.fused_moe.router.rocm_fp32_router_gemm import ( # noqa: E501 + can_use_rocm_fp32_router_gemm, + rocm_fp32_router_gemm, + ) + if x.shape[0] <= _FP32_ROUTER_GEMM_MAX_TOKENS: x = x.contiguous() if can_use_rocm_fp32_router_gemm(x, weight): return rocm_fp32_router_gemm(x, weight) - return torch.nn.functional.linear(x.float(), weight) + if weight.dtype == torch.bfloat16: + return torch.mm(x, weight.T, out_dtype=torch.float32) + return torch.nn.functional.linear(x.float(), weight) + if x.shape[0] <= _FP32_ROUTER_GEMM_MAX_TOKENS: return ops.fp32_router_gemm(x, weight) if allow_bf16x3_router_gemm and x.dtype == torch.bfloat16: diff --git a/vllm/model_executor/layers/fused_moe/router/rocm_fp32_router_gemm.py b/vllm/model_executor/layers/fused_moe/router/rocm_fp32_router_gemm.py index 785bd8a72c..f76318ae66 100644 --- a/vllm/model_executor/layers/fused_moe/router/rocm_fp32_router_gemm.py +++ b/vllm/model_executor/layers/fused_moe/router/rocm_fp32_router_gemm.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Low-token FP32 router GEMM for gfx950.""" +"""Low-token router GEMM with FP32 output for gfx950 (FP32 or BF16 weights).""" import torch @@ -9,15 +9,10 @@ from vllm.platforms.rocm import on_gfx950 from vllm.triton_utils import tl, triton _MAX_TOKENS = 32 -ROCM_FP32_ROUTER_GEMM_SUPPORTED_SHAPES = frozenset( - { - (3072, 256), - (4096, 8), - (4096, 192), - (6144, 128), - (6144, 256), - } -) +# From this many multiply-adds on, F.linear wins on MI355. +_MAX_MACS = 1 << 27 +# Above this num_tokens * num_experts, hipBLASLt wins for BF16 weights. +_MAX_BF16_TOKEN_EXPERTS = 8192 @triton.jit @@ -30,6 +25,7 @@ def _rocm_fp32_router_gemm_kernel( N: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_K: tl.constexpr, + EVEN_K: tl.constexpr, ): pid = tl.program_id(0) pid_m = pid // N @@ -40,12 +36,27 @@ def _rocm_fp32_router_gemm_kernel( for k_start in range(0, K, BLOCK_K): offsets_k_block = k_start + offsets_k - hidden_states = tl.load( - hidden_states_ptr + offsets_m[:, None] * K + offsets_k_block[None, :], - mask=offsets_m[:, None] < M, - other=0.0, - ).to(tl.float32) - router_weight = tl.load(router_weight_ptr + pid_n * K + offsets_k_block) + if EVEN_K: + hidden_states = tl.load( + hidden_states_ptr + offsets_m[:, None] * K + offsets_k_block[None, :], + mask=offsets_m[:, None] < M, + other=0.0, + ).to(tl.float32) + router_weight = tl.load(router_weight_ptr + pid_n * K + offsets_k_block).to( + tl.float32 + ) + else: + k_mask = offsets_k_block < K + hidden_states = tl.load( + hidden_states_ptr + offsets_m[:, None] * K + offsets_k_block[None, :], + mask=(offsets_m[:, None] < M) & k_mask[None, :], + other=0.0, + ).to(tl.float32) + router_weight = tl.load( + router_weight_ptr + pid_n * K + offsets_k_block, + mask=k_mask, + other=0.0, + ).to(tl.float32) partials += hidden_states * router_weight[None, :] accumulator = tl.sum(partials, axis=1) @@ -57,28 +68,36 @@ def _rocm_fp32_router_gemm_kernel( def _launch_config( - hidden_size: int, - num_experts: int, - num_tokens: int, + num_experts: int, num_tokens: int, bf16_weight: bool ) -> tuple[int, int, int]: - if (hidden_size, num_experts) == (4096, 8): - return 1, 1024, 4 + """Return (BLOCK_M, BLOCK_K, num_warps), swept on MI355 for M<=32. + One program computes BLOCK_M tokens for one expert, so narrow routers keep + BLOCK_M=1 to fill the GPU and wide routers grow BLOCK_M to reuse weights. + """ + if bf16_weight: + return (1 if num_tokens <= 4 else 2), 1024, 2 + if num_experts < 64: + return 1, 1024, 4 + wide = num_experts > 256 + if num_tokens <= 2: + return 1, 512, 2 if num_tokens <= 4: - return 1, 1024, 8 + return (2, 1024, 2) if wide else (2, 2048, 4) + if num_tokens <= 8: + return (4, 512, 2) if wide else (4, 1024, 4) if num_tokens <= 16: - block_m = 8 if (hidden_size, num_experts) == (6144, 256) else 4 - return block_m, 1024, 8 - if (hidden_size, num_experts) == (3072, 256): - return 8, 1024, 8 - return 4, 2048, 8 + return (8, 512, 2) if wide else (4, 1024, 4) + if num_experts >= 768: + return 16, 256, 2 + return (8, 256, 2) if wide else (4, 512, 2) def can_use_rocm_fp32_router_gemm( hidden_states: torch.Tensor, router_weight: torch.Tensor, ) -> bool: - """Return whether the tensors match the tuned gfx950 fast path.""" + """Return whether the tensors can use the gfx950 fast path.""" try: _validate_inputs(hidden_states, router_weight) except (RuntimeError, ValueError): @@ -96,26 +115,30 @@ def _validate_inputs( raise ValueError("hidden_states and router_weight must be 2D tensors") if hidden_states.dtype not in (torch.bfloat16, torch.float32): raise ValueError("hidden_states must have dtype bfloat16 or float32") - if router_weight.dtype != torch.float32: - raise ValueError("router_weight must have dtype float32") + if router_weight.dtype not in (torch.float32, torch.bfloat16): + raise ValueError("router_weight must have dtype float32 or bfloat16") if hidden_states.device.type != "cuda" or router_weight.device.type != "cuda": raise ValueError("hidden_states and router_weight must be GPU tensors") if hidden_states.device != router_weight.device: raise ValueError("hidden_states and router_weight must be on the same device") if not hidden_states.is_contiguous() or not router_weight.is_contiguous(): raise ValueError("hidden_states and router_weight must be contiguous") - shape = (hidden_states.shape[1], router_weight.shape[0]) - if ( - shape not in ROCM_FP32_ROUTER_GEMM_SUPPORTED_SHAPES - or router_weight.shape[1] != shape[0] - ): - raise ValueError( - "supported (hidden_size, num_experts) shape pairs are " - "(3072, 256), (4096, 8), (4096, 192), (6144, 128), " - "and (6144, 256)" - ) + if router_weight.shape[1] != hidden_states.shape[1]: + raise ValueError("router_weight and hidden_states hidden sizes differ") if not 0 <= hidden_states.shape[0] <= _MAX_TOKENS: raise ValueError(f"num_tokens must be in [0, {_MAX_TOKENS}]") + if hidden_states.shape[0] * router_weight.numel() >= _MAX_MACS: + raise ValueError( + f"num_tokens * num_experts * hidden_size must be < {_MAX_MACS}" + ) + if ( + router_weight.dtype == torch.bfloat16 + and hidden_states.shape[0] * router_weight.shape[0] > _MAX_BF16_TOKEN_EXPERTS + ): + raise ValueError( + f"num_tokens * num_experts must be <= {_MAX_BF16_TOKEN_EXPERTS} " + "for bfloat16 weights" + ) def rocm_fp32_router_gemm( @@ -131,7 +154,10 @@ def rocm_fp32_router_gemm( if num_tokens == 0: return output - block_m, block_k, num_warps = _launch_config(hidden_size, num_experts, num_tokens) + block_m, block_k, num_warps = _launch_config( + num_experts, num_tokens, router_weight.dtype == torch.bfloat16 + ) + block_k = min(block_k, triton.next_power_of_2(hidden_size)) grid = (triton.cdiv(num_tokens, block_m) * num_experts,) _rocm_fp32_router_gemm_kernel[grid]( hidden_states, @@ -142,6 +168,7 @@ def rocm_fp32_router_gemm( N=num_experts, BLOCK_M=block_m, BLOCK_K=block_k, + EVEN_K=hidden_size % block_k == 0, num_warps=num_warps, num_stages=1, ) diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index b1a6dd5632..776c2064fa 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -325,6 +325,7 @@ def _fp8_quantize( A_scale: torch.Tensor | None, per_act_token: bool, block_shape: list[int] | None = None, + is_scale_transposed: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: """Perform fp8 quantization on the inputs. If a block_shape is provided, the output will be blocked. @@ -339,7 +340,9 @@ def _fp8_quantize( assert not per_act_token assert len(block_shape) == 2 _, block_k = block_shape[0], block_shape[1] - A, A_scale = per_token_group_quant_fp8(A, block_k) + A, A_scale = per_token_group_quant_fp8( + A, block_k, column_major_scales=is_scale_transposed + ) assert cdiv(A.size(-1), block_k) == A_scale.size(-1) return A, A_scale @@ -462,6 +465,7 @@ def moe_kernel_quantize_input( per_act_token_quant: bool, block_shape: list[int] | None = None, is_scale_swizzled: bool = True, + is_scale_transposed: bool = False, quantization_emulation: bool = False, mx_alignment: int = 0, ) -> tuple[torch.Tensor, torch.Tensor | None]: @@ -469,7 +473,9 @@ def moe_kernel_quantize_input( if quantization_emulation: return _fp8_quantize_dequantize(A, A_scale) else: - return _fp8_quantize(A, A_scale, per_act_token_quant, block_shape) + return _fp8_quantize( + A, A_scale, per_act_token_quant, block_shape, is_scale_transposed + ) elif quant_dtype == torch.int8: if quantization_emulation: raise NotImplementedError( diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 5b77caa256..7a97321074 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -31,6 +31,7 @@ from vllm.model_executor.layers.quantization.base_config import ( ) from vllm.model_executor.layers.utils import ( dispatch_unquantized_gemm, + maybe_pad_weight_avoid_cache_aliasing, ) from vllm.model_executor.parameter import ( BasevLLMParameter, @@ -225,6 +226,8 @@ class UnquantizedLinearMethod(LinearMethodBase): and weight.stride(0) != 1 ): layer.weight.data = weight.t().contiguous().t() + elif current_platform.is_rocm(): + layer.weight.data = maybe_pad_weight_avoid_cache_aliasing(layer.weight.data) def apply( self, diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 8f42f23970..c93977d86c 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -479,6 +479,8 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase): e_score_correction_bias=layer.e_score_correction_bias, routed_scaling_factor=layer.routed_scaling_factor, routing_sink=layer.routing_sink, + input_ids=input_ids, + hash_indices_table=layer.hash_indices_table, ) @@ -1132,4 +1134,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): e_score_correction_bias=layer.e_score_correction_bias, routed_scaling_factor=layer.routed_scaling_factor, routing_sink=layer.routing_sink, + input_ids=input_ids, + hash_indices_table=layer.hash_indices_table, ) diff --git a/vllm/model_executor/layers/quantization/online/lm_head.py b/vllm/model_executor/layers/quantization/online/lm_head.py new file mode 100644 index 0000000000..9543aab9c0 --- /dev/null +++ b/vllm/model_executor/layers/quantization/online/lm_head.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Online-quantized copies of an already-loaded lm_head.""" + +from itertools import chain +from typing import Literal +from weakref import WeakKeyDictionary + +import torch +from torch.nn import Module, Parameter + +from vllm.logger import init_logger +from vllm.model_executor.kernels.linear import init_nvfp4_linear_kernel +from vllm.model_executor.layers.linear import UnquantizedLinearMethod +from vllm.model_executor.layers.quantization.online.fp8 import ( + Fp8PtpcOnlineLinearMethod, + OnlineLinearBase, +) +from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( + FLOAT4_E2M1_MAX, + ref_nvfp4_quant, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import weight_amax +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + UnquantizedEmbeddingMethod, +) +from vllm.model_executor.model_loader.reload.layerwise import get_layerwise_info +from vllm.model_executor.utils import replace_parameter +from vllm.utils.torch_utils import set_default_torch_dtype + +logger = init_logger(__name__) + + +def _quantize_nvfp4( + weight: torch.Tensor, chunk_rows: int = 8192 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """NVFP4-quantize the rows of `weight` without the SM100+ quant kernel. + + Returns packed E2M1 codes (low nibble first), E4M3 scales per 16 values and + the fp32 dequantization global scale, in the checkpoint layout. + """ + fp8_max = torch.finfo(torch.float8_e4m3fn).max + global_scale = (weight_amax(weight).float() / (FLOAT4_E2M1_MAX * fp8_max)).clamp( + min=1e-12 + ) + inv_global_scale = (1.0 / global_scale).reshape(()) + # E2M1 magnitudes 0, 0.5, ..., 6 doubled to integers, mapped to their codes. + code_of = torch.zeros(13, dtype=torch.uint8, device=weight.device) + code_of[torch.tensor([0, 1, 2, 3, 4, 6, 8, 12])] = torch.arange( + 8, dtype=torch.uint8 + ).to(weight.device) + n, k = weight.shape + packed = torch.empty(n, k // 2, dtype=torch.uint8, device=weight.device) + scales = torch.empty(n, k // 16, dtype=torch.float8_e4m3fn, device=weight.device) + for start in range(0, n, chunk_rows): + rows = slice(start, start + chunk_rows) + values, block_scales = ref_nvfp4_quant(weight[rows], inv_global_scale, 16) + codes = code_of[(values.abs() * 2).long()] | ((values < 0).to(torch.uint8) << 3) + packed[rows] = codes[:, 0::2] | (codes[:, 1::2] << 4) + scales[rows] = block_scales.to(torch.float8_e4m3fn) + return packed, scales, global_scale + + +class _LMHeadNvfp4Method(OnlineLinearBase): + """Weight-only NVFP4 lm_head copy on the platform's W4A16 kernel; + activations stay in the model dtype.""" + + def process_weights_after_loading(self, layer: Module) -> None: + packed, scales, global_scale = _quantize_nvfp4(layer.weight) + replace_parameter(layer, "weight", packed) + layer.weight_scale = Parameter(scales, requires_grad=False) + layer.weight_global_scale = Parameter(global_scale, requires_grad=False) + self.kernel = init_nvfp4_linear_kernel(use_a16=True) + self.kernel.process_weights_after_loading(layer) + + def apply( + self, layer: Module, x: torch.Tensor, bias: torch.Tensor | None = None + ) -> torch.Tensor: + return self.kernel.apply_weights(layer, x, bias) + + +_LM_HEAD_METHODS: dict[str, type[OnlineLinearBase]] = { + "fp8": Fp8PtpcOnlineLinearMethod, + "nvfp4": _LMHeadNvfp4Method, +} + + +def _quantize_copy(lm_head: ParallelLMHead, quantization: str) -> ParallelLMHead: + if ( + not isinstance( + lm_head.quant_method, (UnquantizedEmbeddingMethod, UnquantizedLinearMethod) + ) + or getattr(lm_head, "bias", None) is not None + or lm_head.weight.dtype not in (torch.float16, torch.bfloat16) + ): + raise ValueError( + "Quantizing an lm_head copy needs an unquantized fp16/bf16 lm_head " + "without bias." + ) + # Online methods take their output dtype from the default dtype, which is + # only the model dtype while the model itself is being built. + dtype = lm_head.weight.dtype + with set_default_torch_dtype(dtype): + head = ParallelLMHead( + lm_head.num_embeddings, + lm_head.embedding_dim, + params_dtype=dtype, + org_num_embeddings=lm_head.org_vocab_size, + padding_size=lm_head.padding_size, + prefix=lm_head.prefix, + disable_tp=lm_head.disable_tp, + quant_method=_LM_HEAD_METHODS[quantization](), + ) + replace_parameter(head, "weight", lm_head.weight.data) + head.quant_method.process_weights_after_loading(head) + # The copy is never loaded. The online loader wrapper captures the layer, + # so keeping it would leave each copy to the cycle collector. + for param in head.parameters(recurse=False): + if hasattr(param, "weight_loader"): + del param.weight_loader + return head + + +# Each quantized copy maps to the head it is derived from. Weak keys, and no +# module attribute, so the copy never owns the source as a submodule. +_SOURCES: WeakKeyDictionary[Module, tuple[ParallelLMHead, str]] = WeakKeyDictionary() + + +def quantized_lm_head_copy( + lm_head: ParallelLMHead, quantization: Literal["fp8", "nvfp4"] +) -> ParallelLMHead: + """A ParallelLMHead holding this rank's shard of `lm_head`, quantized. + + "fp8" uses per-row weight and per-token activation scales on the platform's + FP8 linear kernel; "nvfp4" is weight-only on the + platform's W4A16 kernel. `lm_head` is left + unchanged and no bf16 copy of its weight is made. The copy is derived state: + `refresh_quantized_lm_heads` re-derives it after `lm_head` changes. + """ + head = _quantize_copy(lm_head, quantization) + # Keep it out of layerwise reload, which would re-run + # process_weights_after_loading on the already-quantized weight. + get_layerwise_info(head).reset() + _SOURCES[head] = (lm_head, quantization) + logger.info( + "Quantized a %s copy of the lm_head (%d x %d per rank).", + quantization, + head.output_size_per_partition, + head.input_size_per_partition, + ) + return head + + +@torch.no_grad() +def refresh_quantized_lm_heads(model: Module) -> None: + """Re-derive the quantized lm_head copies in `model` from their source heads. + + Copies in place, so CUDA graphs keep pointing at valid tensors. + """ + for module in model.modules(): + if module not in _SOURCES: + continue + fresh = _quantize_copy(*_SOURCES[module]) + for name, tensor in chain( + fresh.named_parameters(recurse=False), fresh.named_buffers(recurse=False) + ): + getattr(module, name).copy_(tensor) + + +# (parent, attribute, copy) for each copy swapped out by share_source_lm_heads. +_SWAPPED: dict[Module, list[tuple[Module, str, Module]]] = {} + + +def share_source_lm_heads(model: Module) -> None: + """Replace each quantized copy in `model` with its source head. + + `model` then loads weights like a drafter sharing the unquantized head, + until `restore_quantized_lm_heads` swaps the copies back in. + """ + swapped = [ + (module, name, child) + for module in model.modules() + for name, child in module.named_children() + if child in _SOURCES + ] + for module, name, child in swapped: + setattr(module, name, _SOURCES[child][0]) + _SWAPPED[model] = swapped + + +def restore_quantized_lm_heads(model: Module, refresh: bool = True) -> None: + """Undo `share_source_lm_heads`, then re-derive the copies in `model`. + + Args: + model: The model passed to `share_source_lm_heads`, if any. + refresh: Whether to re-derive the copies; False when the source heads + may be mid-reload. + + """ + for module, name, head in _SWAPPED.pop(model, ()): + setattr(module, name, head) + if refresh: + refresh_quantized_lm_heads(model) diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 849b0c9638..439b48a3e6 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -1647,8 +1647,10 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): # If no native backend available, use emulation. if self.mxfp4_backend is Mxfp4MoeBackend.NONE: self.mxfp4_backend = Mxfp4MoeBackend.EMULATION + self.experts_cls = None - self.experts_cls = backend_to_kernel_cls(self.mxfp4_backend)[0] + if self.experts_cls is None: + self.experts_cls = backend_to_kernel_cls(self.mxfp4_backend)[0] logger.info_once( f"Using {self.mxfp4_backend.value} backend for {self.ocp_mx_scheme}" diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index 5a6ee76428..4e113d6912 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py @@ -379,6 +379,79 @@ def align_moe_weights_for_fi( return padded_w13, padded_w2, padded_intermediate +def align_mxfp8_moe_weights_for_fi_trtllm( + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + is_act_and_mul: bool, + min_alignment: int, + hidden_alignment: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, int]: + """align_moe_weights_for_fi for MXFP8 weights and their UE8M0 block scales + for the FlashInfer TRT-LLM MoE kernel, which also pads the hidden size to a + multiple of hidden_alignment. + + Padded weights are zero and padded scales are 127 (UE8M0 for 1.0), so the + padded rows and columns contribute nothing. Scales may be 3D + [E, rows, cols] or flattened [E, rows * cols]; their layout is kept. + Returns the padded intermediate and hidden sizes. + """ + num_experts, hidden_size, intermediate = w2.shape + padded_hidden = round_up(hidden_size, hidden_alignment) + w13, w2, padded_intermediate = align_moe_weights_for_fi( + w13, w2, is_act_and_mul, min_alignment + ) + if padded_intermediate == intermediate and padded_hidden == hidden_size: + return w13, w2, w13_scale, w2_scale, intermediate, hidden_size + + if padded_hidden != hidden_size: + logger.info_once( + "Padding hidden size from %d to %d for MXFP8 TRTLLM MoE kernel.", + hidden_size, + padded_hidden, + ) + padded_w13 = w13.new_zeros((num_experts, w13.shape[1], padded_hidden)) + padded_w13[:, :, :hidden_size] = w13 + padded_w2 = w2.new_zeros((num_experts, padded_hidden, w2.shape[2])) + padded_w2[:, :hidden_size] = w2 + w13, w2 = padded_w13, padded_w2 + + up_mult = 2 if is_act_and_mul else 1 + w13_scale_3d = w13_scale.reshape(num_experts, up_mult * intermediate, -1) + w2_scale_3d = w2_scale.reshape(num_experts, hidden_size, -1) + block = hidden_size // w13_scale_3d.shape[2] + assert padded_hidden % block == 0 and padded_intermediate % block == 0 + + # Pad each of the gate and up halves (or the single non-gated block) on its + # own, like the weights. + padded_w13_scale = w13_scale_3d.new_full( + (num_experts, up_mult * padded_intermediate, padded_hidden // block), 127 + ) + for i in range(up_mult): + dst = i * padded_intermediate + padded_w13_scale[:, dst : dst + intermediate, : w13_scale_3d.shape[2]] = ( + w13_scale_3d[:, i * intermediate : (i + 1) * intermediate] + ) + padded_w2_scale = w2_scale_3d.new_full( + (num_experts, padded_hidden, padded_intermediate // block), 127 + ) + padded_w2_scale[:, :hidden_size, : w2_scale_3d.shape[2]] = w2_scale_3d + + if w13_scale.ndim == 2: + padded_w13_scale = padded_w13_scale.reshape(num_experts, -1) + if w2_scale.ndim == 2: + padded_w2_scale = padded_w2_scale.reshape(num_experts, -1) + return ( + w13, + w2, + padded_w13_scale, + padded_w2_scale, + padded_intermediate, + padded_hidden, + ) + + def _shuffle_deepseek_fp8_moe_weights( w13: torch.Tensor, w2: torch.Tensor, @@ -486,8 +559,9 @@ def prepare_fp8_moe_layer_for_fi( Note that for trtllm we update the model state dict with the scale format needed for these kernels. - Note that for per-tensor, we update the layer's - intermediate size if the weights needed padding. + Note that for per-tensor and for MXFP8 TRT-LLM, we update the layer's + intermediate size (and for MXFP8 TRT-LLM its hidden size) if the weights + needed padding. """ assert hasattr(layer.moe_config, "is_act_and_mul") block_quant = ( @@ -499,6 +573,21 @@ def prepare_fp8_moe_layer_for_fi( # MXFP8 TRT-LLM requires W31 swap + reorder + shuffle. if is_mxfp8 and is_trtllm: + # The kernel needs intermediate % 128 == 0 and hidden % 256 == 0. The + # MoE runner pads the activations to the padded hidden size. + w13, w2, w13_scale, w2_scale, new_intermediate, new_hidden = ( + align_mxfp8_moe_weights_for_fi_trtllm( + w13, + w2, + w13_scale, + w2_scale, + layer.moe_config.is_act_and_mul, + min_alignment=128, + hidden_alignment=256, + ) + ) + layer.moe_config.intermediate_size_per_partition = new_intermediate + layer.moe_config.hidden_dim = new_hidden # FlashInfer TRT-LLM SwiGLU expects [up; gate] but vLLM stores # [gate; up]. Swap both weights and scales before interleaving. if layer.moe_config.is_act_and_mul: diff --git a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py index 45562b12b1..074ee9c95f 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py @@ -13,6 +13,14 @@ from vllm.utils.torch_utils import direct_register_custom_op, is_torch_equal_or_ logger = init_logger(__name__) +def should_use_gfx1250_mx_scale_swizzle() -> bool: + if not current_platform.is_rocm(): + return False + from vllm.platforms.rocm import on_gfx1250 + + return on_gfx1250() + + def should_use_cdna4_mx_scale_swizzle() -> bool: """Whether to use the CDNA4 swizzled scale layout for mxfp4 on gfx950. @@ -27,6 +35,15 @@ def should_use_cdna4_mx_scale_swizzle() -> bool: return on_gfx950() and get_tensor_model_parallel_world_size() <= 2 +def get_aiter_mx_scale_swizzle() -> str | None: + """AITER `swizzle_mx_scale` arg matching the load-time scale swizzle.""" + if should_use_gfx1250_mx_scale_swizzle(): + return "GFX1250_SCALE" + if should_use_cdna4_mx_scale_swizzle(): + return "CDNA4_SCALE" + return None + + def weight_mx_scale(precision_config): """mxfp4 weight scale from a triton_kernels PrecisionConfig: 3.8 exposes it as b_mx_scale, 3.5.1/3.6 as weight_scale. @@ -136,6 +153,17 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps=8): # transpose the tensor so that the quantization axis is on dim1 quant_tensor = quant_tensor.transpose(-2, -1) scale = scale.transpose(-2, -1) + + if should_use_gfx1250_mx_scale_swizzle(): + from aiter.ops.triton.utils.shuffle import shuffle_scale_moe + + assert ( + scale.dim() == 3 and scale.shape[-1] % 32 == 0 and scale.shape[-2] % 4 == 0 + ), f"GFX1250 scale swizzle needs (E,K_SCALE%4,N%32); got {tuple(scale.shape)}" + scale = shuffle_scale_moe( + scale, arch="gfx1250", preshuffle_factor=32, scale_kwidth=4 + ) + quant_tensor = convert_layout( wrap_torch_tensor(quant_tensor, dtype=FP4), value_layout, **value_layout_opts ) diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index 8d285dda75..ef4594e113 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -225,6 +225,8 @@ direct_register_custom_op( def use_aiter_triton_gemm(n, m, k, dtype): + from vllm.platforms.rocm import on_gfx1250 + if ( not rocm_aiter_ops.is_triton_gemm_enabled() # MI300's - fp8nuz=True @@ -233,6 +235,9 @@ def use_aiter_triton_gemm(n, m, k, dtype): ): return False + if on_gfx1250(): + return True + # use hipblaslt for the larger GEMMs if n > 2048 and m > 512: return False @@ -285,6 +290,44 @@ def wvsplitkrc_dispatch(n: int, k: int, m: int, cu_count: int) -> tuple[int, boo return chunkk, fits +# On gfx11 a weight row stride that is a multiple of 2048 B makes the multi-row global +# loads of a GEMM collide on the L2/MALL channel hash. Offsetting the stride by one +# 128 B cache line dodges the collision and keeps each row cache-line aligned. +def cache_aliasing_pad_elems(k: int, element_size: int) -> int: + """Row-stride pad, in elements, that stops an ``[n, k]`` weight aliasing on gfx11, + or 0 when the stride does not alias. + + Args: + k: Number of elements per weight row. + element_size: Size of one element in bytes. + + Returns: + The number of elements to add to the row stride. + + """ + if (k * element_size) % 2048 != 0: + return 0 + return 128 // element_size + + +def maybe_pad_weight_avoid_cache_aliasing(weight: torch.Tensor) -> torch.Tensor: + """Offset a dense weight's row stride so it stops aliasing. + + Returns a view with the original shape whose row stride is one cache line + wider, or ``weight`` unchanged when the pad does not apply. + """ + from vllm.platforms.rocm import on_gfx11 + + if not (current_platform.is_rocm() and on_gfx11()): + return weight + if weight.dim() != 2 or not weight.is_contiguous(): + return weight + pad = cache_aliasing_pad_elems(weight.shape[1], weight.element_size()) + if pad == 0: + return weight + return torch.nn.functional.pad(weight, (0, pad))[..., :-pad] + + def rocm_unquantized_gemm_impl( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None ) -> torch.Tensor: @@ -295,7 +338,7 @@ def rocm_unquantized_gemm_impl( return linear_batch_invariant(x, weight, bias) - from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx950, on_gfx1250 + from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx11, on_gfx950 n = x.numel() // x.size(-1) m = weight.shape[0] @@ -328,38 +371,54 @@ def rocm_unquantized_gemm_impl( x_view = x.reshape(-1, x.size(-1)).contiguous() return ops.wvSplitKrc(x_view, weight, cu_count, bias) - # gfx1250's aiter gemm_a16w16 uses the gluon backend, which requires - # K % 256 == 0 (it walks K with fixed-size descriptors and won't pad a - # partial last tile). Some whitelisted shapes have K=2880 (e.g. gpt-oss-120b - # hidden), so skip aiter there and fall back to the torch GEMM path below. - if use_aiter_triton_gemm(n, m, k, x.dtype) and not (on_gfx1250() and k % 256 != 0): + if use_aiter_triton_gemm(n, m, k, x.dtype): from aiter.ops.triton.gemm_a16w16 import gemm_a16w16 return gemm_a16w16(x, weight, bias) - use_skinny = ( + skinny_gemm_enabled = ( envs.VLLM_ROCM_USE_SKINNY_GEMM and (on_gfx9() or on_gfx1x()) # build (gfx9/gfx11 ISA); fall back to torch GEMM there. # TODO GFX1250: Include once skinny GEMM is supported on gfx1250 and x.dtype in [torch.float16, torch.bfloat16] and k % 8 == 0 - and skinny_operands_compatible ) - if use_skinny: + # wvSplitK allows row strides. Perf impact is only validated on gfx11. + weight_ok_for_wvsplitk = weight.is_contiguous() or ( + on_gfx11() and weight.stride(1) == 1 + ) + + use_wvsplitk = ( + skinny_gemm_enabled + and weight_ok_for_wvsplitk + and (bias is None or bias.is_contiguous()) + and (m == 1 or m > 8) + and 0 < n <= 5 + ) + + use_llmm1 = ( + skinny_gemm_enabled + and weight.is_contiguous() + and bias is None + and m % 4 == 0 + and n == 1 + and k <= 8192 + ) + + if use_wvsplitk: # The skinny kernels assume contiguous K elements. A shape-preserving # reshape can retain a transposed activation's non-contiguous strides. # Note: Only build that view inside the branches that consume it. - if (m == 1 or m > 8) and 0 < n <= 5: - x_view = x.reshape(-1, x.size(-1)).contiguous() - cu_count = num_compute_units() - out = ops.wvSplitK(weight, x_view, cu_count, bias) - return out.reshape(*x.shape[:-1], weight.shape[0]) - elif m % 4 == 0 and n == 1 and k <= 8192 and bias is None: - x_view = x.reshape(-1, x.size(-1)).contiguous() - out = ops.LLMM1(weight, x_view, 4) - return out.reshape(*x.shape[:-1], weight.shape[0]) + x_view = x.reshape(-1, x.size(-1)).contiguous() + out = ops.wvSplitK(weight, x_view, cu_count, bias) + return out.reshape(*x.shape[:-1], weight.shape[0]) + + if use_llmm1: + x_view = x.reshape(-1, x.size(-1)).contiguous() + out = ops.LLMM1(weight, x_view, 4) + return out.reshape(*x.shape[:-1], weight.shape[0]) if rocm_aiter_ops.is_tgemm_enabled(): from aiter.tuned_gemm import tgemm diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index 92e65c839d..574aed3090 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/vllm/model_executor/layers/vocab_parallel_embedding.py @@ -607,6 +607,7 @@ class ParallelLMHead(VocabParallelEmbedding): prefix: str = "", *, disable_tp: bool = False, + quant_method: QuantizeMethodBase | None = None, ): self.has_bias = bias super().__init__( @@ -618,6 +619,7 @@ class ParallelLMHead(VocabParallelEmbedding): quant_config, prefix, disable_tp=disable_tp, + quant_method=quant_method, ) self.quant_config = quant_config if bias: diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index f796bcf615..12ba253671 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -992,7 +992,7 @@ class Qwen4ExpForConditionalGenerationConfig(Qwen3_5ForConditionalGenerationConf parallel_config = vllm_config.parallel_config # MoE SP also enables HC SP without the explicit HC flag. if ( - current_platform.is_cuda() + current_platform.is_cuda_alike() and parallel_config.use_sequence_parallel_moe and parallel_config.pipeline_parallel_size != 1 ): diff --git a/vllm/model_executor/models/d3.py b/vllm/model_executor/models/d3.py index 2af055b7cd..da9019d66f 100644 --- a/vllm/model_executor/models/d3.py +++ b/vllm/model_executor/models/d3.py @@ -81,3 +81,8 @@ class D3ForSequenceClassification(_D3ForSequenceClassificationBase): allow_patterns_overrides=[READOUT_WEIGHTS], ) ] + + def is_unused_checkpoint_weight(self, name: str) -> bool: + # The readout's only tensor, loaded through `secondary_weights`. A + # checkpoint without a weight index lists its file with the backbone's. + return name == "weight" diff --git a/vllm/model_executor/models/deepseek_vl2.py b/vllm/model_executor/models/deepseek_vl2.py index c720f62335..154cdc5747 100644 --- a/vllm/model_executor/models/deepseek_vl2.py +++ b/vllm/model_executor/models/deepseek_vl2.py @@ -4,9 +4,10 @@ # adapted from https://github.com/deepseek-ai/DeepSeek-VL2/blob/faf18023f24b962b32d9f0a2d89e402a8d383a78/deepseek_vl2/models/modeling_deepseek_vl_v2.py """Inference-only Deepseek-VL2 model compatible with HuggingFace weights.""" +import functools import math -from collections.abc import Iterable, Mapping, Sequence -from typing import Annotated, Literal, TypeAlias +from collections.abc import Hashable, Iterable, Mapping, Sequence +from typing import Annotated, Any, Literal, TypeAlias import torch import torch.nn as nn @@ -51,7 +52,12 @@ from vllm.transformers_utils.configs.deepseek_vl2 import ( from vllm.utils.tensor_schema import TensorSchema, TensorShape from vllm.utils.torch_utils import set_default_torch_dtype -from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP +from .interfaces import ( + MultiModalEmbeddings, + SupportsEncoderCudaGraph, + SupportsMultiModal, + SupportsPP, +) from .utils import ( AutoWeightsLoader, WeightsMapper, @@ -319,7 +325,9 @@ class DeepseekVL2MultiModalProcessor( info=DeepseekVL2ProcessingInfo, dummy_inputs=DeepseekVL2DummyInputsBuilder, ) -class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP): +class DeepseekVLV2ForCausalLM( + nn.Module, SupportsMultiModal, SupportsPP, SupportsEncoderCudaGraph +): hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ "language.": "language_model.", @@ -340,6 +348,7 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP): multimodal_config = vllm_config.model_config.multimodal_config self.config = config + self.model_config = vllm_config.model_config self.multimodal_config = multimodal_config self.vision_config = config.vision_config @@ -583,6 +592,186 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP): pixel_values=pixel_values, images_spatial_crop=images_spatial_crop ) + # -- SupportsEncoderCudaGraph protocol methods -- + + @functools.cached_property + def _proj_hw(self) -> tuple[int, int]: + """Return (h, w) of the projector output per tile.""" + h = w = math.ceil( + (self.vision_config.image_size // self.vision_config.patch_size) + / self.projector_config.downsample_ratio + ) + return h, w + + def get_encoder_cudagraph_config(self): + from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphConfig + + return EncoderCudaGraphConfig( + modalities=["image"], + buffer_keys=["pixel_values"], + out_hidden_size=self.projector_config.n_embed, + ) + + def get_encoder_cudagraph_budget_range( + self, vllm_config: "VllmConfig" + ) -> tuple[int, int]: + h, w = self._proj_hw + # Smallest valid image: 1 global + 1×1 local tile → 2 tiles total. + min_budget = h * (w + 1) * 2 + 1 + max_budget = min( + vllm_config.scheduler_config.max_num_batched_tokens, + self.model_config.max_model_len, + ) + return min_budget, max_budget + + def get_encoder_cudagraph_item_specs(self, mm_kwargs: dict[str, Any]): + from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec + + h, w = self._proj_hw + specs = [] + for row in mm_kwargs["images_spatial_crop"].tolist(): + tw, th = int(row[0]), int(row[1]) + if tw == 0 or th == 0: + break + n_tiles = 1 + tw * th + output_tokens = h * (w + 1) + th * h * (tw * w + 1) + 1 + specs.append( + EncoderItemSpec(input_size=n_tiles, output_tokens=output_tokens) + ) + return specs + + def select_encoder_cudagraph_items( + self, mm_kwargs: dict[str, Any], indices: list[int] + ) -> dict[str, Any]: + pixel_values = mm_kwargs["pixel_values"] + images_spatial_crop = mm_kwargs["images_spatial_crop"] + + if not indices: + return { + "pixel_values": pixel_values[:0], + "images_spatial_crop": images_spatial_crop[:0], + } + + # Build per-image tile counts (1 global + tw*th local). + tiles_per_image = [] + for row in images_spatial_crop.tolist(): + tw, th = int(row[0]), int(row[1]) + if tw == 0 or th == 0: + break + tiles_per_image.append(1 + tw * th) + + cum_tiles = [0] + for t in tiles_per_image: + cum_tiles.append(cum_tiles[-1] + t) + + selected_pv = torch.cat( + [pixel_values[cum_tiles[i] : cum_tiles[i + 1]] for i in indices] + ) + return { + "pixel_values": selected_pv, + "images_spatial_crop": images_spatial_crop[list(indices)], + } + + def prepare_encoder_cudagraph_capture_inputs( + self, + token_budget: int, + max_batch_size: int, + max_frames_per_batch: int, + device: torch.device, + dtype: torch.dtype, + path: str = "default", + axis_keys: tuple[Hashable, ...] | None = None, + ): + from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphCaptureInputs + + h, w = self._proj_hw + # Each tile contributes at least h*w output tokens (before newlines). + # Using h*w (not h*(w+1)) guarantees we never under-allocate. + max_tiles = (token_budget + h * w - 1) // (h * w) + dummy_pixel_values = torch.randn( + max_tiles, + 3, + self.vision_config.image_size, + self.vision_config.image_size, + device=device, + dtype=dtype, + ) + return EncoderCudaGraphCaptureInputs( + values={"pixel_values": dummy_pixel_values} + ) + + def prepare_encoder_cudagraph_replay_buffers( + self, + mm_kwargs: dict[str, Any], + max_batch_size: int, + max_frames_per_batch: int, + path: str = "default", + ): + from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers + + return EncoderCudaGraphReplayBuffers( + values={"pixel_values": mm_kwargs["pixel_values"]} + ) + + def encoder_cudagraph_forward( + self, values: dict[str, torch.Tensor], path: str = "default" + ) -> torch.Tensor: + pixel_values = values["pixel_values"] + features = self.vision.forward_features(pixel_values) + return self.projector(features) + + def postprocess_encoder_output( + self, + outputs: dict[str, torch.Tensor], + indices: list[int], + per_item_out_tokens: list[int], + dest: dict[int, torch.Tensor] | list[torch.Tensor | None], + clone: bool = False, + batch_mm_kwargs: dict[str, Any] | None = None, + ) -> None: + """Assemble per-image embeddings from raw per-tile projector output. + + ``outputs["default"]`` has shape ``[N_tiles_captured, hw, n_embed]``. + This method re-applies newlines and the view-separator that the + eager ``_pixel_values_to_embedding`` path inserts in-line. + """ + assert batch_mm_kwargs is not None + output = outputs["default"] + images_spatial_crop = batch_mm_kwargs["images_spatial_crop"] + _, hw, n_dim = output.shape + h = w = int(hw**0.5) + + tile_offset = 0 + for rank, img_idx in enumerate(indices): + tw, th = ( + int(images_spatial_crop[rank][0]), + int(images_spatial_crop[rank][1]), + ) + n_tiles = 1 + tw * th + tiles = output[tile_offset : tile_offset + n_tiles] + tile_offset += n_tiles + + # Global tile: add one newline per row. + global_f = tiles[0].view(h, w, n_dim) + newline_g = self.image_newline.view(1, 1, n_dim).expand(h, 1, n_dim) + global_f = torch.cat([global_f, newline_g], dim=1).view(-1, n_dim) + + # Local tiles: rearrange grid, then add one newline per row. + local_f = tiles[1:].view(th, tw, h, w, n_dim) + local_f = local_f.permute(0, 2, 1, 3, 4).reshape(th * h, tw * w, n_dim) + newline_l = self.image_newline.view(1, 1, n_dim).expand(th * h, 1, n_dim) + local_f = torch.cat([local_f, newline_l], dim=1).view(-1, n_dim) + + if self.global_view_pos == "head": + emb = torch.cat([global_f, self.view_seperator[None], local_f]) + else: + emb = torch.cat([local_f, self.view_seperator[None], global_f]) + + if isinstance(dest, dict): + dest[img_idx] = emb.clone() if clone else emb + else: + dest[rank] = emb.clone() if clone else emb + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: image_input = self._parse_and_validate_image_input(**kwargs) if image_input is None: diff --git a/vllm/model_executor/warmup/deep_gemm_warmup.py b/vllm/model_executor/warmup/deep_gemm_warmup.py index 9f7d7c7e4e..c96f1f09bf 100644 --- a/vllm/model_executor/warmup/deep_gemm_warmup.py +++ b/vllm/model_executor/warmup/deep_gemm_warmup.py @@ -14,7 +14,6 @@ from tqdm import tqdm import vllm.envs as envs from vllm.distributed.parallel_state import ( get_dp_group, - get_ep_group, is_global_first_rank, ) from vllm.forward_context import ForwardContext, override_forward_context @@ -356,7 +355,6 @@ def _get_fp8_fp4_mega_moe_m_values( dg = _import_deep_gemm() gate, experts = module.gate, module.experts uses_mega_gate = _mega_moe_uses_mega_gate(module) - num_max_tokens_per_rank = experts.get_symm_buffer().num_max_tokens_per_rank if module.use_sequence_parallel: max_tokens = cdiv(max_tokens, module.tp_size) @@ -369,14 +367,7 @@ def _get_fp8_fp4_mega_moe_m_values( if uses_mega_gate and hasattr(dg, "get_bf16_mega_gate_config") else {} ) - block_m = dg.get_block_m_for_mega_moe( - get_ep_group().world_size, - experts.num_experts, - num_max_tokens_per_rank, - num_tokens, - experts.top_k, - "fp8xfp4", - ) + block_m = experts.get_block_m(num_tokens) # Keep the largest M: DeepSeek-V4 only calls bf16_mega_gate above a # small token threshold, and a config's largest M is always above it. m_values[(tuple(gate_config.items()), block_m)] = num_tokens diff --git a/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py b/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py index 91ae7c44c2..5c01f4ca2e 100644 --- a/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py +++ b/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py @@ -30,6 +30,7 @@ _DEEPSEEK_V4_SPARSE_MLA_BACKENDS = frozenset( "ROCM_FLASHMLA_SPARSE_DSV4", "FLASHMLA_SPARSE_DSV41", "FLASHINFER_MLA_SPARSE_DSV41", + "ROCM_FLASHMLA_SPARSE_DSV41", "DEEPSEEK_SPARSE_SWA", } ) diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index e7894f319f..f933cffcda 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -80,7 +80,7 @@ from vllm.models.deepseek_v4.amd.mega_moe import ( ) from vllm.models.deepseek_v4.amd.rocm import DeepseekV4ROCMAiterMLAAttention from vllm.platforms import current_platform -from vllm.platforms.rocm import on_gfx950 +from vllm.platforms.rocm import on_gfx950, on_gfx1250 from vllm.sequence import IntermediateTensors from vllm.v1.worker.ubatching import dbo_current_ubatch_token_offset @@ -142,6 +142,11 @@ class DeepseekV4MLP(nn.Module): # B-preshuffle the gate_up_proj weight in place (single weight). if not self._gateup: return + # aiter's gemm_a8w8_blockscale_bpreshuffle currently disabled on gfx1250. + # Leaving _gateup_scale as None -> "not preshuffled" + # forward() falls back to the standard gate_up_proj linear + if on_gfx1250(): + return from vllm.model_executor.layers.quantization.utils.fp8_utils import ( _upcast_e8m0_to_fp32, get_fp8_block_weight_scale, @@ -157,11 +162,13 @@ class DeepseekV4MLP(nn.Module): return if ws.dtype == torch.float8_e8m0fnu: ws = _upcast_e8m0_to_fp32(ws).contiguous() - replace_parameter( - self.gate_up_proj, - "weight", - rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)), - ) + # Skip when the linear kernel already B-preshuffled this weight. + if not getattr(self.gate_up_proj, "aiter_bpreshuffled", False): + replace_parameter( + self.gate_up_proj, + "weight", + rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)), + ) self._gateup_scale = ws def forward(self, x): diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index 3dad09dd08..5353b13ebc 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -23,7 +23,7 @@ from vllm.models.deepseek_v4.sparse_mla import ( DeepseekV4SparseMLAMetadataBuilder, ) from vllm.platforms import current_platform -from vllm.platforms.rocm import _ON_GFX950 +from vllm.platforms.rocm import _ON_GFX950, on_gfx1250 from vllm.triton_utils import tl, triton from vllm.utils.multi_stream_utils import execute_in_parallel from vllm.v1.attention.backend import ( @@ -912,6 +912,11 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): if not rocm_aiter_ops.is_enabled(): return + # aiter's gemm_a8w8_blockscale_bpreshuffle currently disabled on gfx1250. + # Leaving the block scales as None -> "not preshuffled" + # _fused_wqa_wkv_gemm and _o_proj both fall back to the standard linear path + if on_gfx1250(): + return from vllm.model_executor.layers.quantization.utils.fp8_utils import ( _upcast_e8m0_to_fp32, get_fp8_block_weight_scale, @@ -930,12 +935,14 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): return None if ws.dtype == torch.float8_e8m0fnu: ws = _upcast_e8m0_to_fp32(ws).contiguous() - # Shuffle the weight in place (single weight, no unshuffled copy). - replace_parameter( - linear, - "weight", - rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)), - ) + # Shuffle the weight in place (single weight, no unshuffled copy), + # unless the linear kernel already B-preshuffled it at load time. + if not getattr(linear, "aiter_bpreshuffled", False): + replace_parameter( + linear, + "weight", + rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)), + ) return ws self._wqa_wkv_scale = _prep(self.fused_wqa_wkv) @@ -1300,9 +1307,12 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): fuse_inv_rope = self._wo_a_fp8_weight is None rotated = 0 if num_decodes > 0: + fuse_decode_inv_rope = fuse_inv_rope and not on_gfx1250() rotated = self._forward_decode( q=q[:num_decode_tokens], - positions=positions[:num_decode_tokens] if fuse_inv_rope else None, + positions=( + positions[:num_decode_tokens] if fuse_decode_inv_rope else None + ), kv_cache=self_kv_cache, swa_metadata=swa_metadata, attn_metadata=rocm_metadata, diff --git a/vllm/models/deepseek_v4/cpu/model.py b/vllm/models/deepseek_v4/cpu/model.py index 02ea09bcc4..71b9a9810d 100644 --- a/vllm/models/deepseek_v4/cpu/model.py +++ b/vllm/models/deepseek_v4/cpu/model.py @@ -216,14 +216,6 @@ class DeepseekV4MoE(nn.Module): swiglu_limit=self.swiglu_limit, router_logits_dtype=torch.float32, ) - if is_hash_moe: - # apply_monolithic() never sees hash_indices_table via the - # router path (modular-only) -- stash it on RoutedExperts - # directly. object.__setattr__ avoids nn.Module re-registering - # (and duplicating in the state_dict) this alias. - object.__setattr__( - self.experts.routed_experts, "hash_indices_table", self.gate.tid2eid - ) def forward( self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index ef24943f02..71decb5453 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -16,7 +16,6 @@ from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_world_size, ) -from vllm.distributed.eplb.eplb_state import EplbLayerState from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger from vllm.model_executor.kernels.mhc.tilelang import ( @@ -32,10 +31,9 @@ from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe import ( - get_deep_gemm_mega_moe_backend, -) -from vllm.model_executor.layers.fused_moe.router.base_router import ( - eplb_map_to_physical_and_record, + MEGA_MOE_MXFP4, + DeepGemmMegaMoEExperts, + MegaMoEActivation, ) from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( fused_topk_bias, @@ -70,7 +68,6 @@ from vllm.model_executor.models.utils import ( maybe_prefix, spec_decode_needs_target_embed, ) -from vllm.model_executor.utils import set_weight_attrs from vllm.models.common.ops.sequence_parallel import ( sp_all_gather, sp_padding_mask, @@ -84,12 +81,11 @@ from vllm.models.deepseek_v4.nvidia.flashinfer_sparse import ( DeepseekV4FlashInferSM120Attention, ) from vllm.models.deepseek_v4.nvidia.flashmla import DeepseekV4FlashMLAAttention -from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors +from vllm.utils.deep_gemm import bf16_mega_gate from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.registry import AttentionBackendEnum -from vllm.v1.worker.ubatching import dbo_current_ubatch_id from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID @@ -206,326 +202,6 @@ def make_deepseek_v4_expert_params_mapping( ] -class DeepseekV4MegaMoEExperts(nn.Module): - _symm_buffer_cache: dict[ - tuple[int, int, int, int, int, int, int, int, str], object - ] = {} - - def __init__( - self, - vllm_config: VllmConfig, - *, - num_experts: int, - num_local_experts: int, - experts_start_idx: int, - top_k: int, - hidden_size: int, - intermediate_size: int, - num_shared_experts: int = 0, - shared_experts: nn.Module | None = None, - prefix: str = "", - num_logical_experts: int | None = None, - ): - super().__init__() - self.prefix = prefix - self.capture_fn: Callable[[torch.Tensor], None] | None = None - self.num_experts = num_experts - self.num_local_experts = num_local_experts - self.experts_start_idx = experts_start_idx - self.experts_end_idx = experts_start_idx + num_local_experts - self.top_k = top_k - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens - - self.num_logical_experts = ( - num_logical_experts if num_logical_experts is not None else num_experts - ) - - self.eplb_state = EplbLayerState() - - device = torch.get_default_device() - if device.type != current_platform.device_type: # meta under ipc_cache - device = torch.device( - current_platform.device_type, torch.accelerator.current_device_index() - ) - self.backend = get_deep_gemm_mega_moe_backend( - device, hidden_size, intermediate_size - ) - if num_shared_experts and not self.backend.supports_shared_experts( - shared_experts - ): - num_shared_experts = 0 - self.num_shared_experts = num_shared_experts - self.backend.create_weights( - self, - num_local_experts=num_local_experts, - hidden_size=hidden_size, - intermediate_size=intermediate_size, - num_shared_experts=num_shared_experts, - ) - for name, param in self.named_parameters(recurse=False): - shared = name.startswith("shared_") - loader = self.shared_weight_loader if shared else self.weight_loader - set_weight_attrs(param, {"weight_loader": loader}) - self.w13_weight_scale.quant_method = "block" - self.w2_weight_scale.quant_method = "block" - - # Register in the static forward context so the custom-op wrapper - # can look up this module by name from within a torch.compile graph. - compilation_config = vllm_config.compilation_config - if prefix in compilation_config.static_forward_context: - raise ValueError(f"Duplicate layer name: {prefix}") - compilation_config.static_forward_context[prefix] = self - - def _map_global_expert_id(self, expert_id: int) -> list[int]: - """Return local (per-rank) slot offsets where logical expert - `expert_id` should land on this rank. - """ - physical_ids: list[int] = [] - for p in range(self.experts_start_idx, self.experts_end_idx): - if p % self.num_logical_experts == expert_id: - physical_ids.append(p - self.experts_start_idx) - return physical_ids - - def weight_loader( - self, - param: nn.Parameter, - loaded_weight: torch.Tensor, - weight_name: str, - shard_id: str, - expert_id: int, - return_success: bool = False, - ) -> bool | None: - local_expert_ids = self._map_global_expert_id(expert_id) - param_prefix = "w2_" if shard_id == "w2" else "w13_" - if not local_expert_ids or param_prefix not in weight_name: - return False if return_success else None - is_scale = weight_name.endswith("weight_scale") - for local_expert_id in local_expert_ids: - self.backend.load_weight( - param.data[local_expert_id], loaded_weight, shard_id, is_scale - ) - return True if return_success else None - - def load_shared_expert_weight(self, name: str, loaded_weight: torch.Tensor) -> str: - """Load one shared-expert checkpoint tensor; return the parameter name.""" - if ".w2." in name or ".down_proj." in name: - param_name = "shared_l2_weight" - elif ".w1." in name or ".w3." in name: - param_name = "shared_l1_weight" - else: - raise ValueError(f"Unexpected shared-expert tensor: {name}") - if "weight_scale" in name: - param_name += "_scale" - param = getattr(self, param_name) - param.weight_loader(param, loaded_weight, name) - return f"{name.partition('shared_experts.')[0]}experts.{param_name}" - - def shared_weight_loader( - self, param: nn.Parameter, loaded_weight: torch.Tensor, weight_name: str - ) -> None: - shard_id = ( - "w1" if ".w1." in weight_name else "w3" if ".w3." in weight_name else "w2" - ) - is_scale = "weight_scale" in weight_name - self.backend.load_weight(param.data, loaded_weight, shard_id, is_scale) - - @property - def _transformed_l1_weights(self) -> tuple[torch.Tensor, torch.Tensor]: - return self.w13_weight, self.w13_weight_scale - - @property - def _transformed_l2_weights(self) -> tuple[torch.Tensor, torch.Tensor]: - return self.w2_weight, self.w2_weight_scale - - @property - def has_fused_shared_experts(self) -> bool: - return self.num_shared_experts > 0 - - def get_symm_buffer(self): - from vllm.utils.deep_gemm import _import_deep_gemm - - deep_gemm = _import_deep_gemm() - - backend = self.backend - group = get_ep_group().device_group - device = torch.accelerator.current_device_index() - key = ( - id(group), - device, - self.num_experts, - self.max_num_tokens, - self.top_k, - self.hidden_size, - self.intermediate_size, - self.num_shared_experts if self.has_fused_shared_experts else 0, - backend.mma_type, - ) - symm_buffer = self._symm_buffer_cache.get(key) - if symm_buffer is None: - symm_buffer = deep_gemm.get_symm_buffer_for_mega_moe( - group, - self.num_experts, - self.max_num_tokens, - self.top_k, - self.hidden_size, - self.intermediate_size, - num_shared_experts=( - self.num_shared_experts if self.has_fused_shared_experts else 0 - ), - mma_type=backend.mma_type, - ) - self._symm_buffer_cache[key] = symm_buffer - return symm_buffer - - def set_eplb_state( - self, - moe_layer_idx: int, - expert_load_view: torch.Tensor, - logical_to_physical_map: torch.Tensor, - logical_replica_count: torch.Tensor, - ) -> None: - self.eplb_state.set_layer_state( - moe_layer_idx, - expert_load_view, - logical_to_physical_map, - logical_replica_count, - ) - - def get_expert_weights(self) -> list[torch.Tensor]: - def _to_eplb_view(name: str, t: torch.Tensor) -> torch.Tensor: - """Return a (num_local_experts, -1) view with contiguous memory layout.""" - assert t.shape[0] == self.num_local_experts - if t.is_contiguous(): - return t.view(self.num_local_experts, -1) - elif t.dim() == 3 and t.stride(1) == 1 and t.stride(2) == t.shape[1]: - # scales have shape (E, M, N) with memory layout (E, N, M) - back = torch.transpose(t, 1, 2) - assert back.is_contiguous() - return back.view(self.num_local_experts, -1) - - raise AssertionError( - f"DSv4 EPLB {name}: non-contiguous expert tensor with " - f"unexpected layout shape={tuple(t.shape)} " - f"stride={tuple(t.stride())} dtype={t.dtype}" - ) - - return [ - _to_eplb_view("l1_packed", self._transformed_l1_weights[0]), - _to_eplb_view("l1_scale", self._transformed_l1_weights[1]), - _to_eplb_view("l2_weight", self._transformed_l2_weights[0]), - _to_eplb_view("l2_scale", self._transformed_l2_weights[1]), - ] - - def update_expert_map(self) -> None: - pass - - @property - def layer_id(self) -> int: - return extract_layer_index(self.prefix) - - def forward( - self, - hidden_states: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - *, - activation_clamp: float | None, - fast_math: bool = True, - ) -> torch.Tensor: - if hidden_states.shape[0] > self.max_num_tokens: - raise ValueError( - f"DeepSeek V4 MegaMoE got {hidden_states.shape[0]} tokens, " - f"but the symmetric buffer was sized for {self.max_num_tokens}." - ) - y = torch.empty_like(hidden_states, dtype=torch.bfloat16) - - from vllm.utils.deep_gemm import _import_deep_gemm - - deep_gemm = _import_deep_gemm() - - num_tokens = hidden_states.shape[0] - is_padding = None - if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): - is_padding = get_forward_context().is_padding - if is_padding is not None: - is_padding = is_padding[:num_tokens] - - if self.capture_fn is not None: - self.capture_fn(topk_ids) - - # EPLB: map logical expert IDs to physical replicas and record load. - eplb_state = self.eplb_state - if eplb_state.logical_to_physical_map is not None: - assert eplb_state.expert_load_view is not None - assert eplb_state.logical_replica_count is not None - assert eplb_state.should_record_tensor is not None - if is_padding is not None: - topk_ids = torch.where(is_padding.unsqueeze(1), -1, topk_ids) - topk_ids = eplb_map_to_physical_and_record( - topk_ids=topk_ids, - expert_load_view=eplb_state.expert_load_view, - logical_to_physical_map=eplb_state.logical_to_physical_map, - logical_replica_count=eplb_state.logical_replica_count, - record_enabled=eplb_state.should_record_tensor, - num_unpadded_tokens=eplb_state.num_unpadded_tokens_tensors[ - dbo_current_ubatch_id() - ] - if eplb_state.num_unpadded_tokens_tensors is not None - else None, - ) - - backend = self.backend - symm_buffer = self.get_symm_buffer() - shared_x_sf = None - shared_block_m = None - if self.has_fused_shared_experts: - shared_x_sf = symm_buffer.shared_l1_acts_sf - shared_block_m = deep_gemm.get_block_m_for_mega_moe( - get_ep_group().world_size, - self.num_experts, - symm_buffer.num_max_tokens_per_rank, - num_tokens, - self.top_k, - backend.mma_type, - ) - - prepare_megamoe_inputs( - hidden_states, - topk_weights, - topk_ids, - symm_buffer.x[:num_tokens], - symm_buffer.x_sf[:num_tokens], - symm_buffer.topk_idx[:num_tokens], - symm_buffer.topk_weights[:num_tokens], - is_padding=is_padding, - shared_x_sf=shared_x_sf, - shared_block_m=shared_block_m, - hidden_quant=backend.hidden_quant, - ) - - fused = self.has_fused_shared_experts - backend.run_mega_moe( - y=y, - l1_weights=self._transformed_l1_weights, - l2_weights=self._transformed_l2_weights, - symm_buffer=symm_buffer, - shared_l1_weights=( - (self.shared_l1_weight, self.shared_l1_weight_scale) if fused else None - ), - shared_l2_weights=( - (self.shared_l2_weight, self.shared_l2_weight_scale) if fused else None - ), - activation_clamp=activation_clamp, - fast_math=fast_math, - ) - return y - - -DeepseekV4MegaMoEExperts.weight_loader.supports_moe_loading = True # type: ignore[attr-defined] - - class DeepseekV4MoE(nn.Module): def __init__( self, @@ -701,7 +377,8 @@ class DeepseekV4MoE(nn.Module): and (self.use_sequence_parallel or self.tp_size == 1) ) - expert_kwargs: dict[str, typing.Any] = dict( + self.experts = DeepGemmMegaMoEExperts( + vllm_config, num_experts=self.n_physical_experts, num_local_experts=self.n_local_physical_experts, experts_start_idx=self.physical_expert_start, @@ -709,11 +386,14 @@ class DeepseekV4MoE(nn.Module): top_k=self.n_activated_experts, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, - num_shared_experts=(self.n_shared_experts if fuse_shared_experts else 0), - shared_experts=self.shared_experts, + weight_format=MEGA_MOE_MXFP4, + activation=MegaMoEActivation( + "swiglu", + clamp=None if self.swiglu_limit is None else float(self.swiglu_limit), + ), + shared_experts=self.shared_experts if fuse_shared_experts else None, prefix=f"{prefix}.experts", ) - self.experts = DeepseekV4MegaMoEExperts(vllm_config, **expert_kwargs) if self.experts.has_fused_shared_experts: self.shared_experts = None @@ -783,13 +463,9 @@ class DeepseekV4MoE(nn.Module): return self._forward_fused_moe(hidden_states, input_ids) org_shape = hidden_states.shape - mega_backend = self.experts.backend # Small local padded batches favor GateLinear; 128-expert gates cross earlier. gate_threshold = 1 if self.gate.weight.shape[0] == 128 else 16 - if ( - not mega_backend.supports_bf16_mega_gate() - or hidden_states.shape[0] <= gate_threshold - ): + if hidden_states.shape[0] <= gate_threshold: router_logits, _ = self.gate(hidden_states) topk_weights, topk_ids = fused_topk_bias( hidden_states=hidden_states, @@ -825,10 +501,10 @@ class DeepseekV4MoE(nn.Module): ] fix_routing_mask = mega_gate_metadata.hash_token_mask - topk_weights, topk_ids = mega_backend.bf16_mega_gate( - x=hidden_states, - weight=self.gate.weight, - num_topk=self.n_activated_experts, + topk_weights, topk_ids = bf16_mega_gate( + hidden_states, + self.gate.weight, + self.n_activated_experts, scoring_func=self.scoring_func, routed_scaling_factor=self.routed_scaling_factor, ep_rank=self.ep_rank, @@ -840,20 +516,9 @@ class DeepseekV4MoE(nn.Module): fix_routing_mask=fix_routing_mask, unmapped_topk_idx=unmapped_topk_idx, ) - activation_clamp = ( - float(self.swiglu_limit) if self.swiglu_limit is not None else None - ) final_hidden_states = self.experts( - hidden_states, - topk_weights, - topk_ids, - activation_clamp=activation_clamp, + hidden_states, topk_weights, topk_ids, self.shared_experts ) - - if self.shared_experts is not None: - shared_output = self.shared_experts(hidden_states) - final_hidden_states += shared_output - return final_hidden_states.view(org_shape) def _forward_fused_moe( diff --git a/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py b/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py index d97c54a011..a714aa96c5 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py +++ b/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py @@ -52,6 +52,12 @@ class DequantGatherKCacheKernel( data_dim = fp8_dim + bf16_dim * 2 scale_bytes = head_dim // group_size block_size = compile_key.block_size + # V4.1 stores each token's scales right after its data; V4 puts all of + # a block's data rows ahead of its scale rows. + interleaved = bf16_dim == 0 + data_stride = data_dim + scale_bytes if interleaved else data_dim + scale_offset = data_dim if interleaved else block_size * data_dim + scale_stride = data_stride if interleaved else scale_bytes @cute.jit def load_g2s( @@ -274,23 +280,22 @@ class DequantGatherKCacheKernel( offset: Int32, stream: CUstream, ): - # Split k_cache into k_data and k_scale. Each [block_size, head_bytes] - # block is actually a concat of - # [block_size, fp8_dim + bf16_dim * 2] and [block_size, scale_bytes]. + # Split k_cache into k_data [block_size, fp8_dim + bf16_dim * 2] and + # k_scale [block_size, scale_bytes] views of each block. k_data = cute.make_tensor( k_cache.iterator, layout=cute.make_layout( (k_cache.shape[0], block_size, data_dim), - stride=(k_cache.stride[0], data_dim, 1), + stride=(k_cache.stride[0], data_stride, 1), ), ) # k_data_slice: [num_blocks, block_size, (16, data_dim/16)] # s_kdata_slice: [(4, data_dim/16), num_stages] k_scale = cute.make_tensor( - k_cache.iterator + (block_size * data_dim), + k_cache.iterator + scale_offset, layout=cute.make_layout( (k_cache.shape[0], block_size, scale_bytes), - stride=(k_cache.stride[0], scale_bytes, 1), + stride=(k_cache.stride[0], scale_stride, 1), ), ) diff --git a/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py b/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py deleted file mode 100644 index 2d93b0d430..0000000000 --- a/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py +++ /dev/null @@ -1,323 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Triton input-staging kernel for DeepSeek V4 MegaMoE. - -Quantizes hidden states to fp8 with E8M0 group scales and repacks the -routing top-k tensors into the int64/float32 layout that the DeepGEMM -MegaMoE kernels consume. -""" - -import torch - -from vllm.model_executor.layers.quantization.utils.quant_utils import ( - FP8_DTYPE, - QuantKey, - kMxfp8Dynamic, -) -from vllm.triton_utils import tl, triton - -_E8M0_SCALE_DTYPES = (torch.uint8, torch.float8_e8m0fnu) - - -@triton.jit(do_not_specialize=["num_tokens"]) -def _prepare_megamoe_inputs_kernel( - hidden_states, - x_fp8, - x_sf, - shared_x_sf, - topk_ids, - topk_weights, - is_padding, - topk_idx_out, - topk_weights_out, - hidden_stride_m: tl.constexpr, - hidden_stride_k: tl.constexpr, - x_stride_m: tl.constexpr, - x_stride_k: tl.constexpr, - x_sf_stride_m: tl.constexpr, - x_sf_stride_k: tl.constexpr, - shared_x_sf_stride_m: tl.constexpr, - shared_x_sf_stride_k: tl.constexpr, - topk_ids_stride_m: tl.constexpr, - topk_ids_stride_k: tl.constexpr, - topk_weights_stride_m: tl.constexpr, - topk_weights_stride_k: tl.constexpr, - is_padding_stride_m: tl.constexpr, - topk_idx_stride_m: tl.constexpr, - topk_idx_stride_k: tl.constexpr, - topk_weights_out_stride_m: tl.constexpr, - topk_weights_out_stride_k: tl.constexpr, - num_tokens, - hidden_size: tl.constexpr, - top_k: tl.constexpr, - BLOCK_M: tl.constexpr, - BLOCK_K: tl.constexpr, - GROUP_K: tl.constexpr, - BLOCK_TOPK: tl.constexpr, - SHARED_BLOCK_M: tl.constexpr, - USE_UE8M0: tl.constexpr, -) -> None: - token_id = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) - token_mask = token_id < num_tokens - k_block_id = tl.program_id(1) - - k_offsets = k_block_id * BLOCK_K + tl.arange(0, BLOCK_K) - k_mask = token_mask[:, None] & (k_offsets[None, :] < hidden_size) - hidden = tl.load( - hidden_states - + token_id[:, None] * hidden_stride_m - + k_offsets[None, :] * hidden_stride_k, - mask=k_mask, - other=0.0, - ).to(tl.float32) - - num_groups: tl.constexpr = BLOCK_K // GROUP_K - hidden_groups = tl.reshape(tl.abs(hidden), [BLOCK_M, num_groups, GROUP_K]) - amax = tl.max(hidden_groups, axis=2) - amax = tl.maximum(amax, 1.0e-4) - - scale = amax / 448.0 - if USE_UE8M0: - scale_bits = scale.to(tl.uint32, bitcast=True) - scale_exp = ((scale_bits >> 23) & 0xFF) + ((scale_bits & 0x7FFFFF) != 0).to( - tl.uint32 - ) - scale_exp = tl.minimum(tl.maximum(scale_exp, 1), 254) - quant_scale = (scale_exp << 23).to(tl.float32, bitcast=True) - else: - quant_scale = scale - - hidden_groups = tl.reshape(hidden, [BLOCK_M, num_groups, GROUP_K]) - scaled = hidden_groups * (1.0 / quant_scale)[:, :, None] - scaled = tl.reshape(scaled, [BLOCK_M, BLOCK_K]) - fp8 = scaled.to(tl.float8e4nv) - tl.store( - x_fp8 + token_id[:, None] * x_stride_m + k_offsets[None, :] * x_stride_k, - fp8, - mask=k_mask, - ) - - if USE_UE8M0: - scale_offsets = tl.arange(0, num_groups) - packed_scale = tl.sum(scale_exp << (scale_offsets[None, :] * 8), axis=1).to( - tl.int32 - ) - tl.store( - x_sf + token_id * x_sf_stride_m + k_block_id * x_sf_stride_k, - packed_scale, - mask=token_mask, - ) - else: - sf_offset = tl.arange(0, num_groups) - tl.store( - x_sf - + token_id[:, None] * x_sf_stride_m - + k_block_id * x_sf_stride_k - + sf_offset[None, :], - quant_scale, - mask=token_mask[:, None], - ) - - # DeepGEMM's SM100 shared-expert TMA loads require the activation scales - # in an MN-major layout whose row permutation depends on the MegaMoE - # scheduler's runtime BLOCK_M. Write that view while the packed UE8M0 scale - # is already resident, avoiding another kernel and temporary tensor. - if USE_UE8M0 and shared_x_sf is not None: - m_block_id = token_id // SHARED_BLOCK_M - m_in_block = token_id % SHARED_BLOCK_M - aligned_block_m: tl.constexpr = triton.cdiv(SHARED_BLOCK_M, 128) * 128 - transposed_m = ( - (m_in_block // 128) * 128 + (m_in_block % 32) * 4 + (m_in_block % 128) // 32 - ) - shared_row = m_block_id * aligned_block_m + transposed_m - tl.store( - shared_x_sf - + shared_row * shared_x_sf_stride_m - + k_block_id * shared_x_sf_stride_k, - packed_scale, - mask=token_mask, - ) - - if k_block_id == 0: - topk_offsets = tl.arange(0, BLOCK_TOPK) - topk_mask = token_mask[:, None] & (topk_offsets[None, :] < top_k) - token_is_padding = tl.full((BLOCK_M,), False, tl.int1) - if is_padding is not None: - token_is_padding = tl.load( - is_padding + token_id * is_padding_stride_m, - mask=token_mask, - other=True, - ) - - ids = tl.load( - topk_ids - + token_id[:, None] * topk_ids_stride_m - + topk_offsets[None, :] * topk_ids_stride_k, - mask=topk_mask, - other=0, - ).to(tl.int64) - ids = tl.where(token_is_padding[:, None], -1, ids) - tl.store( - topk_idx_out - + token_id[:, None] * topk_idx_stride_m - + topk_offsets[None, :] * topk_idx_stride_k, - ids, - mask=topk_mask, - ) - - weights = tl.load( - topk_weights - + token_id[:, None] * topk_weights_stride_m - + topk_offsets[None, :] * topk_weights_stride_k, - mask=topk_mask, - other=0.0, - ) - weights = tl.where(token_is_padding[:, None], 0.0, weights) - tl.store( - topk_weights_out - + token_id[:, None] * topk_weights_out_stride_m - + topk_offsets[None, :] * topk_weights_out_stride_k, - weights, - mask=topk_mask, - ) - - -def _resolve_hidden_quant(hidden_quant: QuantKey, block_k: int) -> tuple[int, bool]: - """Map a hidden-state ``QuantKey`` to the kernel's ``(GROUP_K, USE_UE8M0)``.""" - scale = hidden_quant.scale - group_shape = scale.group_shape - if ( - hidden_quant.dtype != FP8_DTYPE - or not hidden_quant.symmetric - or hidden_quant.scale2 is not None - or scale.static - or not group_shape.is_per_group() - or block_k % group_shape.col != 0 - or scale.dtype not in (*_E8M0_SCALE_DTYPES, torch.float32) - ): - raise ValueError( - "DeepSeek V4 MegaMoE input staging requires symmetric dynamic fp8 " - f"per-group quantization with a group size dividing {block_k} and " - f"E8M0 or fp32 scales, got {hidden_quant}." - ) - return group_shape.col, scale.dtype in _E8M0_SCALE_DTYPES - - -def prepare_megamoe_inputs( - hidden_states: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - x_fp8: torch.Tensor, - x_sf: torch.Tensor, - topk_idx_out: torch.Tensor, - topk_weights_out: torch.Tensor, - is_padding: torch.Tensor | None = None, - shared_x_sf: torch.Tensor | None = None, - shared_block_m: int | None = None, - hidden_quant: QuantKey = kMxfp8Dynamic, -) -> None: - """Quantize hidden states and repack top-k routing for DeepGEMM MegaMoE. - - Args: - hidden_states: Input activations of shape ``[num_tokens, hidden_dim]``. - topk_weights: Router top-k weights of shape ``[num_tokens, top_k]``. - topk_ids: Router top-k expert ids of shape ``[num_tokens, top_k]``. - x_fp8: Output buffer for the fp8-quantized hidden states. - x_sf: Output buffer for the hidden-state scale factors. - topk_idx_out: Output buffer for the repacked top-k expert ids. - topk_weights_out: Output buffer for the repacked top-k weights. - is_padding: Optional per-token mask; padded tokens are not routed. - shared_x_sf: Optional output buffer for shared-expert scale factors. - Must be given together with ``shared_block_m``. - shared_block_m: Block M used to lay out ``shared_x_sf``. - hidden_quant: Hidden-state quantization scheme. E8M0 scales are packed - four per int32 into ``x_sf``; fp32 scales are stored directly. - - """ - num_tokens, hidden_size = hidden_states.shape - if num_tokens == 0: - return - if hidden_size % 128 != 0: - raise ValueError( - "DeepSeek V4 MegaMoE input staging requires hidden_size to be " - "a multiple of 128." - ) - block_k = 128 - group_k, use_ue8m0 = _resolve_hidden_quant(hidden_quant, block_k) - top_k = topk_ids.shape[1] - if topk_weights.shape != topk_ids.shape: - raise ValueError( - "DeepSeek V4 MegaMoE input staging requires topk_weights and " - "topk_ids to have the same shape." - ) - if (shared_x_sf is None) != (shared_block_m is None): - raise ValueError( - "DeepSeek V4 MegaMoE shared input staging requires both " - "shared_x_sf and shared_block_m." - ) - if shared_x_sf is not None and not use_ue8m0: - raise ValueError( - "DeepSeek V4 MegaMoE shared input staging currently requires " - "UE8M0-packed hidden scales." - ) - if shared_x_sf is not None: - assert shared_block_m is not None - if shared_block_m <= 0: - raise ValueError("MegaMoE shared_block_m must be positive.") - expected_sf_k = hidden_size // 128 - if shared_x_sf.ndim != 2 or shared_x_sf.shape[1] != expected_sf_k: - raise ValueError( - "MegaMoE shared_x_sf must have shape " - f"(*, {expected_sf_k}), got {tuple(shared_x_sf.shape)}." - ) - aligned_block_m = triton.cdiv(shared_block_m, 128) * 128 - required_rows = triton.cdiv(num_tokens, shared_block_m) * aligned_block_m - if shared_x_sf.shape[0] < required_rows: - raise ValueError( - "MegaMoE shared_x_sf has insufficient rows: requires " - f"{required_rows}, got {shared_x_sf.shape[0]}." - ) - - # On GB200, eight-row tiles win from 64 tokens; keep smaller batches untiled. - block_m = 8 if num_tokens >= 64 else 1 - grid = (triton.cdiv(num_tokens, block_m), triton.cdiv(hidden_size, block_k)) - block_topk = triton.next_power_of_2(top_k) - padding_stride_m = is_padding.stride(0) if is_padding is not None else 0 - _prepare_megamoe_inputs_kernel[grid]( - hidden_states, - x_fp8, - x_sf, - shared_x_sf, - topk_ids, - topk_weights, - is_padding, - topk_idx_out, - topk_weights_out, - hidden_states.stride(0), - hidden_states.stride(1), - x_fp8.stride(0), - x_fp8.stride(1), - x_sf.stride(0), - x_sf.stride(1), - shared_x_sf.stride(0) if shared_x_sf is not None else 0, - shared_x_sf.stride(1) if shared_x_sf is not None else 0, - topk_ids.stride(0), - topk_ids.stride(1), - topk_weights.stride(0), - topk_weights.stride(1), - padding_stride_m, - topk_idx_out.stride(0), - topk_idx_out.stride(1), - topk_weights_out.stride(0), - topk_weights_out.stride(1), - num_tokens, - hidden_size, - top_k, - BLOCK_M=block_m, - BLOCK_K=block_k, - GROUP_K=group_k, - BLOCK_TOPK=block_topk, - SHARED_BLOCK_M=shared_block_m or 1, - USE_UE8M0=use_ue8m0, - num_warps=4, - ) diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index 639bea27ad..57366f2453 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -140,13 +140,25 @@ class DeepseekV4MoE(DeepseekV4MoEBase): def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]: """Pick the ROCm sparse-MLA implementation for DeepSeek V4.1.""" + if vllm_config.parallel_config.prefill_context_parallel_size > 1: + raise NotImplementedError( + "DeepSeek-V4.1 PCP currently supports NVIDIA CUDA only." + ) backend = vllm_config.attention_config.backend - if backend is not None and ( - backend is not AttentionBackendEnum.ROCM_FLASHMLA_SPARSE_DSV4 + if backend is not None and backend not in ( + AttentionBackendEnum.ROCM_FLASHMLA_SPARSE_DSV4, + AttentionBackendEnum.ROCM_FLASHMLA_SPARSE_DSV41, ): raise ValueError( f"{backend.name} is not supported for DeepSeek V4.1 on ROCm; " - "use ROCM_FLASHMLA_SPARSE_DSV4." + "use ROCM_FLASHMLA_SPARSE_DSV41." + ) + if backend is AttentionBackendEnum.ROCM_FLASHMLA_SPARSE_DSV4: + # TODO: Remove DSV4 compatibility after a deprecation period. + logger.warning_once( + "ROCM_FLASHMLA_SPARSE_DSV4 is deprecated for DeepSeek V4.1 " + "and will be removed in a future release. " + "Use ROCM_FLASHMLA_SPARSE_DSV41 instead." ) return DeepseekV41ROCMAiterMLAAttention diff --git a/vllm/models/deepseek_v41/amd/mono_decode.py b/vllm/models/deepseek_v41/amd/mono_decode.py index ddeb6b2631..2dd17cc6d2 100644 --- a/vllm/models/deepseek_v41/amd/mono_decode.py +++ b/vllm/models/deepseek_v41/amd/mono_decode.py @@ -39,7 +39,7 @@ from vllm.platforms import current_platform if TYPE_CHECKING: from vllm.models.deepseek_v41.amd.model import DeepseekV4DecoderLayer from vllm.models.deepseek_v41.amd.mono.runner import MonoLayerWeights - from vllm.models.deepseek_v41.amd.rocm import DeepseekV4ROCMAiterSparseSWAMetadata + from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterSparseSWAMetadata from vllm.models.deepseek_v41.sparse_mla import DeepseekV4FlashMLAMetadata logger = init_logger(__name__) @@ -301,7 +301,7 @@ class MonoDecodeLayer: swa, comp = None, None if isinstance(md, dict): swa = cast( - "DeepseekV4ROCMAiterSparseSWAMetadata | None", + "DeepseekV41ROCMAiterSparseSWAMetadata | None", md.get(attn.swa_cache_layer.prefix), ) if attn.compressed_cache_prefix is not None: diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index 5802a54b3f..44b9336632 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -280,7 +280,7 @@ _TopkRagged = tuple[torch.Tensor, torch.Tensor, torch.Tensor] @dataclass -class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata): +class DeepseekV41ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata): decode_swa_ragged_indices: torch.Tensor | None = None decode_swa_ragged_indptr: torch.Tensor | None = None # Built only for the aiter prefill, which reads the SWA cache in place. @@ -303,7 +303,7 @@ class DeepseekV41ROCMAiterMLASparseMetadataBuilder(DeepseekV4SparseMLAMetadataBu return super().get_cudagraph_support(vllm_config, kv_cache_spec) -class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBuilder): +class DeepseekV41ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBuilder): @classmethod def get_cudagraph_support( cls, @@ -348,7 +348,7 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBu common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, replay_start: torch.Tensor | None = None, - ) -> DeepseekV4ROCMAiterSparseSWAMetadata: + ) -> DeepseekV41ROCMAiterSparseSWAMetadata: base = super().build( common_prefix_len=common_prefix_len, common_attn_metadata=common_attn_metadata, @@ -394,7 +394,7 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBu base.decode_swa_width, ) - return DeepseekV4ROCMAiterSparseSWAMetadata( + return DeepseekV41ROCMAiterSparseSWAMetadata( **vars(base), decode_swa_ragged_indices=ragged_indices, decode_swa_ragged_indptr=ragged_indptr, @@ -529,10 +529,10 @@ class DeepseekV41RocmMxfp4Indexer(DeepseekV4Indexer): return q, q_scale, weights -class DeepseekV4ROCMAiterMLASparseBackend(DeepseekV4SparseMLABackend): +class DeepseekV41ROCMAiterMLASparseBackend(DeepseekV4SparseMLABackend): @staticmethod def get_name() -> str: - return "ROCM_FLASHMLA_SPARSE_DSV4" + return "ROCM_FLASHMLA_SPARSE_DSV41" @staticmethod def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: @@ -545,14 +545,14 @@ class DeepseekV4ROCMAiterMLASparseBackend(DeepseekV4SparseMLABackend): class DeepseekV41ROCMAiterSparseSWABackend(DeepseekSparseSWABackend): @staticmethod - def get_builder_cls() -> type["DeepseekV4ROCMAiterSparseSWAMetadataBuilder"]: - return DeepseekV4ROCMAiterSparseSWAMetadataBuilder + def get_builder_cls() -> type["DeepseekV41ROCMAiterSparseSWAMetadataBuilder"]: + return DeepseekV41ROCMAiterSparseSWAMetadataBuilder class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): """ROCm sparse MLA attention layer for DeepSeek V4.1.""" - backend_cls = DeepseekV4ROCMAiterMLASparseBackend + backend_cls = DeepseekV41ROCMAiterMLASparseBackend swa_backend_cls = DeepseekV41ROCMAiterSparseSWABackend _use_aiter_sparse_mla = False @@ -1021,7 +1021,7 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): else None, ) swa_metadata = cast( - DeepseekV4ROCMAiterSparseSWAMetadata | None, + DeepseekV41ROCMAiterSparseSWAMetadata | None, attn_metadata.get(self.swa_cache_layer.prefix), ) assert swa_metadata is not None @@ -1100,7 +1100,7 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): def _decode_topk_ragged( self, - swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata, + swa_metadata: DeepseekV41ROCMAiterSparseSWAMetadata, attn_metadata: DeepseekV4FlashMLAMetadata | None, num_decodes: int, num_decode_tokens: int, @@ -1139,7 +1139,7 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): q: torch.Tensor, positions: torch.Tensor, kv_cache: torch.Tensor | None, - swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata, + swa_metadata: DeepseekV41ROCMAiterSparseSWAMetadata, attn_metadata: DeepseekV4FlashMLAMetadata | None, swa_only: bool, output: torch.Tensor | None, @@ -1233,7 +1233,7 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): swa_k_cache: torch.Tensor, output: torch.Tensor | None, attn_metadata: DeepseekV4FlashMLAMetadata | None, - swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata, + swa_metadata: DeepseekV41ROCMAiterSparseSWAMetadata, mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> None: """Sparse prefill into bf16 ``output``, or into ``mxfp8_out``. @@ -1399,7 +1399,7 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): def _prefill_topk_ragged( self, - swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata, + swa_metadata: DeepseekV41ROCMAiterSparseSWAMetadata, attn_metadata: DeepseekV4FlashMLAMetadata, compressed_k_cache: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: @@ -1449,7 +1449,7 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): swa_k_cache: torch.Tensor, output: torch.Tensor | None, attn_metadata: DeepseekV4FlashMLAMetadata | None, - swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata, + swa_metadata: DeepseekV41ROCMAiterSparseSWAMetadata, mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> None: """Prefill straight off the paged caches, with no bf16 gather. diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index 7009b466a2..76300d57eb 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -27,7 +27,11 @@ from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer from vllm.model_executor.layers.sparse_mqa_indexer import SparseMQAIndexer from vllm.model_executor.model_loader.weight_utils import make_attention_sink from vllm.models.common.ops import fused_q_kv_rmsnorm -from vllm.models.common.ops.sequence_parallel import sp_reduce_scatter +from vllm.models.common.ops.sequence_parallel import ( + sp_all_gather, + sp_reduce_scatter, + sp_shard, +) from vllm.models.deepseek_v41.common.ops import ( MXFP4_BLOCK_SIZE, fused_indexer_q_rope_quant, @@ -49,7 +53,7 @@ from vllm.config import ( get_current_vllm_config, ) from vllm.config.cache import CacheDType -from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.distributed import get_pcp_group, get_tensor_model_parallel_world_size from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase @@ -82,6 +86,10 @@ from vllm.v1.kv_cache_interface import ( logger = init_logger(__name__) +# From this many tokens, a layer with shard_fused_wqa_wkv set projects one token +# shard per TP rank and all-gathers the result. +_SHARD_FUSED_WQA_WKV_MIN_TOKENS = 2048 + def _replace_layer_index(prefix: str, layer_id: int) -> str: """Swap the layer index inside a ``...layers....`` prefix.""" @@ -122,10 +130,10 @@ def _fill_short_context_topk_indices( # Which packed fp8_ds_mla record a V4.1 layer writes. FlashMLA decodes # DeepSeek's V4.1 record -- all 512 dims (RoPE included) as fp8 e4m3 with one -# UE8M0 scale per 32 dims, 512 data bytes + 16 scale bytes per token, pages -# rounded to the kernel's 512 B TMA stride -- only in its SM100 sparse-decode -# kernels. Every other arch keeps the V4 record: 448 fp8 NoPE + 64 bf16 RoPE -# plus 7 UE8M0 scales of 64 dims and a pad byte (584 B, 576 B pages). +# UE8M0 scale per 32 dims, each token's 512 data bytes followed by its 16 scale +# bytes -- only in its SM100 sparse-decode kernels. Every other arch keeps the +# V4 record: 448 fp8 NoPE + 64 bf16 RoPE plus 7 UE8M0 scales of 64 dims and a +# pad byte (584 B, 576 B pages). def _use_v41_mxfp8_kv_record() -> bool: return current_platform.is_device_capability_family(100) @@ -289,6 +297,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): self.layer_id = layer_id self.prefix = prefix # Alias for compatibility with compressor + self.use_pcp = vllm_config.parallel_config.prefill_context_parallel_size > 1 self.hidden_size = config.hidden_size self.n_heads = config.num_attention_heads assert self.n_heads % tp_size == 0 @@ -402,6 +411,8 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): # Set by ``bind_gemm_rs`` when the decoder layer runs sequence # parallel and the fused GEMM + reduce-scatter kernel accepts wo_b. self.gemm_rs: GemmRsAr | None = None + # Set by the NVIDIA decoder layer from _shard_fused_wqa_wkv. + self.shard_fused_wqa_wkv = False # Initialize rotary embedding before the indexer/compressor consume it. self.rotary_emb = build_deepseek_v4_rope( @@ -515,8 +526,8 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): 288 if self.kv_cache_dtype == "nvfp4_ds_mla" else self.swa_bytes_per_token ) # One alignment for every page in the block: the block stride is their - # sum, and 512 satisfies both TMA strides in play (512 for the V4.1 - # fp8 record, 256 for NVFP4). + # sum. V4 needs its 576 B TMA stride; V4.1's is on the SWA block + # stride, and 512 is free (its pages are already 512 B multiples). self.kv_page_alignment = 512 if self.kv_mxfp8 else 576 if self.kv_cache_dtype == "nvfp4_ds_mla" and not self.kv_mxfp8: raise ValueError( @@ -562,6 +573,13 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): block_size=32, packed_bytes_per_token=self.swa_bytes_per_token, packed_page_alignment=self.kv_page_alignment, + # FlashMLA reads the interleaved V4.1 records as rows of the whole + # cache, so the packed block has to hold whole rows of both. + packed_block_stride_alignment=( + math.lcm(self.swa_bytes_per_token, self.compressed_bytes_per_token) + if self.kv_mxfp8 + else None + ), bounded_replay=swa_bounded_replay, ) @@ -646,7 +664,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): if self.backend_cls.get_name() in ( "FLASHMLA_SPARSE_DSV41", "FLASHMLA_MEGA_ATTN_DSV41", - "ROCM_FLASHMLA_SPARSE_DSV4", + "ROCM_FLASHMLA_SPARSE_DSV41", ): from vllm.models.deepseek_v41.common.ops.cache_utils import ( _COMBINE_TOPK_SWA_INDICES_KERNEL, @@ -861,7 +879,10 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): index_q_scale: torch.Tensor | None = None index_weights_out: torch.Tensor | None = None latent: torch.Tensor | None = None - aux_stream = aux_streams[0] if aux_streams is not None else None + # Keep PCP collectives on one stream in the same order on every rank. + aux_stream = ( + aux_streams[0] if aux_streams is not None and not self.use_pcp else None + ) if compressor is not None: # Q projection / KV insertion on the default stream overlaps the @@ -876,6 +897,17 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): else: q = project_query_and_cache_kv() + cache_latent = latent + cache_positions = positions + if self.use_pcp and latent is not None: + assert isinstance(attn_metadata, dict) + swa_metadata = cast( + "DeepseekSparseSWAMetadata", attn_metadata[self.swa_cache_layer.prefix] + ) + assert swa_metadata.cache_positions is not None + cache_latent = get_pcp_group().all_gather(latent, dim=0) + cache_positions = swa_metadata.cache_positions + def prepare_indexer(): if indexer is None: return None, None, None @@ -886,12 +918,16 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): positions, self.indexer_rotary_emb, qr_scale, + cache_latent=cache_latent, + cache_positions=cache_positions, ) if compressor is not None: indexer_result, _ = maybe_execute_in_parallel( prepare_indexer, - lambda: compressor.insert_cache(latent, positions, self.rotary_emb), + lambda: compressor.insert_cache( + cache_latent, cache_positions, self.rotary_emb + ), self.ln_events[0], self.ln_events[1], aux_stream, @@ -914,7 +950,11 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): def _fused_wqa_wkv_gemm(self, hidden_states: torch.Tensor) -> torch.Tensor: # Override point: the ROCm layer preshuffles this weight in place, so # it cannot go through fused_wqa_wkv directly. + num_tokens = hidden_states.shape[0] # MergedColumnParallelLinear returns (output, bias); bias is None. + if self.shard_fused_wqa_wkv and num_tokens >= _SHARD_FUSED_WQA_WKV_MIN_TOKENS: + qr_kv, _ = self.fused_wqa_wkv(sp_shard(hidden_states)) + return sp_all_gather(qr_kv)[:num_tokens] qr_kv, _ = self.fused_wqa_wkv(hidden_states) return qr_kv @@ -1050,6 +1090,25 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): assert positions.dtype == torch.int64 cos_sin_cache = self.rotary_emb.cos_sin_cache cache_dtype = swa_kv_cache.dtype + slot_mapping = swa_metadata.slot_mapping + if self.use_pcp: + cache_kv = get_pcp_group().all_gather(kv.contiguous(), dim=0) + assert swa_metadata.cache_slot_mapping is not None + assert swa_metadata.cache_positions is not None + torch.ops._C.fused_deepseek_v4_kv_rope_insert( + cache_kv, + swa_kv_cache, + swa_metadata.cache_slot_mapping, + swa_metadata.cache_positions, + cos_sin_cache, + swa_metadata.block_size, + self._flashinfer_fp8_kv_scale + if cache_dtype == torch.float8_e4m3fn + else None, + self.kv_mxfp8, + ) + # The fused op now only prepares the rank-local queries. + slot_mapping = slot_mapping[:0] # kv is unchanged; attention reads kv solely via swa_kv_cache. if cache_dtype == torch.uint8: @@ -1071,7 +1130,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): q, kv, swa_kv_cache_2d, - swa_metadata.slot_mapping, + slot_mapping, positions, cos_sin_cache, pad_to, @@ -1099,7 +1158,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): q, kv, swa_kv_cache_3d, - swa_metadata.slot_mapping, + slot_mapping, positions, cos_sin_cache, self.eps, @@ -1115,7 +1174,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): kv, q_fp8, swa_kv_cache_3d, - swa_metadata.slot_mapping, + slot_mapping, positions, cos_sin_cache, self._flashinfer_fp8_kv_scale, @@ -1222,15 +1281,18 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase): tokens_per_state=self.compress_ratio, # The indexer page shares a physical block with the main KV cache # (the block stride is the sum of every page in the group), so it - # has to carry the same padding: the packed record's TMA stride, + # has to carry the same padding: the packed KV pages' alignment, # else 512B for FlashInfer sparse (#44577). alignment=page_alignment, - # DeepGEMM's paged sparse MQA logits address pages by the block - # stride, which packs every layer's page in this layout, and need - # it 512B-aligned; the main KV record read from the same blocks - # needs its own TMA stride, so keep both. + # DeepGEMM's paged sparse logits need a 256B-aligned block stride; + # V4 records keep their TMA stride (V4.1's is on the SWA spec). block_stride_alignment=( - math.lcm(512, page_alignment) + math.lcm( + 256, + 1 + if uses_fp8_ds_mla_layout and _use_v41_mxfp8_kv_record() + else page_alignment, + ) if self.sparse_logits and not self.rocm_mxfp4 else None ), @@ -1473,7 +1535,14 @@ class DeepseekV4Indexer(nn.Module): positions: torch.Tensor, rotary_emb: nn.Module, qr_scale: torch.Tensor | None = None, + *, + cache_latent: torch.Tensor | None = None, + cache_positions: torch.Tensor | None = None, ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + if cache_latent is None: + cache_latent = latent + if cache_positions is None: + cache_positions = positions attn_metadata = get_forward_context().attn_metadata if isinstance(attn_metadata, dict): indexer_metadata = cast(Any, attn_metadata[self.k_cache.prefix]) @@ -1484,7 +1553,7 @@ class DeepseekV4Indexer(nn.Module): # candidates num smaller than topk, every candidate is selected # but we still need to build k cache if self.owns_k: - self._produce_k(latent, positions, rotary_emb) + self._produce_k(cache_latent, cache_positions, rotary_emb) assert self.topk_indices_buffer is not None num_tokens = ( indexer_metadata.num_decode_tokens @@ -1504,7 +1573,7 @@ class DeepseekV4Indexer(nn.Module): if self.owns_k: # K write must land before indexer_op reads the cache # (skip_k_cache_insert=True). - self._produce_k(latent, positions, rotary_emb) + self._produce_k(cache_latent, cache_positions, rotary_emb) return self.forward_q(qr, qr_scale, indexer_weights, positions, rotary_emb) diff --git a/vllm/models/deepseek_v41/common/ops/cache_utils.py b/vllm/models/deepseek_v41/common/ops/cache_utils.py index 000fb03195..b7c8a42fd2 100644 --- a/vllm/models/deepseek_v41/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v41/common/ops/cache_utils.py @@ -36,16 +36,17 @@ from vllm.triton_utils import tl, triton from vllm.utils.import_utils import has_cutedsl from vllm.utils.math_utils import cdiv, next_power_of_2 -# Per-token byte width of the two paged fp8 records. Both put a page's whole -# data region ahead of its whole scale region, so ``k_cache.shape[-1]`` -- the +# Per-token byte width of the paged records; ``k_cache.shape[-1]`` -- the # width the KV-cache spec publishes -- identifies the record: # V4 (584 B): 448 fp8 e4m3 NoPE + 64 bf16 RoPE, then 7 UE8M0 scales of 64 -# dims each plus a pad byte. +# dims each plus a pad byte. A page puts its whole data region +# ahead of its whole scale region. # V4.1 (528 B): all 512 dims as fp8 e4m3 (RoPE included), then 16 UE8M0 # scales of 32 dims each. FlashMLA's ``ModelType::V41``. # V4.1 NVFP4 (288 B): 512 e2m1 values packed two per byte (even element in # the low nibble), then 32 e4m3 scales of 16 dims each. # FlashMLA's ``ModelType::V41_FP4``, compressed cache only. +# The V4.1 records are interleaved: each token's scales follow its own data. V4_BYTES_PER_TOKEN = 584 V41_BYTES_PER_TOKEN = 528 V41_QUANT_BLOCK = 32 @@ -212,10 +213,8 @@ def _quantize_and_insert_k_mxfp8_kernel( # int64: block_idx * block_stride can exceed 2^31 with many KV-cache blocks. cache_block_ptr = k_cache_ptr + block_idx.to(tl.int64) * block_stride - token_data_ptr = cache_block_ptr + pos_in_block * head_dim - token_scale_ptr = ( - cache_block_ptr + cache_block_size * head_dim + pos_in_block * scale_dim - ) + token_data_ptr = cache_block_ptr + pos_in_block * (head_dim + scale_dim) + token_scale_ptr = token_data_ptr + head_dim d = tl.arange(0, head_dim) x = tl.load(k_ptr + pid.to(tl.int64) * head_dim + d).to(tl.float32) @@ -568,10 +567,8 @@ def _dequantize_and_gather_k_mxfp8_kernel( # int64: physical_block_idx * block_stride can exceed 2^31. cache_block_ptr = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride - token_data_ptr = cache_block_ptr + pos_in_block * head_dim - token_scale_ptr = ( - cache_block_ptr + cache_block_size * head_dim + pos_in_block * scale_dim - ) + token_data_ptr = cache_block_ptr + pos_in_block * (head_dim + scale_dim) + token_scale_ptr = token_data_ptr + head_dim x_uint8 = tl.load(token_data_ptr + d) if use_fnuz: @@ -629,10 +626,9 @@ def _dequantize_and_gather_k_nvfp4_kernel( block_table_ptr + batch_idx * max_blocks_per_seq + block_in_seq ) page = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride + token_ptr = page + pos_in_block * (packed_bytes + scale_dim) - packed = tl.load( - page + pos_in_block * packed_bytes + tl.arange(0, packed_bytes) - ) + packed = tl.load(token_ptr + tl.arange(0, packed_bytes)) # Even element in the low nibble, odd in the high nibble. codes = tl.interleave((packed & 0xF).to(tl.int32), (packed >> 4).to(tl.int32)) # e2m1 magnitudes: 0, 0.5, 1, 1.5, 2, 3, 4, 6. @@ -642,12 +638,7 @@ def _dequantize_and_gather_k_nvfp4_kernel( mag = tl.where(mag_code < 2, m * 0.5, (1.0 + m * 0.5) * tl.exp2(e - 1.0)) vals = tl.where(codes >= 8, -mag, mag) - sf = tl.load( - page - + cache_block_size * packed_bytes - + pos_in_block * scale_dim - + tl.arange(0, scale_dim) - ) + sf = tl.load(token_ptr + packed_bytes + tl.arange(0, scale_dim)) scale = sf.to(tl.float8e4nv, bitcast=True).to(tl.float32) tiles = tl.reshape(vals, (scale_dim, quant_block)) dequant = tl.reshape(tiles * tl.reshape(scale, (scale_dim, 1)), (head_dim,)) diff --git a/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py index 2ba9dbea3d..daf2fffc85 100644 --- a/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py +++ b/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py @@ -30,6 +30,8 @@ def fused_save_compress_norm( rms_norm_eps: float, compress_ratio: int, latent_out: torch.Tensor, + prev_rows: torch.Tensor | None = None, + prev_row_indices: torch.Tensor | None = None, ) -> None: """Pool each closed group into a normalized BF16 latent; save FP32 states. @@ -58,6 +60,11 @@ def fused_save_compress_norm( rms_norm_eps: RMSNorm epsilon. compress_ratio: Group size, either 1 or 2. latent_out: BF16 [tokens, 512], written only at valid group boundaries. + prev_rows: FP32 [rows, 1024] rows gathered from all PCP ranks (CR2). A + chunk reads its predecessor here where ``prev_row_indices`` is set, + and the ring is left to ``save_ring_rows``. + prev_row_indices: [num_reqs] row in ``prev_rows`` holding the token + before each chunk's first token, or -1 to read the ring. """ assert compress_ratio in (1, 2) @@ -82,6 +89,11 @@ def fused_save_compress_norm( state_cache.shape[1], ) num_reqs = query_start_loc.numel() - 1 + if prev_rows is not None: + assert prev_row_indices is not None + assert prev_row_indices.numel() >= num_reqs + assert prev_rows.dtype == torch.float32 and prev_rows.shape[1] == 1024 + assert prev_rows.stride(1) == 1 and prev_rows.stride(0) % 16 == 0 else: state_cache = query_start_loc = token_to_req_indices = None state_stride = state_row_stride = state_block = 1 @@ -90,6 +102,8 @@ def fused_save_compress_norm( assert num_tokens <= min(kv_score.shape[0], positions.numel()) if num_tokens == 0: return + has_prev_rows = prev_rows is not None + prev_row_stride = prev_rows.stride(0) if prev_rows is not None else 1 grid = num_reqs + triton.cdiv(num_tokens, compress_ratio) _fused_save_compress_norm_kernel[(grid,)]( kv_score, @@ -100,15 +114,19 @@ def fused_save_compress_norm( token_to_req_indices, rms_norm_weight, latent_out, + prev_rows if has_prev_rows else kv_score, + prev_row_indices if has_prev_rows else slot_mapping, num_tokens, num_reqs, RAW_STRIDE=kv_score.stride(0), + PREV_ROW_STRIDE=prev_row_stride, STATE_STRIDE=state_stride, STATE_ROW_STRIDE=state_row_stride, STATE_BLOCK=state_block, COMPRESS_RATIO=compress_ratio, EPS=rms_norm_eps, JOIN_ROW_PTRS=_JOIN_ROW_PTRS, + HAS_PREV_ROWS=has_prev_rows, num_warps=4, **({"launch_pdl": False} if current_platform.is_cuda() else {}), ) @@ -124,15 +142,19 @@ def _fused_save_compress_norm_kernel( req_ids, norm_weight, latent, + prev_rows, + prev_row_indices, num_tokens, num_reqs, RAW_STRIDE: tl.constexpr, + PREV_ROW_STRIDE: tl.constexpr, STATE_STRIDE: tl.constexpr, STATE_ROW_STRIDE: tl.constexpr, STATE_BLOCK: tl.constexpr, COMPRESS_RATIO: tl.constexpr, EPS: tl.constexpr, JOIN_ROW_PTRS: tl.constexpr, + HAS_PREV_ROWS: tl.constexpr, ): pid = tl.program_id(0) d = tl.arange(0, 512) @@ -152,6 +174,10 @@ def _fused_save_compress_norm_kernel( if (position + 1) % 2 == 0: ring = state + (state_slot // STATE_BLOCK).to(tl.int64) * STATE_STRIDE prev = ring + ((position - 1) % STATE_BLOCK) * STATE_ROW_STRIDE + if HAS_PREV_ROWS: + prev_idx = tl.load(prev_row_indices + pid) + if prev_idx >= 0: + prev = prev_rows + prev_idx.to(tl.int64) * PREV_ROW_STRIDE current = raw + start.to(tl.int64) * RAW_STRIDE if JOIN_ROW_PTRS: # Joining pointers hides row alignment from Triton; restate it. @@ -177,19 +203,20 @@ def _fused_save_compress_norm_kernel( weight_prev + weight_current ) _store_latent(pooled, start, norm_weight, latent, EPS) - num_rows = tl.minimum(end - start, STATE_BLOCK) - for k in tl.range(0, num_rows): - token = end - num_rows + k - slot = tl.load(state_slots + token) - if slot >= 0: - row = ( - state - + (slot // STATE_BLOCK).to(tl.int64) * STATE_STRIDE - + (slot % STATE_BLOCK) * STATE_ROW_STRIDE - ) - src = raw + token.to(tl.int64) * RAW_STRIDE - tl.store(row + d, tl.load(src + d)) - tl.store(row + 512 + d, tl.load(src + 512 + d)) + if not HAS_PREV_ROWS: + num_rows = tl.minimum(end - start, STATE_BLOCK) + for k in tl.range(0, num_rows): + token = end - num_rows + k + slot = tl.load(state_slots + token) + if slot >= 0: + row = ( + state + + (slot // STATE_BLOCK).to(tl.int64) * STATE_STRIDE + + (slot % STATE_BLOCK) * STATE_ROW_STRIDE + ) + src = raw + token.to(tl.int64) * RAW_STRIDE + tl.store(row + d, tl.load(src + d)) + tl.store(row + 512 + d, tl.load(src + 512 + d)) return # Group program: one token (ratio 1) or the token of the pair that ends @@ -212,6 +239,50 @@ def _fused_save_compress_norm_kernel( _store_latent(pooled, t, norm_weight, latent, EPS) +def save_ring_rows( + rows: torch.Tensor, slot_mapping: torch.Tensor, state_cache: torch.Tensor +) -> None: + """Store FP32 [kv, score] rows to their ring slots, skipping slot -1.""" + assert rows.dtype == torch.float32 and rows.shape[1] == 1024 + assert rows.stride(1) == 1 and state_cache.stride(2) == 1 + assert slot_mapping.numel() == rows.shape[0] + if rows.shape[0] == 0: + return + _save_ring_rows_kernel[(rows.shape[0],)]( + rows, + slot_mapping, + state_cache, + ROW_STRIDE=rows.stride(0), + STATE_STRIDE=state_cache.stride(0), + STATE_ROW_STRIDE=state_cache.stride(1), + STATE_BLOCK=state_cache.shape[1], + num_warps=4, + ) + + +@triton.jit +def _save_ring_rows_kernel( + rows, + slot_mapping, + state, + ROW_STRIDE: tl.constexpr, + STATE_STRIDE: tl.constexpr, + STATE_ROW_STRIDE: tl.constexpr, + STATE_BLOCK: tl.constexpr, +): + pid = tl.program_id(0) + slot = tl.load(slot_mapping + pid) + if slot < 0: + return + d = tl.arange(0, 1024) + dst = ( + state + + (slot // STATE_BLOCK).to(tl.int64) * STATE_STRIDE + + (slot % STATE_BLOCK) * STATE_ROW_STRIDE + ) + tl.store(dst + d, tl.load(rows + pid.to(tl.int64) * ROW_STRIDE + d)) + + @triton.jit def _store_latent(pooled, t, norm_weight, latent, EPS: tl.constexpr): d = tl.arange(0, 512) @@ -239,7 +310,8 @@ def rope_quant_insert( UE8M0 scale bytes, including one zero padding scale), 528 B for V4.1 (512 MXFP8 value bytes covering the RoPE dims too, then 16 UE8M0 scales of 32 dims each), or 288 B for V4.1 NVFP4 (256 bytes of e2m1 pairs then 32 - e4m3 scales of 16 dims each), which only the compressed cache uses. + e4m3 scales of 16 dims each), which only the compressed cache uses. The + two V4.1 records keep each token's scales right after its own values. ``bfloat16`` and ``float8_e4m3fn`` are the plain [448 NoPE | 64 RoPE] rows read by FlashInfer, the latter scaled by the per-tensor ``fp8_scale``. """ @@ -390,8 +462,8 @@ def _rope_quant_insert_mxfp8_kernel( rotated = tl.where(rotated == rotated, rotated, 0.0) page = cache + (slot // CACHE_BLOCK).to(tl.int64) * CACHE_STRIDE - values = page + (slot % CACHE_BLOCK) * 512 - scales = page + CACHE_BLOCK * 512 + (slot % CACHE_BLOCK) * 16 + values = page + (slot % CACHE_BLOCK) * 528 + scales = values + 512 quant = tl.reshape(rotated, (16, 32)) amax = tl.maximum(tl.max(tl.abs(quant), 1), 1e-4) @@ -453,12 +525,9 @@ def _rope_quant_insert_nvfp4_kernel( lo, hi = tl.split(tl.reshape(tl.reshape(scaled, (512,)), (256, 2))) page = cache + (slot // CACHE_BLOCK).to(tl.int64) * CACHE_STRIDE - tl.store( - page + (slot % CACHE_BLOCK) * 256 + tl.arange(0, 256), - _fp32x2_to_fp4x2(lo, hi), - ) - scales = page + CACHE_BLOCK * 256 + (slot % CACHE_BLOCK) * 32 - tl.store(scales + tl.arange(0, 32), scale.to(tl.uint8, bitcast=True)) + values = page + (slot % CACHE_BLOCK) * 288 + tl.store(values + tl.arange(0, 256), _fp32x2_to_fp4x2(lo, hi)) + tl.store(values + 256 + tl.arange(0, 32), scale.to(tl.uint8, bitcast=True)) @triton.jit diff --git a/vllm/models/deepseek_v41/compressor.py b/vllm/models/deepseek_v41/compressor.py index 6f56ad8d5b..324123fb54 100644 --- a/vllm/models/deepseek_v41/compressor.py +++ b/vllm/models/deepseek_v41/compressor.py @@ -8,6 +8,7 @@ import torch from torch import nn from vllm.config import VllmConfig, get_current_vllm_config +from vllm.distributed import get_pcp_group from vllm.forward_context import get_forward_context from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.layernorm import RMSNorm @@ -15,6 +16,7 @@ from vllm.model_executor.layers.linear import MergedColumnParallelLinear from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import ( fused_save_compress_norm, rope_quant_insert, + save_ring_rows, ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton @@ -36,6 +38,10 @@ class CompressorBackend(AttentionBackend): def get_name() -> str: return "CompressorBackend" + @classmethod + def supports_pcp(cls) -> bool: + return True + @staticmethod def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: return [MultipleOf(1)] @@ -55,6 +61,12 @@ class CompressorMetadata: slot_mapping: torch.Tensor query_start_loc: torch.Tensor # [num_reqs + 1] token_to_req_indices: torch.Tensor # [num_tokens] + # PCP only: every rank gathers the last row of each rank's local requests. + gather_tokens: torch.Tensor | None = None # [rows] this rank's rows + # [pcp * rows] ring slot of each request's newest row if it is open, else -1. + ring_write_slots: torch.Tensor | None = None + # [num_reqs] gathered row preceding each local chunk, or -1 to use the ring. + prev_row_indices: torch.Tensor | None = None @triton.jit(do_not_specialize=["block_table_stride", "num_actual_tokens", "num_tokens"]) @@ -91,6 +103,7 @@ class CompressorMetadataBuilder(AttentionMetadataBuilder): super().__init__(*args, **kwargs) assert isinstance(self.kv_cache_spec, CircularBufferSpec) self.capacity = self.kv_cache_spec.block_size + self.pcp_size = self.vllm_config.parallel_config.prefill_context_parallel_size max_num_batched_tokens = ( self.vllm_config.scheduler_config.max_num_batched_tokens ) @@ -100,6 +113,8 @@ class CompressorMetadataBuilder(AttentionMetadataBuilder): self.slot_mapping_buffer = torch.empty( max_num_batched_tokens, dtype=torch.int64, device=self.device ) + # The ring keeps only the open row under PCP, too few for spec decoding. + assert self.pcp_size == 1 or self.vllm_config.speculative_config is None def build( self, @@ -107,7 +122,7 @@ class CompressorMetadataBuilder(AttentionMetadataBuilder): common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, ) -> CompressorMetadata: - num_tokens = common_attn_metadata.slot_mapping.numel() + num_tokens = common_attn_metadata.slot_mapping.numel() // self.pcp_size positions = common_attn_metadata.positions assert positions is not None token_to_req_indices = common_attn_metadata.token_to_req_indices( @@ -126,11 +141,66 @@ class CompressorMetadataBuilder(AttentionMetadataBuilder): CAPACITY=self.capacity, BLOCK=256, ) - return CompressorMetadata( + metadata = CompressorMetadata( slot_mapping=slot_mapping, query_start_loc=common_attn_metadata.query_start_loc, token_to_req_indices=token_to_req_indices, ) + if self.pcp_size > 1: + self._link_pcp_chunks(metadata, common_attn_metadata.num_reqs, positions) + return metadata + + def _link_pcp_chunks( + self, metadata: CompressorMetadata, num_reqs: int, positions: torch.Tensor + ) -> None: + """Link ratio-2 groups split across PCP chunks and pick the ring writes. + + Every rank gathers ``(ring block, position)`` of each local request's + last token; the ring block names a request on every rank. Sorted, they + give each local chunk's predecessor row and each request's newest row. + """ + num_tokens = metadata.slot_mapping.numel() + # Equal on every rank, so the gathers line up. + num_rows = min(num_tokens, 2 * self.vllm_config.scheduler_config.max_num_seqs) + num_reqs = min(num_reqs, num_rows) + query_start_loc = metadata.query_start_loc.long() + starts = query_start_loc[:num_reqs].clamp(max=num_tokens - 1) + last = (query_start_loc[1 : num_reqs + 1] - 1).clamp(min=0) + nonempty = query_start_loc[1 : num_reqs + 1] > query_start_loc[:num_reqs] + + def key(slots: torch.Tensor, pos: torch.Tensor) -> torch.Tensor: + return torch.where(slots >= 0, (slots // self.capacity) << 32 | pos, -1) + + local_keys = torch.full((num_rows,), -1, dtype=torch.int64, device=self.device) + local_keys[:num_reqs] = key( + torch.where(nonempty, metadata.slot_mapping[last], -1), positions[last] + ) + sorted_keys, order = get_pcp_group().all_gather(local_keys, dim=0).sort() + + blocks, sorted_positions = sorted_keys >> 32, sorted_keys & 0xFFFFFFFF + newest = torch.ones_like(sorted_keys, dtype=torch.bool) + newest[:-1] = blocks[:-1] != blocks[1:] + is_open = (sorted_keys >= 0) & newest & (sorted_positions % 2 == 0) + ring_write_slots = torch.empty_like(sorted_keys) + ring_write_slots[order] = torch.where( + is_open, blocks * self.capacity + sorted_positions % self.capacity, -1 + ) + + first_positions = positions[starts] + wanted = key(metadata.slot_mapping[starts], first_positions - 1) + found = torch.searchsorted(sorted_keys, wanted).clamp_(max=order.numel() - 1) + has_prev = ( + nonempty + & (wanted >= 0) + & (first_positions % 2 == 1) + & (sorted_keys[found] == wanted) + ) + + gather_tokens = torch.zeros_like(local_keys) + gather_tokens[:num_reqs] = last + metadata.gather_tokens = gather_tokens + metadata.ring_write_slots = ring_write_slots + metadata.prev_row_indices = torch.where(has_prev, order[found], -1) class CompressorStateCache(torch.nn.Module, AttentionLayerBase): @@ -205,6 +275,8 @@ class DeepseekCompressor(nn.Module): self.rotate = rotate self.prefix = prefix self.k_cache_prefix = k_cache_prefix + self.pcp_size = vllm_config.parallel_config.prefill_context_parallel_size + self.pcp_rank = get_pcp_group().rank_in_group if self.pcp_size > 1 else 0 # Ratio 1 pools single tokens, so the checkpoint carries no gate. self.has_gate = compress_ratio > 1 @@ -256,15 +328,20 @@ class DeepseekCompressor(nn.Module): ) -> torch.Tensor | None: """Save states and return the BF16 latent for cache insertion and indexing. - Only valid group-boundary rows are written. + Only valid group-boundary rows are written. Under PCP each rank pools + its own rows; the last row of every rank's chunks is gathered, both to + close groups split across chunks and to keep every rank's ring equal. """ attn_metadata = get_forward_context().attn_metadata if not isinstance(attn_metadata, dict): return None + prev_rows = prev_row_indices = ring_write_slots = None if self.state_cache is None: state_cache = query_start_loc = token_to_req_indices = None slot_mapping = cast(Any, attn_metadata[self.k_cache_prefix]).slot_mapping + if self.pcp_size > 1: + slot_mapping = slot_mapping.chunk(self.pcp_size)[self.pcp_rank] else: state_metadata = cast( CompressorMetadata, attn_metadata[self.state_cache.prefix] @@ -273,6 +350,13 @@ class DeepseekCompressor(nn.Module): slot_mapping = state_metadata.slot_mapping query_start_loc = state_metadata.query_start_loc token_to_req_indices = state_metadata.token_to_req_indices + if self.pcp_size > 1: + assert state_metadata.gather_tokens is not None + prev_rows = get_pcp_group().all_gather( + kv_score[state_metadata.gather_tokens], dim=0 + ) + prev_row_indices = state_metadata.prev_row_indices + ring_write_slots = state_metadata.ring_write_slots latent = torch.empty( kv_score.shape[0], @@ -291,7 +375,12 @@ class DeepseekCompressor(nn.Module): self.rms_norm_eps, self.compress_ratio, latent, + prev_rows=prev_rows, + prev_row_indices=prev_row_indices, ) + if prev_rows is not None: + assert state_cache is not None and ring_write_slots is not None + save_ring_rows(prev_rows, ring_write_slots, state_cache) return latent def insert_cache( diff --git a/vllm/models/deepseek_v41/nvidia/flashmla.py b/vllm/models/deepseek_v41/nvidia/flashmla.py index ce3047ceca..d995058216 100644 --- a/vllm/models/deepseek_v41/nvidia/flashmla.py +++ b/vllm/models/deepseek_v41/nvidia/flashmla.py @@ -42,6 +42,10 @@ class DeepseekSparseSWAFlashMLAMetadataBuilder(DeepseekV41SparseSWAMetadataBuild class DeepseekSparseSWAFlashMLABackend(DeepseekSparseSWABackend): + @classmethod + def supports_pcp(cls) -> bool: + return True + @staticmethod def get_builder_cls() -> type[DeepseekSparseSWAFlashMLAMetadataBuilder]: return DeepseekSparseSWAFlashMLAMetadataBuilder @@ -132,6 +136,14 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention): ) assert swa_metadata is not None + if self.use_pcp: + num_tokens = ( + swa_metadata.num_decode_tokens + swa_metadata.num_prefill_tokens + ) + output[num_tokens:].zero_() + if num_tokens == 0: + return + swa_only = self.compress_ratio == 0 # SWA-only layers (compress_ratio == 0) don't have their own KV cache # allocation; consumers read the kv source's cache, which is bound on diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index 3757809dbc..84b3413eb6 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -14,8 +14,10 @@ from vllm.config import VllmConfig from vllm.config.kernel import MEGA_MOE_BACKENDS, NATIVE_MEGA_MOE_BACKENDS from vllm.distributed import ( get_engram_dp_size, + get_pcp_group, get_pp_group, get_tensor_model_parallel_world_size, + get_tp_group, tensor_model_parallel_all_reduce, ) from vllm.forward_context import ( @@ -203,6 +205,17 @@ def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]: FlashMLA path. """ backend = vllm_config.attention_config.backend + if vllm_config.parallel_config.prefill_context_parallel_size > 1: + if backend not in ( + None, + AttentionBackendEnum.FLASHMLA_SPARSE, + AttentionBackendEnum.FLASHMLA_SPARSE_DSV4, + AttentionBackendEnum.FLASHMLA_SPARSE_DSV41, + ): + raise NotImplementedError( + "DeepSeek-V4.1 PCP requires FlashMLA sparse attention." + ) + return DeepseekV4FlashMLAAttention device_capability = current_platform.get_device_capability() if backend in ( AttentionBackendEnum.FLASHINFER_MLA_SPARSE, @@ -255,6 +268,21 @@ def _use_sequence_parallel(vllm_config: VllmConfig) -> bool: ) +def _shard_fused_wqa_wkv(vllm_config: VllmConfig) -> bool: + """Whether large batches shard fused_wqa_wkv by token and all-gather the output.""" + ca_comm = getattr(get_tp_group().device_communicator, "ca_comm", None) + return ( + current_platform.is_device_capability(90) + and not _use_sequence_parallel(vllm_config) + and vllm_config.lora_config is None + and not vllm_config.parallel_config.use_ubatching + and not envs.VLLM_BATCH_INVARIANT + and ca_comm is not None + and not ca_comm.disabled + and ca_comm.fully_connected + ) + + def maybe_init_gemm_rs(vllm_config: VllmConfig, use_sequence_parallel: bool) -> bool: """Set up the fused ``wo_b`` GEMM + reduce-scatter when opted in. @@ -329,6 +357,7 @@ class DeepseekV4DecoderLayer(nn.Module): # Binds only when wo_b's kernel and shape qualify; otherwise # forward keeps the separate reduce-scatter below. self.attn.bind_gemm_rs() + self.attn.shard_fused_wqa_wkv = _shard_fused_wqa_wkv(vllm_config) self.ffn = DeepseekV4MoE( vllm_config, prefix=f"{prefix}.ffn", @@ -667,6 +696,21 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): self.config = config self.quant_config = quant_config self.parallel_config = vllm_config.parallel_config + if self.parallel_config.prefill_context_parallel_size > 1 and ( + self.parallel_config.pipeline_parallel_size != 1 + or self.parallel_config.data_parallel_size != 1 + or self.parallel_config.decode_context_parallel_size != 1 + or self.parallel_config.use_ubatching + or vllm_config.speculative_config is not None + or vllm_config.kv_transfer_config is not None + or vllm_config.cache_config.enable_prefix_caching + or vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS + ): + raise NotImplementedError( + "DeepSeek-V4.1 PCP currently requires PP=DP=DCP=1, with no " + "microbatching, speculative decoding, KV transfer, prefix " + "caching, or MegaMoE." + ) self.use_native_mega_moe = ( vllm_config.kernel_config.moe_backend in NATIVE_MEGA_MOE_BACKENDS ) @@ -1002,6 +1046,13 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): hidden_states = sp_shard(hidden_states) input_ids = sp_shard(input_ids) + if ( + input_ids is not None + and self.parallel_config.prefill_context_parallel_size > 1 + and not self.parallel_config.enable_expert_parallel + ): + input_ids = get_pcp_group().all_gather(input_ids, dim=0) + residual, post_mix, res_mix = None, None, None pre_mix: torch.Tensor | None = None if not get_pp_group().is_first_rank: diff --git a/vllm/models/deepseek_v41/nvidia/model_state.py b/vllm/models/deepseek_v41/nvidia/model_state.py index e8bbaef64d..89b881be98 100644 --- a/vllm/models/deepseek_v41/nvidia/model_state.py +++ b/vllm/models/deepseek_v41/nvidia/model_state.py @@ -41,7 +41,8 @@ from vllm.v1.worker.utils import AttentionGroup def _gather_lookback_kernel( lookback_ptr, idx_mapping_ptr, - num_computed_tokens_ptr, + query_start_loc_ptr, + positions_ptr, all_token_ids_ptr, all_token_ids_stride, num_reqs, @@ -52,13 +53,16 @@ def _gather_lookback_kernel( batch_idx = tl.program_id(0) in_batch = batch_idx < num_reqs req_state_idx = tl.load(idx_mapping_ptr + batch_idx, mask=in_batch, other=0) - num_computed = tl.load(num_computed_tokens_ptr + req_state_idx) + begin = tl.load(query_start_loc_ptr + batch_idx, mask=in_batch, other=0) + end = tl.load(query_start_loc_ptr + batch_idx + 1, mask=in_batch, other=0) + in_batch = in_batch & (end > begin) + num_computed = tl.load(positions_ptr + begin, mask=in_batch, other=0) offs = tl.arange(0, BLOCK_DEPTH) pos = num_computed - 1 - offs valid = in_batch & (offs < DEPTH) & (pos >= 0) ids = tl.load( - all_token_ids_ptr + req_state_idx * all_token_ids_stride + pos, + all_token_ids_ptr + req_state_idx.to(tl.int64) * all_token_ids_stride + pos, mask=valid, other=-1, ) @@ -200,8 +204,11 @@ class DeepseekV41ModelState(DefaultModelState): self.lookback_token_ids: torch.Tensor | None = None if depth > 0: # Persistent so a captured graph can read it on replay. + max_num_reqs = self.max_num_reqs + if vllm_config.parallel_config.prefill_context_parallel_size > 1: + max_num_reqs *= 2 self.lookback_token_ids = torch.full( - (self.max_num_reqs, depth), -1, dtype=torch.int32, device=device + (max_num_reqs, depth), -1, dtype=torch.int32, device=device ) # Per request state index; batches gather from it (see prepare_attn). @@ -273,7 +280,7 @@ class DeepseekV41ModelState(DefaultModelState): params = new_req_data.sampling_params self._keeps_rows_np[req_index] = ( params is not None and params.prompt_logprobs is not None - ) + ) or self.vllm_config.aux_output_config.enable_return_routed_experts def prepare_inputs( self, input_batch: InputBatch, req_states: RequestState @@ -287,7 +294,8 @@ class DeepseekV41ModelState(DefaultModelState): _gather_lookback_kernel[(window.shape[0],)]( window, input_batch.idx_mapping, - req_states.num_computed_tokens.gpu, + input_batch.query_start_loc, + input_batch.positions, all_token_ids, all_token_ids.stride(0), input_batch.idx_mapping.shape[0], diff --git a/vllm/models/deepseek_v41/sparse_mla.py b/vllm/models/deepseek_v41/sparse_mla.py index 1130eec090..3a3c26f438 100644 --- a/vllm/models/deepseek_v41/sparse_mla.py +++ b/vllm/models/deepseek_v41/sparse_mla.py @@ -9,6 +9,7 @@ import torch from vllm.config import VllmConfig from vllm.config.cache import CacheDType +from vllm.distributed import get_pcp_group from vllm.platforms import current_platform from vllm.platforms.interface import DeviceCapability from vllm.v1.attention.backend import ( @@ -24,6 +25,7 @@ from vllm.v1.attention.backends.mla.sparse_swa import ( _LAYER_TYPE_C1A, _LAYER_TYPE_C2A, _LAYER_TYPE_SWAONLY, + DeepseekSparseSWAMetadata, DeepseekSparseSWAMetadataBuilder, ) from vllm.v1.kv_cache_interface import AttentionSpec @@ -60,6 +62,12 @@ class DeepseekV41SparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) + self.pcp_size = self.vllm_config.parallel_config.prefill_context_parallel_size + self.pcp_rank = get_pcp_group().rank_in_group if self.pcp_size > 1 else 0 + if self.pcp_size > 1: + self.no_replay_start = self.no_replay_start.new_zeros( + 2 * self.vllm_config.scheduler_config.max_num_seqs + 1 + ) compress_ratios = getattr( self.vllm_config.model_config.hf_config, "compress_ratios", None ) or [0] @@ -67,6 +75,24 @@ class DeepseekV41SparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder): deepseek_v41_layer_type(int(ratio)) for ratio in compress_ratios } + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + replay_start: torch.Tensor | None = None, + ) -> DeepseekSparseSWAMetadata: + cm = common_attn_metadata + if self.pcp_size > 1: + local_slots = cm.slot_mapping.chunk(self.pcp_size)[self.pcp_rank] + cm = cm.replace(slot_mapping=local_slots) + metadata = super().build(common_prefix_len, cm, fast_build, replay_start) + if self.pcp_size > 1: + metadata.cache_slot_mapping = common_attn_metadata.slot_mapping + assert cm.positions is not None + metadata.cache_positions = get_pcp_group().all_gather(cm.positions, dim=0) + return metadata + class DeepseekV4SparseMLABackend(AttentionBackend): """DeepSeek-V4.1 sparse-MLA backend base. @@ -155,6 +181,8 @@ class DeepseekV4SparseMLAMetadataBuilder( ) -> None: super().__init__(kv_cache_spec, layer_names, vllm_config, device) self.model_config = vllm_config.model_config + self.pcp_size = vllm_config.parallel_config.prefill_context_parallel_size + self.pcp_rank = get_pcp_group().rank_in_group if self.pcp_size > 1 else 0 # Classify single-token queries (plus num_speculative_tokens via # supports_spec_as_decode=True) as decodes; longer queries go to prefill. self._init_reorder_batch_threshold(1, supports_spec_as_decode=True) @@ -193,9 +221,12 @@ class DeepseekV4SparseMLAMetadataBuilder( slot_mapping = cm.slot_mapping if self.compress_ratio > 1: + local_slots = slot_mapping + if self.pcp_size > 1: + local_slots = slot_mapping.chunk(self.pcp_size)[self.pcp_rank] slot_mapping = get_compressed_slot_mapping( cm.num_actual_tokens, - cm.slot_mapping, + local_slots, cm.query_start_loc, cm.seq_lens, cm.block_table_tensor.clamp_(min=0), @@ -203,6 +234,10 @@ class DeepseekV4SparseMLAMetadataBuilder( self.compress_ratio, out=self.compressed_slot_mapping_buffer, ) + if self.pcp_size > 1: + slot_mapping = get_pcp_group().all_gather( + self.compressed_slot_mapping_buffer[: local_slots.numel()], dim=0 + ) return DeepseekV4FlashMLAMetadata( num_reqs=cm.num_reqs, @@ -223,6 +258,10 @@ class DeepseekV4FlashMLAMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder): class DeepseekV4FlashMLABackend(DeepseekV4SparseMLABackend): + @classmethod + def supports_pcp(cls) -> bool: + return True + @staticmethod def get_name() -> str: return "FLASHMLA_SPARSE_DSV41" @@ -250,6 +289,10 @@ class FlashMLAMegaAttnBackend(DeepseekV4FlashMLABackend): "nvfp4_ds_mla", # V4.1 fp8 SWA cache + NVFP4 compressed cache ] + @classmethod + def supports_pcp(cls) -> bool: + return False + @staticmethod def get_name() -> str: return "FLASHMLA_MEGA_ATTN_DSV41" diff --git a/vllm/models/kimi_k3/amd/kda.py b/vllm/models/kimi_k3/amd/kda.py index 96529734bb..611a0c5d96 100644 --- a/vllm/models/kimi_k3/amd/kda.py +++ b/vllm/models/kimi_k3/amd/kda.py @@ -8,6 +8,7 @@ from einops import rearrange from torch import nn from vllm import _custom_ops as ops +from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import VllmConfig from vllm.distributed import divide @@ -182,16 +183,19 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): # into one kernel, which wants a width-major fp32 conv weight staged at # load time. Everything else keeps the [channel, width] layout. conv_state_dtype, _ = self.get_state_dtype() - decode_conv1d_weight = None - if is_fused_kda_decode_supported( + self.use_aiter_decode = bool(rocm_aiter_ops.is_kda_gluon_decode_available()) + use_hip_decode = not self.use_aiter_decode and is_fused_kda_decode_supported( self.local_num_heads, self.head_dim, self.conv_size, self.num_spec, vllm_config.model_config.dtype, conv_state_dtype, - ): + ) + if use_hip_decode: logger.info_once("Fused KDA decode kernel (conv+KDA+norm) is enabled.") + decode_conv1d_weight = None + if use_hip_decode or self.use_aiter_decode: decode_conv1d_weight = torch.empty( 3, self.conv_size, @@ -257,7 +261,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid") decode_norm_weight = None - if decode_conv1d_weight is not None: + if use_hip_decode: # Upcast once at load time; a BF16 norm weight slows the fused # decode kernel's epilogue. decode_norm_weight = torch.empty( @@ -388,6 +392,66 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): if not is_conv_state_dim_first(): conv_state = conv_state.transpose(-1, -2) + if ( + self.use_aiter_decode + and spec_sequence_masks is None + and m.num_prefills == 0 + and m.num_decodes > 0 + ): + assert non_spec_state_indices_tensor is not None + rocm_aiter_ops.kda_fused_recurrent_packed_decode( + mixed_qkv=mixed_qkv, + g=g1[0], + beta=beta[0], + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=recurrent_state, + ssm_state_indices=non_spec_state_indices_tensor[:num_actual_tokens], + conv_state=conv_state, + conv_weight=self.decode_conv1d_weight, + out_gate=g2[:num_actual_tokens], + norm_weight=self.o_norm.weight, + norm_eps=self.o_norm.eps, + out=core_attn_out[0, :num_actual_tokens], + ) + return + + if ( + self.use_aiter_decode + and spec_sequence_masks is not None + and m.num_prefills == 0 + ): + # Spec-only batch (gdn_attn never mixes spec and plain decodes). + assert spec_state_indices_tensor is not None + assert spec_query_start_loc is not None + q, k, v = ( + rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim) + for x in mixed_qkv.split(self.local_projection_size, dim=-1) + ) + n_spec = q.shape[1] + rocm_aiter_ops.kda_fused_recurrent_spec_decode( + q=q, + k=k, + v=v, + g=g1, + beta=beta, + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=recurrent_state, + cu_seqlens=spec_query_start_loc[: m.num_spec_decodes + 1], + ssm_state_indices=spec_state_indices_tensor, + num_accepted_tokens=num_accepted_tokens, + conv_state=conv_state, + conv_weight=self.decode_conv1d_weight, + out_gate=g2[:n_spec].unsqueeze(0), + norm_weight=self.o_norm.weight, + norm_eps=self.o_norm.eps, + out=core_attn_out[:, :n_spec], + ) + return + if ( self.decode_conv1d_weight is not None and self.decode_norm_weight is not None diff --git a/vllm/models/kimi_k3/nvidia/kda.py b/vllm/models/kimi_k3/nvidia/kda.py index b1956c3720..39e404c086 100644 --- a/vllm/models/kimi_k3/nvidia/kda.py +++ b/vllm/models/kimi_k3/nvidia/kda.py @@ -456,7 +456,7 @@ def _flashinfer_kda_prefill( beta_is_logit=True, seq_order=seq_order, prefill_workspace=prefill_workspace, - backend="cute-dsl", + backend="auto", ) return output, initial_state diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 993422f26a..a49741aec7 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -4,7 +4,7 @@ import math from collections.abc import Hashable, Iterable -from typing import Any, cast +from typing import TYPE_CHECKING, Any, cast import torch from torch import nn @@ -23,8 +23,10 @@ from vllm.model_executor.layers.fused_moe import ( FusedMoEFactory, fused_moe_make_expert_params_mapping, ) -from vllm.model_executor.layers.fused_moe.router.base_router import ( - eplb_map_to_physical_and_record, +from vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe import ( + MEGA_MOE_MXFP4, + DeepGemmMegaMoEExperts, + MegaMoEActivation, ) from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import ( @@ -95,8 +97,6 @@ from vllm.models.common.ops.sequence_parallel import ( sp_reduce_scatter, sp_shard, ) -from vllm.models.deepseek_v4.nvidia.model import DeepseekV4MegaMoEExperts -from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention from vllm.models.kimi_k3.nvidia.latent_moe_runner import ( LatentMoERunner, @@ -115,13 +115,18 @@ from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig from vllm.utils.math_utils import cdiv from vllm.utils.multi_stream_utils import maybe_execute_in_parallel from vllm.utils.torch_utils import aux_stream, is_meta_module -from vllm.v1.worker.ubatching import dbo_current_ubatch_id from ..common.mm_preprocess import ( KimiK3DummyInputsBuilder, KimiK3MultiModalProcessor, KimiK3ProcessingInfo, ) +from .specialized_decode.runtime import specialized_decode_config + +if TYPE_CHECKING: + from vllm.v1.worker.gpu.model_states.interface import ModelState + + from .specialized_decode.runtime.decode import SpecializedDecode logger = init_logger(__name__) @@ -330,143 +335,6 @@ class KimiRoutedOutputTransform(nn.Module): return hidden_states -class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts): - """Kimi K3 adapter for the DeepGEMM MegaMoE kernel.""" - - _kimi_symm_buffer_cache: dict[tuple[object, ...], object] = {} - _synchronized_ep_groups: set[tuple[int, int]] = set() - - def __init__( - self, - *args, - activation: str, - activation_beta: float | None, - activation_linear_beta: float | None, - **kwargs, - ): - super().__init__(*args, **kwargs) - self.activation = activation - self.activation_beta = activation_beta - self.activation_linear_beta = activation_linear_beta - - def synchronize_first_launch(self) -> None: - ep_group = get_ep_group() - device = torch.accelerator.current_device_index() - key = (id(ep_group.cpu_group), device) - if key in self._synchronized_ep_groups: - return - torch.accelerator.synchronize() - torch.distributed.barrier(group=ep_group.cpu_group) - self._synchronized_ep_groups.add(key) - - def get_symm_buffer(self): - from vllm.utils.deep_gemm import _import_deep_gemm - - deep_gemm = _import_deep_gemm() - backend = self.backend - group = get_ep_group().device_group - device = torch.accelerator.current_device_index() - key = ( - id(group), - device, - self.num_experts, - self.max_num_tokens, - self.top_k, - self.hidden_size, - self.intermediate_size, - self.activation, - backend.mma_type, - ) - symm_buffer = self._kimi_symm_buffer_cache.get(key) - if symm_buffer is None: - symm_buffer = deep_gemm.get_symm_buffer_for_mega_moe( - group, - self.num_experts, - self.max_num_tokens, - self.top_k, - self.hidden_size, - self.intermediate_size, - activation=self.activation, - mma_type=backend.mma_type, - ) - self._kimi_symm_buffer_cache[key] = symm_buffer - return symm_buffer - - def forward( - self, - hidden_states: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - *, - activation_clamp: float | None, - fast_math: bool = True, - ) -> torch.Tensor: - self.synchronize_first_launch() - if hidden_states.shape[0] > self.max_num_tokens: - raise ValueError( - f"Kimi K3 MegaMoE got {hidden_states.shape[0]} tokens, " - f"but its symmetric buffer supports {self.max_num_tokens}." - ) - y = torch.empty_like(hidden_states, dtype=torch.bfloat16) - num_tokens = hidden_states.shape[0] - is_padding = None - if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): - is_padding = get_forward_context().is_padding - if is_padding is not None: - is_padding = is_padding[:num_tokens] - - if self.capture_fn is not None: - self.capture_fn(topk_ids) - - eplb_state = self.eplb_state - if eplb_state.logical_to_physical_map is not None: - assert eplb_state.expert_load_view is not None - assert eplb_state.logical_replica_count is not None - assert eplb_state.should_record_tensor is not None - if is_padding is not None: - topk_ids = torch.where(is_padding.unsqueeze(1), -1, topk_ids) - topk_ids = eplb_map_to_physical_and_record( - topk_ids=topk_ids, - expert_load_view=eplb_state.expert_load_view, - logical_to_physical_map=eplb_state.logical_to_physical_map, - logical_replica_count=eplb_state.logical_replica_count, - record_enabled=eplb_state.should_record_tensor, - num_unpadded_tokens=eplb_state.num_unpadded_tokens_tensors[ - dbo_current_ubatch_id() - ] - if eplb_state.num_unpadded_tokens_tensors is not None - else None, - ) - - backend = self.backend - symm_buffer = self.get_symm_buffer() - prepare_megamoe_inputs( - hidden_states, - topk_weights, - topk_ids, - symm_buffer.x[:num_tokens], - symm_buffer.x_sf[:num_tokens], - symm_buffer.topk_idx[:num_tokens], - symm_buffer.topk_weights[:num_tokens], - is_padding=is_padding, - hidden_quant=backend.hidden_quant, - ) - backend.run_mega_moe( - y=y, - l1_weights=self._transformed_l1_weights, - l2_weights=self._transformed_l2_weights, - symm_buffer=symm_buffer, - activation_clamp=activation_clamp, - fast_math=fast_math, - # DeepGEMM names the SiTU gate tanh scale `activation_alpha` and the - # linear/up tanh scale `activation_beta`; beta=0 leaves up untouched. - activation=self.activation, - activation_alpha=self.activation_beta or 1.0, - activation_beta=self.activation_linear_beta or 0.0, - ) - return y - - def make_kimi_k3_mega_moe_expert_params_mapping( num_experts: int, ) -> list[tuple[str, str, int, str]]: @@ -650,7 +518,7 @@ class KimiMoE(nn.Module): f"EP size {ep_size}." ) num_local_experts = num_experts // ep_size - self.experts = KimiK3MegaMoEExperts( + self.experts = DeepGemmMegaMoEExperts( vllm_config, num_experts=num_experts, num_local_experts=num_local_experts, @@ -658,10 +526,13 @@ class KimiMoE(nn.Module): top_k=num_experts_per_token, hidden_size=self.moe_hidden_size, intermediate_size=self.padded_moe_intermediate_size, + activation=MegaMoEActivation( + "situ", + alpha=activation_situ_beta or 1.0, + beta=activation_situ_linear_beta or 0.0, + ), + weight_format=MEGA_MOE_MXFP4, prefix=f"{prefix}.experts", - activation="situ", - activation_beta=activation_situ_beta, - activation_linear_beta=activation_situ_linear_beta, ) else: self.experts = FusedMoEFactory( @@ -777,11 +648,10 @@ class KimiMoE(nn.Module): if self.use_mega_moe: assert self.routed_output_transform is not None assert topk_ids is not None + # The shared expert reads the un-projected hidden states, so it + # stays outside the latent-space MegaMoE launch. final_hidden_states = self.experts( - routed_hidden_states, - router_output, - topk_ids, - activation_clamp=None, + routed_hidden_states, router_output, topk_ids ) # The shared output is folded into the up-projection GEMM's beta-add # epilogue, so combining the two branches costs no extra kernel. @@ -1557,8 +1427,34 @@ class KimiLinearForCausalLM( self.logits_processor = LogitsProcessor( self.config.vocab_size, scale=logit_scale ) + # decode_backend=latency_specialized: a specialized decode executor over + # self.model, built by the model state once the weights are loaded. + self.specialized_decode_config = specialized_decode_config(vllm_config) + self.specialized_decode: SpecializedDecode | None = None + # The final norm is applied in compute_logits unless the specialized decode + # applies it in forward. + self.needs_norm_in_logits = True + + def build_specialized_decode(self) -> None: + from .specialized_decode.runtime.decode import SpecializedDecode + + self.specialized_decode = SpecializedDecode( + self.model, self.specialized_decode_config + ) + # From here on forward() returns normalized hidden states on every step: the + # specialized decode's last AttnRes applies the norm, and forward() applies it + # to each step that falls back to the stock forward (prefills, larger batches, + # PIECEWISE steps). + self.needs_norm_in_logits = False + + def get_model_state_cls(self) -> "type[ModelState]": + from .model_state import KimiK3ModelState + + return KimiK3ModelState def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + if self.specialized_decode is not None: + return self.specialized_decode.embed_input_ids(input_ids) return self.model.embed_input_ids(input_ids) def make_empty_intermediate_tensors( @@ -1577,9 +1473,18 @@ class KimiLinearForCausalLM( inputs_embeds: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: - return self.model( + decode = self.specialized_decode + if decode is not None: + step = decode.plan(positions, intermediate_tensors) + if step is not None: + return decode.run(step, input_ids, inputs_embeds) + hidden_states = self.model( input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs ) + if not self.needs_norm_in_logits: + assert decode is not None + hidden_states = decode.normalize(hidden_states) + return hidden_states @classmethod def get_mamba_state_dtype_from_config( @@ -1638,7 +1543,8 @@ class KimiLinearForCausalLM( ) -> torch.Tensor | None: # The model's final norm is applied here (not at the end of forward) so # that the pre-norm hidden states can be fed to the MTP draft model. - hidden_states = self.model.norm(hidden_states, None) + if self.needs_norm_in_logits: + hidden_states = self.model.norm(hidden_states, None) return self.logits_processor(self.lm_head, hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: @@ -2062,6 +1968,9 @@ class KimiK3ForConditionalGeneration( def compute_logits(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor: return self.language_model.compute_logits(hidden_states) + def get_model_state_cls(self) -> "type[ModelState]": + return self.language_model.get_model_state_cls() + def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs): return self.language_model.mamba_cache.copy_inputs_before_cuda_graphs( input_buffers, **kwargs diff --git a/vllm/models/kimi_k3/nvidia/model_state.py b/vllm/models/kimi_k3/nvidia/model_state.py new file mode 100644 index 0000000000..d1d0e489b5 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/model_state.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kimi K3's model runner state: the hybrid (KDA + MLA) state, plus the lifecycle of +the specialized decode when ``decode_backend=latency_specialized`` selects it.""" + +from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState + + +class KimiK3ModelState(MambaHybridModelState): + """Builds the specialized decode once the weights are loaded and the aux layers + set, and prepares it for every bound KV cache (the profiling cache, then the + real one).""" + + def __init__(self, vllm_config, model, encoder_cache, device): + super().__init__(vllm_config, model, encoder_cache, device) + if hasattr(model, "get_language_model"): + model = model.get_language_model() + self.language_model = model + if model.specialized_decode_config is not None: + model.build_specialized_decode() + + def initialize_kv_cache(self, kv_cache_config, block_tables): + super().initialize_kv_cache(kv_cache_config, block_tables) + decode = self.language_model.specialized_decode + if decode is not None: + decode.prepare_kv() diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/__init__.py b/vllm/models/kimi_k3/nvidia/specialized_decode/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/__init__.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/attn_res/__init__.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/attn_res/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/attn_res/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/attn_res/kernel.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/attn_res/kernel.py new file mode 100644 index 0000000000..5af01ab6fd --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/attn_res/kernel.py @@ -0,0 +1,276 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Seven-CTA cluster AttnRes: one cluster per token, 1024 hidden columns per CTA. + +The delta is a tensor (read after the PDL wait) or ``slots`` Lamport mailbox slots +summed in slot order, rounded to BF16 and re-armed. The final boundary rounds the +mixture to BF16 before its RMSNorm. Prefix and saved blocks must be complete at launch. +""" + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Uint32, Uint64 + +from ..common.mailbox import poll, rearm +from ..common.primitives import ( + alloc, + at, + bf16, + bf16x2, + dot, + fma, + group_barrier, + launch, + load_words, + push_float, + scatter_reduce, + store_words, + wait, + zeros, +) + + +@cute.jit +def round_pair(t, row: cutlass.Constexpr, j: cutlass.Constexpr): + word = bf16x2(t[row, 2 * j], t[row, 2 * j + 1]) + t[row, 2 * j], t[row, 2 * j + 1] = bf16(word, 0), bf16(word, 1) + + +@cute.jit +def unpack(words, out, row: cutlass.Constexpr): + for i in cutlass.range_constexpr(8): + out[row, i] = bf16(words[i // 2], i % 2) + + +def block_ptr(blocks, token, block, col): + """``&blocks[token, block, col]`` from the tensor's strides.""" + return blocks.iterator + token * blocks.stride[0] + block * blocks.stride[1] + col + + +@cute.jit +def rstd(sum_sq, eps): + """Reciprocal RMS of a 7168-wide row from its sum of squares.""" + return cute.math.rsqrt(fma(sum_sq, Float32(1 / 7168), eps), fastmath=True) + + +@cute.jit +def dot8(a, ar: cutlass.Constexpr, b, br: cutlass.Constexpr): + return dot((a[ar, i] for i in range(8)), (b[br, i] for i in range(8)), Float32(0)) + + +@cute.jit +def arm_exchanges(bars, counts): + """Thread 0 arms barrier ``r`` for ``counts[r]`` FP32 pushes from each of 7 CTAs.""" + if cute.arch.thread_idx()[0] == 0: + for r, count in enumerate(counts): + if cutlass.const_expr(count > 0): + cute.arch.mbarrier_init(bars.iterator + r, 1) + cute.arch.mbarrier_arrive_and_expect_tx(bars.iterator + r, 28 * count) + cute.arch.mbarrier_init_fence() + + +class ClusterAttnRes: + def __init__(self, nb, *, slots=1, has_delta=True, write=-1, final=False): + self.nb, self.slots, self.has_delta = nb, slots, has_delta + self.write, self.final = write, final + self.npre = 2 * nb + self.npost = (2 if final else 4) if nb else 1 + self.kpre = 1 << max(self.npre - 1, 0).bit_length() + self.kpost = 1 << max(self.npost - 1, 0).bit_length() + + @cute.jit + def __call__(self, *args): + launch(self.kernel, args, (7, args[1].shape[0], 1), 128, 7) + + @cute.jit + def exchange( + self, values, count: cutlass.Constexpr, kpad: cutlass.Constexpr, bufs, r, ids + ): + """Cluster-sum ``count`` per-thread partials into round ``r``'s total row.""" + wb, peer_buf, total, bars = bufs + peer_buf, total, bar = peer_buf[r, None], total[r, None], bars.iterator + r + rank, _, tid = ids + lane, warp = tid % 32, tid // 32 + reduced, index = scatter_reduce(values, lane, kpad) + if (lane & (32 // kpad - 1)) == 0 and index < count: + wb[warp * 16 + index] = reduced + group_barrier(0) + if tid < count: + value = (wb[tid] + wb[16 + tid]) + (wb[32 + tid] + wb[48 + tid]) + for peer in cutlass.range_constexpr(7): + push_float(peer_buf.iterator + rank * 16 + tid, value, bar, Int32(peer)) + wait(bar) + if tid < count: + sums = [Float32(0), Float32(0)] + for peer in cutlass.range_constexpr(7): + sums[peer % 2] += peer_buf[peer * 16 + tid] + total[tid] = sums[0] + sums[1] + group_barrier(0) + + @cute.kernel + def kernel(self, args): + delta, prefix, blocks, norm, query, out_norm, output, eps, out_eps = args[0:9] + rank, token, _ = cute.arch.block_idx() + tid, _, _ = cute.arch.thread_idx() + bufs = self.setup(cutlass.utils.SmemAllocator()) + cute.arch.cluster_arrive_relaxed() + cute.arch.griddepcontrol_launch_dependents() + state = self.state() + rows, ids = (prefix, blocks, norm, query, out_norm), (rank, token, tid) + self.pre(state, bufs, rows, eps, ids) + self.post(state, bufs, delta, rows, output, eps, out_eps, ids, args[9:]) + cute.arch.griddepcontrol_wait() + + def state(self): + """Carried from pre() to post(); the rest reloads so a GEMM can run between.""" + return cute.make_rmem_tensor((1, 8), Float32), cute.make_rmem_tensor(2, Float32) + + @cute.jit + def setup(self, smem): + warp_buf = alloc(smem, Float32, 64) + peer_buf = alloc(smem, Float32, (3, 112)) + total = alloc(smem, Float32, (3, 16)) + bars = alloc(smem, Uint64, 3, 8) + arm_exchanges(bars, (self.npre, self.npost, int(self.final))) + return warp_buf, peer_buf, total, bars + + @cute.jit + def pre(self, state, bufs, rows, eps: Float32, ids): + """Inputs complete at launch, then the saved-block statistics exchange.""" + prefix, blocks, norm, query, out_norm = rows + rank, token, tid = ids + warp_buf, peer_buf, total, bars = bufs + mixture, scal = state + # Rows: saved blocks, combined weight, gamma, old prefix, new prefix. + v = cute.make_rmem_tensor((self.nb + 4, 8), Float32) + col = (rank * 128 + tid) * 8 + W = self.nb + for source in cutlass.range_constexpr(self.nb): + words = load_words( + block_ptr(blocks, token, source, col), count=4, mutable=True + ) + unpack(words, v, source) + self.weights(v, norm, query, col) + cute.arch.cluster_wait() + maximum, exp_sum = Float32(0), Float32(0) + mixture.fill(0) + if cutlass.const_expr(self.nb > 0): + values = zeros(self.kpre, Float32) + for source in cutlass.range_constexpr(self.nb): + values[source] = dot8(v, source, v, source) + values[self.nb + source] = dot8(v, source, v, W) + self.exchange(values, self.npre, self.kpre, bufs, 0, ids) + logits = cute.make_rmem_tensor(self.nb, Float32) + for b in cutlass.range_constexpr(self.nb): + logits[b] = total[0, self.nb + b] * rstd(total[0, b], eps) + maximum = logits[0] + for source in cutlass.range_constexpr(1, self.nb): + maximum = cute.arch.fmax(maximum, logits[source]) + for source in cutlass.range_constexpr(self.nb): + e = cute.math.exp(logits[source] - maximum, fastmath=True) + exp_sum += e + for i in cutlass.range_constexpr(8): + mixture[0, i] = fma(e, v[source, i], mixture[0, i]) + scal[0], scal[1] = maximum, exp_sum + + @cute.jit + def weights(self, v, norm, query, col): + if cutlass.const_expr(self.nb > 0): + nw = load_words(norm.iterator + col, count=4) + qw = load_words(query.iterator + col, count=4) + for i in cutlass.range_constexpr(8): + v[self.nb, i] = bf16(nw[i // 2], i % 2) * bf16(qw[i // 2], i % 2) + + @cute.jit + def post( + self, state, bufs, delta, rows, output, eps: Float32, out_eps: Float32, ids, aux + ): + """Add the delta, update prefix (and its ``aux`` copy) and saved block, write + the normalized output.""" + prefix, blocks, norm, query, out_norm = rows + rank, token, tid = ids + warp_buf, peer_buf, total, bars = bufs + mixture, scal = state + maximum, exp_sum = scal[0], scal[1] + v = cute.make_rmem_tensor((self.nb + 4, 8), Float32) + col = (rank * 128 + tid) * 8 + W, G, OLD, NEW = self.nb, self.nb + 1, self.nb + 2, self.nb + 3 + self.weights(v, norm, query, col) + unpack(load_words(out_norm.iterator + col, count=4), v, G) + if cutlass.const_expr(self.has_delta): + unpack(load_words(at(prefix, token, col), count=4, mutable=True), v, OLD) + cur = cute.make_rmem_tensor((1, 8), Float32) + if cutlass.const_expr(self.slots == 0): + cute.arch.griddepcontrol_wait() + unpack(load_words(at(delta, token, col), count=4, mutable=True), cur, 0) + else: + cur.fill(0) + base = delta.iterator + token * 7168 + col + got = poll(base, delta.shape[1] * 7168, self.slots) + for slot in cutlass.range_constexpr(self.slots): + rearm(base + slot * delta.shape[1] * 7168) + part = cute.make_rmem_tensor((1, 8), Float32) + unpack(tuple(got[slot, k] for k in range(4)), part, 0) + for i in cutlass.range_constexpr(8): + cur[0, i] = cur[0, i] + part[0, i] + if cutlass.const_expr(self.slots > 1): + for j in cutlass.range_constexpr(4): # all-reduce output boundary + round_pair(cur, 0, j) + words = cute.make_rmem_tensor(4, Uint32) + for j in cutlass.range_constexpr(4): + lo, hi = cur[0, 2 * j], cur[0, 2 * j + 1] + if cutlass.const_expr(self.has_delta): + lo, hi = v[OLD, 2 * j] + lo, v[OLD, 2 * j + 1] + hi + words[j] = bf16x2(lo, hi) + unpack(words, v, NEW) + values = zeros(self.kpost, Float32) + values[0] = dot8(v, NEW, v, NEW) + if cutlass.const_expr(self.nb > 0): + values[1] = dot8(v, NEW, v, W) + if cutlass.const_expr(not self.final): + values[2] = dot8(mixture, 0, v, NEW) + values[3] = dot8(mixture, 0, mixture, 0) + self.exchange(values, self.npost, self.kpost, bufs, 1, ids) + row = (words[0], words[1], words[2], words[3]) + store_words(at(prefix, token, col), row) + for stream_out in aux: + store_words(at(stream_out, token, col), row) + if cutlass.const_expr(self.write >= 0): + store_words(block_ptr(blocks, token, self.write, col), row) + result = cute.make_rmem_tensor((1, 8), Float32) + if cutlass.const_expr(self.nb == 0): + scale = rstd(total[1, 0], out_eps) + for i in cutlass.range_constexpr(8): + result[0, i] = v[NEW, i] * scale * v[G, i] + else: + ss, dot = total[1, 0], total[1, 1] + logit = dot * rstd(ss, eps) + top = cute.arch.fmax(maximum, logit) + alpha = cute.math.exp(maximum - top, fastmath=True) + ep = cute.math.exp(logit - top, fastmath=True) + denominator = fma(alpha, exp_sum, ep) + if cutlass.const_expr(self.final): + inverse = Float32(1) / denominator + for j in cutlass.range_constexpr(4): + for i in (2 * j, 2 * j + 1): + result[0, i] = fma( + alpha * inverse, mixture[0, i], ep * inverse * v[NEW, i] + ) + round_pair(result, 0, j) + fin = cute.make_rmem_tensor(1, Float32) + fin[0] = dot8(result, 0, result, 0) + self.exchange(fin, 1, 1, bufs, 2, ids) + scale = rstd(total[2, 0], out_eps) + for i in cutlass.range_constexpr(8): + result[0, i] = result[0, i] * scale * v[G, i] + else: + a2, t3 = alpha * alpha, total[1, 3] + ae2 = Float32(2) * alpha * ep + quad = fma(a2, t3, fma(ae2, total[1, 2], ep * ep * ss)) + scale = rstd(quad, out_eps * denominator * denominator) + ca, cp = alpha * scale, ep * scale + for i in cutlass.range_constexpr(8): + result[0, i] = fma(ca, mixture[0, i], cp * v[NEW, i]) * v[G, i] + for j in cutlass.range_constexpr(4): + words[j] = bf16x2(result[0, 2 * j], result[0, 2 * j + 1]) + store_words(at(output, token, col), tuple(words[j] for j in range(4))) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/attn_res/prepared.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/attn_res/prepared.py new file mode 100644 index 0000000000..36c65a77c0 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/attn_res/prepared.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Prepared seven-CTA AttnRes consumers for ordinary or Lamport mailbox input.""" + +from functools import cache +from typing import Any + +import cutlass +import cutlass.cute as cute +from cutlass.cute.runtime import make_fake_stream + +from ..common.host import ( + architecture, + dynamic_fake, + dynamic_view, + fake, + stream, + to_cute, +) +from .kernel import ClusterAttnRes + + +@cache +def compile_consumer(nb, delta, write, final, slots, cap, arch, aux=False): + row, weight, eps = dynamic_fake(2), fake((7168,), (1,)), cutlass.Float32(1e-6) + source = fake((slots, cap, 7168), (cap * 7168, 7168, 1)) if slots else row + kernel = ClusterAttnRes(nb, slots=slots, has_delta=delta, write=write, final=final) + args = (source, row, dynamic_fake(3), weight, weight, weight, row, eps, eps) + args += (row,) * aux # the updated prefix row, also stored to an aux stream + return cute.compile(kernel, *args, make_fake_stream(), options=f"--gpu-arch={arch}") + + +class PreparedAttnRes: + """Ordinary-input code plus the Lamport variant of its bound mailbox ring.""" + + def __init__(self, weights, key, eps): + self.device = weights[0].device + self.views = tuple(to_cute(w.reshape(-1)) for w in weights) + self.key, self.eps = key, tuple(map(cutlass.Float32, eps)) + self.code = self._compile(0, 0) + self.ring: Any = None + self.aux: Any = None + + def _compile(self, slots, cap, aux=False): + arch = architecture(self.device) + return compile_consumer(*self.key, slots, cap, arch, aux) + + @classmethod + def for_target(cls, target, *, final=False, has_delta=True): + if final: + norm, query = target.output_attn_res_norm, target.output_attn_res_proj + output, nb, write = target.norm, target.num_attn_res_blocks, -1 + else: + norm = target.self_attention_res_norm + query, output = target.self_attention_res_proj, target.input_layernorm + nb = target.prev_valid_blocks + write = target.block_write_idx if target.is_block_write_layer else -1 + weights = (norm.weight, query.weight, output.weight) + eps = (norm.variance_epsilon, output.variance_epsilon) + return cls(weights, (nb, has_delta, write, final), eps) + + def bind_ring(self, mailbox): + """Lamport input: sum the [slots, capacity, 7168] mailbox rows, then re-arm.""" + shape = tuple(mailbox.shape[:2]) + self.ring = self._compile(*shape), to_cute(mailbox), shape + + def bind_aux(self): + """Variants that also store the updated prefix rows (a stock aux stream).""" + ring = None + if self.ring is not None: + slots, cap = self.ring[2] + ring = self._compile(slots, cap, True) + self.aux = self._compile(0, 0, True), ring + + def run_ring(self, prefix, blocks, out, aux=None): + code = self.ring[0] if aux is None else self.aux[1] + return self._launch(code, self.ring[1], prefix, blocks, out, aux) + + def run(self, delta, prefix, blocks, out, aux=None): + code = self.code if aux is None else self.aux[0] + return self._launch(code, dynamic_view(delta), prefix, blocks, out, aux) + + def _launch(self, code, source, prefix, blocks, out, aux=None): + rows = (prefix, blocks, out) + (() if aux is None else (aux,)) + prefix, blocks, view, *aux = map(dynamic_view, rows) + args = *self.views, view, *self.eps, *aux + code(source, prefix, blocks, *args, stream(self.device)) + return out diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/__init__.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/embedding.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/embedding.py new file mode 100644 index 0000000000..eebf56b41a --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/embedding.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Local CuTe vocabulary gather plus the ordinary framework TP all-reduce.""" + +from functools import cache + +import cutlass as c +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import make_fake_stream + +from vllm.distributed import get_tp_group + +from .host import architecture, fake +from .primitives import launch + + +class EmbeddingGather: + @cute.jit + def __call__(self, *args): + launch(self.kernel, args, (7, args[0].shape[0], 1), 128, pdl=False) + + @cute.kernel + def kernel(self, args): + ids, weight, out, start, vocab = args[0:5] + lane, _, _ = cute.arch.thread_idx() + tile, token, _ = cute.arch.block_idx() + token_id = c.Int64(ids[token]) + row, column = token_id - start, tile * 1024 + lane * 8 + values = cute.make_rmem_tensor((8,), c.BFloat16) + values.fill(c.BFloat16(0)) + if (token_id >= 0) & (token_id < vocab) & (row >= 0) & (row < weight.shape[0]): + source = weight.iterator + row * 7168 + column + cute.autovec_copy(cute.make_tensor(source, cute.make_layout((8,))), values) + destination = out.iterator + token * 7168 + column + cute.autovec_copy(values, cute.make_tensor(destination, cute.make_layout((8,)))) + + +@cache +def compile_gather(dtype, shard, arch): + batch = cute.sym_int32() + args: tuple = (fake((batch,), None, dtype, dtype.itemsize), fake((shard, 7168))) + args += (fake((batch, 7168)), c.Int64(0), c.Int64(0)) + stream = make_fake_stream(use_tvm_ffi_env_stream=True) + options = f"--gpu-arch={arch} --enable-tvm-ffi" + return cute.compile(EmbeddingGather(), *args, stream, options=options) + + +class PreparedEmbedding: + def __init__(self, embedding, tp, rank): + weight = embedding.weight + # The gather moves 7 x 1024 BF16 columns of contiguous rows. + if weight.dim() != 2 or weight.shape[1] != 7168 or not weight.is_contiguous(): + raise ValueError( + f"embedding shard must be a contiguous [V, 7168] tensor, got " + f"{tuple(weight.shape)}" + ) + if weight.dtype != torch.bfloat16: + raise ValueError(f"embedding shard must be BF16, got {weight.dtype}") + self.group, self.weight = get_tp_group(), weight.detach() + self.start, self.vocab = rank * weight.shape[0], embedding.org_vocab_size + with torch.accelerator.device_index(weight.device.index): + arch, ids = architecture(weight.device), (torch.int32, torch.int64) + self.kernels = {t: compile_gather(t, weight.shape[0], arch) for t in ids} + + def run(self, ids, out): + if ids.dim() != 1 or ids.dtype not in self.kernels: + raise ValueError( + f"token ids must be 1-D int32/int64, got {ids.dtype} {tuple(ids.shape)}" + ) + if out.shape != (ids.shape[0], 7168) or not out.is_contiguous(): + raise ValueError( + f"embedding output must be contiguous [{ids.shape[0]}, 7168]" + ) + self.kernels[ids.dtype](ids, self.weight, out, self.start, self.vocab) + reduced = self.group.all_reduce(out) + if reduced is not out: + out.copy_(reduced) + return out diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/host.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/host.py new file mode 100644 index 0000000000..2833c3a37d --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/host.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Host-side helpers for the prepared CuTe providers.""" + +import ctypes +import math + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import make_fake_tensor + +from vllm.models.kimi_k3.nvidia.ops.cute_dsl.latent_moe_tail.primitives import to_cute + +_NAMES = ("BFloat16", "Float32", "Uint8", "Uint32", "Int32", "Int64") +_DTYPES = {getattr(torch, name.lower()): getattr(cutlass, name) for name in _NAMES} + + +def layout(tensor): + return tuple(tensor.shape), tuple(tensor.stride()), tensor.dtype + + +def stream(device): + return cuda.CUstream(torch.cuda.current_stream(device).cuda_stream) + + +def architecture(device): + major, minor = torch.cuda.get_device_capability(device) + return f"sm_{major}{minor}a" + + +def fake(shape, stride=None, dtype=torch.bfloat16, align=16): + """Row-major unless ``stride`` is given.""" + stride = stride or tuple(math.prod(shape[i + 1 :]) for i in range(len(shape))) + dtype = _DTYPES.get(dtype, dtype) + return make_fake_tensor(dtype, shape, stride, assumed_align=align) + + +def dynamic_fake(rank, dtype=torch.bfloat16): + shape = tuple(cute.sym_int() for _ in range(rank)) + return fake(shape, (*(cute.sym_int(64) for _ in range(rank - 1)), 1), dtype) + + +def dynamic_view(tensor, align=16): + return to_cute(tensor, align).mark_layout_dynamic(leading_dim=tensor.ndim - 1) + + +def tensor_maps(specs, device): + """TMA descriptors (128 B each) in one device tensor. + Each spec: (pointer, dtype, dims, strides, box, swizzle, l2 promotion). + + The bytes are uploaded with a copy on the current stream: the new tensor may + reuse a block freed by a kernel still pending on that stream (e.g. a warmup + launch's output), which a copy outside the stream would race with.""" + encoded = bytearray() + with torch.accelerator.device_index(torch.device(device).index): + for pointer, dtype, dims, strides, box, swizzle, l2 in specs: + u64 = ([cuda.cuuint64_t(v) for v in values] for values in (dims, strides)) + u32 = ([cuda.cuuint32_t(v) for v in box], [cuda.cuuint32_t(1)] * len(dims)) + interleave = cuda.CUtensorMapInterleave.CU_TENSOR_MAP_INTERLEAVE_NONE + fill = cuda.CUtensorMapFloatOOBfill.CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE + status, tensor_map = cuda.cuTensorMapEncodeTiled( + dtype, len(dims), pointer, *u64, *u32, interleave, swizzle, l2, fill + ) + if status != cuda.CUresult.CUDA_SUCCESS: + raise RuntimeError(f"cuTensorMapEncodeTiled failed: {status}") + encoded += (ctypes.c_ubyte * 128).from_address(tensor_map.getPtr()) + return torch.frombuffer(encoded, dtype=torch.uint8).to(device) + + +def _checked(result, operation="CUDA driver call"): + """The values of a cuda-python driver result, raising on a failed status.""" + status, *values = result + if status != cuda.CUresult.CUDA_SUCCESS: + raise RuntimeError(f"{operation} failed: {status}") + return tuple(values) + + +class PreparedLaunchAttributes: + def __init__(self, compiled, executor): + self.compiled, self.executor = compiled, executor + + def __call__(self, *args, **kwargs): + return self.compiled(*args, **kwargs) + + +def resolve_kernel(compiled, device): + executor = compiled.to(device) + library = compiled.library + (count,) = _checked(cuda.cuLibraryGetKernelCount(library)) + if count != 1: # every prepared module compiles one kernel + raise RuntimeError( + f"expected one kernel in the compiled library, found {count}" + ) + (kernels,) = _checked(cuda.cuLibraryEnumerateKernels(count, library)) + (function,) = _checked(cuda.cuKernelGetFunction(kernels[0])) + return executor, function + + +def check_residency(compiled, device_index: int, grid: int, threads: int, what: str): + """Persistent grids whose CTAs wait on one another (not only within a cluster) must + all be resident at once: raise when the device cannot hold ``grid`` CTAs of + ``compiled`` (``threads`` per CTA, its dynamic shared memory) at the same time.""" + _, function = resolve_kernel(compiled, device_index) + attribute = cuda.CUfunction_attribute + max_dynamic = attribute.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES + (smem,) = _checked(cuda.cuFuncGetAttribute(max_dynamic, function)) + (per_sm,) = _checked( + cuda.cuOccupancyMaxActiveBlocksPerMultiprocessor(function, threads, smem) + ) + sms = torch.cuda.get_device_properties(device_index).multi_processor_count + if grid > per_sm * sms: + raise RuntimeError( + f"{what}: {grid} CTAs cannot be co-resident ({per_sm} per SM x {sms} SMs)" + ) + + +def prepare_shared_carveout(compiled, device_index: int, value: int): + """Restore ``value`` (CuTe's launcher drops preferred_smem_carveout) as a driver + function attribute; call once before capture.""" + executor, function = resolve_kernel(compiled, device_index) + attribute = cuda.CUfunction_attribute + attribute = attribute.CU_FUNC_ATTRIBUTE_PREFERRED_SHARED_MEMORY_CARVEOUT + _checked(cuda.cuFuncSetAttribute(function, attribute, value)) + return PreparedLaunchAttributes(compiled, executor) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/mailbox.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/mailbox.py new file mode 100644 index 0000000000..75da9e77c7 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/mailbox.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Lamport mailboxes: BF16 rows published in 16-byte fragments; a reader polls until no +fragment holds the -0.0 sentinel, then re-arms.""" + +import cutlass +import cutlass.cute as cute +from cutlass import Uint32 + +from .primitives import load_words, published, store_relaxed + +SENTINEL = 0x80000000 + + +def rearm(pointer): + store_relaxed(pointer, (Uint32(SENTINEL),) * 4) + + +@cute.jit +def fetch(values, i: cutlass.Constexpr[int], pointer): + """Load fragment ``i`` into ``values[i]``; returns its dirty flag at bit ``i``.""" + got = load_words(pointer, count=4, volatile=True) + for k in cutlass.range_constexpr(4): + values[i, k] = got[k] + return Uint32(not published(got)) << i + + +@cute.jit +def poll(base, stride, count: cutlass.Constexpr[int]): + """Lamport-poll ``count`` 16-byte fragments at base + i*stride (elements); all + loads issue before any check and only dirty fragments re-poll.""" + values = cute.make_rmem_tensor((count, 4), Uint32) + pending = Uint32(0) + for i in cutlass.range_constexpr(count): + pending = pending | fetch(values, i, base + i * stride) + while pending != 0: + again = Uint32(0) + for i in cutlass.range_constexpr(count): + if (pending >> i) & 1: + again = again | fetch(values, i, base + i * stride) + pending = again + return values diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/policy.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/policy.py new file mode 100644 index 0000000000..02938c1067 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/policy.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Per-batch policy tables.""" + + +def batch_table(spec): + """Per-batch table of integers ("1-2:96,3-8:120"), tuples ("16:64.16.4.8") or names + ("16-64:trtllm"); later entries override earlier ones.""" + table = {} + for item in filter(None, spec.split(",")): + span, value = item.split(":") + lo, _, hi = span.partition("-") + if "." in value: + value = tuple(map(int, value.split("."))) + elif value.lstrip("-").isdigit(): + value = int(value) + table.update({b: value for b in range(int(lo), int(hi or lo) + 1)}) + return table diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/prefetch.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/prefetch.py new file mode 100644 index 0000000000..2db36e42bb --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/prefetch.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""L2 weight prefetch: each warp walks a persistent device table of (address, bytes), +spreading 16-KiB chunks over the grid in table order. The owner retains the weights +and table through graph replay; activations and mutable caches may not be ranges +before their producer completes.""" + +import functools + +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import make_fake_stream + +from .host import fake +from .primitives import launch, prefetch_bulk, thread + + +class L2Prefetch: + def __init__(self, grid, carveout): + self.grid, self.carveout = grid, carveout + + @cute.jit + def __call__(self, *args): + options = dict(pdl=False, preferred_smem_carveout=self.carveout) + launch(self.kernel, args, self.grid, 32, **options) + + @cute.kernel + def kernel(self, args): + (ranges,) = args + cute.arch.griddepcontrol_launch_dependents() + tid, _, _, block = thread() + worker = cutlass.Int64(block * 32 + tid) + workers = cutlass.Int64(self.grid * 32) + phase, row = cutlass.Int64(0), cutlass.Int32(0) + while row < ranges.shape[0]: + base, nbytes = ranges[row, 0], ranges[row, 1] + if (base != 0) & (nbytes > 0): + chunks = (nbytes + 16384 - 1) // 16384 + index = worker - phase % workers + if index < 0: + index += workers + while index < chunks: + offset = index * 16384 + size = nbytes - offset + if size > 16384: + size = cutlass.Int64(16384) + size = size & cutlass.Int64(-16) + if size > 0: + address = (base + offset) & cutlass.Int64(-16) + prefetch_bulk(address, size.to(cutlass.Int32)) + index += workers + phase += chunks + row += 1 + # Early exit must not let a later grid bypass its producer. + cute.arch.griddepcontrol_wait() + + +@functools.cache +def compile_kernel(count, grid, carveout, device_index): + ranges = fake((count, 2), None, torch.int64) + with torch.accelerator.device_index(device_index): + return cute.compile(L2Prefetch(grid, carveout), ranges, make_fake_stream()) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py new file mode 100644 index 0000000000..66b1d6e0f0 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py @@ -0,0 +1,928 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inline-PTX, memory, synchronization and math helpers of the decode kernels.""" + +import math +from functools import partial, wraps + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Float64, Int32, Int64, Uint8, Uint16, Uint32, Uint64 +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T, dsl_user_op +from cutlass.experimental import primitives as nvvm + + +def ptx(assembly, constraints, operands=(), result=None, *, effects, loc=None, ip=None): + """One verbatim asm block; a tuple ``result`` (even a 1-tuple) returns a tuple.""" + if isinstance(result, tuple) and len(result) == 1: + args = (assembly, constraints, operands, result[0]) + return (ptx(*args, effects=effects, loc=loc, ip=ip),) + if result is None: + ir_type = None + elif isinstance(result, tuple): + ir_type = llvm.StructType.get_literal([dtype.mlir_type for dtype in result]) + else: + ir_type = result.mlir_type + value = llvm.inline_asm( + ir_type, + [operand.ir_value(loc=loc, ip=ip) for operand in operands], + assembly, + constraints, + has_side_effects=effects, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + if result is None: + return None + if isinstance(result, tuple): + return tuple( + dtype(llvm.extractvalue(dtype.mlir_type, value, [i], loc=loc, ip=ip)) + for i, dtype in enumerate(result) + ) + return result(value) + + +def asm(result=None, *, pure=False): + """``f(*args) -> (code, constraints, operands[, result, pure])`` as one inline-PTX + op. Pointer operands pass their address: 32-bit for "r" (shared), 64-bit for "l".""" + + def wrap(fn): + @dsl_user_op + @wraps(fn) + def op(*args, loc=None, ip=None, **kwargs): + code, constraints, operands, *rest = fn(*args, **kwargs) + res, fx = (*rest, pure)[:2] if rest else (result, pure) + kinds = [c for c in constraints.split(",") if c[:1] not in ("=", "~", "")] + values = [] + for kind, v in zip(kinds, operands): + if hasattr(v, "toint"): + v = v.toint(loc=loc, ip=ip) + v = Int32(v) if kind == "r" else v + values.append(v) + return ptx(code, constraints, values, res, effects=not fx, loc=loc, ip=ip) + + return op + + return wrap + + +def _vector(count, first=0): + """(".vN", "{$first,...}") of an N-word access; ("", "$first") for one word.""" + names = ",".join(f"${first + i}" for i in range(count)) + return (f".v{count}", "{" + names + "}") if count > 1 else ("", names) + + +div_rn = partial(cute.math.div, rounding=cute.math.RoundingMode.NEAREST_EVEN) + + +@asm(Float32, pure=True) +def add_rn(a, b): + return "add.rn.f32 $0,$1,$2;", "=f,f,f", (Float32(a), Float32(b)) + + +@asm(Float32, pure=True) +def routing_add(a: Float32, b: Float32): + return "add.ftz.f32 $0,$1,$2;", "=f,f,f", (a, b) + + +@asm(Float32, pure=True) +def fma(a, b, c, *, ftz=False): + return f"fma.rn{'.ftz' if ftz else ''}.f32 $0, $1, $2, $3;", "=f,f,f,f", (a, b, c) + + +def dot(a, b, acc=None, **kwargs): + """``fma`` chain over ``zip(a, b)`` from ``acc``, or from ``a[0] * b[0]``.""" + for x, y in zip(a, b): + acc = x * y if acc is None else fma(x, y, acc, **kwargs) + return acc + + +@asm(Float32, pure=True) +def fma_bf16(a, b, c, *, high=False): + half = "h" if high else "l" + code = "{ .reg .b16 al,ah,bl,bh; mov.b32 {al,ah},$1; mov.b32 {bl,bh},$2; " + code += f"fma.rn.f32.bf16 $0,a{half},b{half},$3; }}" + return code, "=f,r,r,f", (Uint32(a), Uint32(b), Float32(c)) + + +@asm(Uint32, pure=True) +def round_pair_square(word): + return "mul.rn.bf16x2 $0,$1,$1;", "=r,r", (Uint32(word),) + + +@asm(Float32, pure=True) +def inv_rms(value, eps, *, m8=False): + """rsqrt(value / 3584 + eps); ``m8``: the mean as a reciprocal product refined by + its FMA residual.""" + mean = "div.rn.f32 r,$1,0f45600000;" + if m8: + mean = "mul.f32 q,$1,0f39924925; neg.f32 r,q; fma.rn.f32 r,r,0f45600000,$1; " + mean += "fma.rn.f32 r,r,0f39924925,q;" + code = f"{{ .reg .f32 {'q,' * m8}r; {mean} add.f32 r,r,$2; " + code += "rsqrt.approx.ftz.f32 $0,r; }" + return code, "=f,f,f", (Float32(value), Float32(eps)) + + +@asm(Float32, pure=True) +def query_norm_rsqrt(total: Float32, eps: Float32): + code = "{ .reg .f32 mean; div.rn.f32 mean,$1,0f44c00000; " + code += "add.f32 mean,mean,$2; rsqrt.approx.f32 $0,mean; }" + return code, "=f,f,f", (total, eps) + + +@asm(Float64, pure=True) +def reciprocal64(value: Float64): + step = "fma.rn.f64 e,n,r,0d3ff0000000000000; fma.rn.f64 {},r,e,r; " + code = "{ .reg .f64 r,e,n; neg.f64 n,$1; rcp.approx.ftz.f64 r,$1; " + code += step.format("r") * 2 + step.format("$0") + "}" + return code, "=d,d", (value,) + + +@asm(Float32, pure=True) +def sigmoid_reciprocal(denominator: Float32, *, fast_math): + """1 / (1 + exp(-x)) from its denominator: an approximate divide with ``fast_math``, + else the approximate reciprocal refined by two IEEE FMAs (out-of-range: div.rn).""" + if fast_math: + return "div.approx.ftz.f32 $0, 0f3f800000, $1;", "=f,f", (denominator,) + code = ( + "{\n.reg .f32 r,error,neg_error; .reg .b32 bits; .reg .pred regular;\n" + "mov.b32 bits,$1; add.u32 bits,bits,0x01800000; and.b32 bits,bits,0x7f800000;\n" + "setp.gt.u32 regular,bits,0x01ffffff; rcp.approx.ftz.f32 r,$1;\n" + "fma.rn.f32 error,$1,r,0fbf800000; neg.f32 neg_error,error;\n" + "fma.rn.f32 $0,r,neg_error,r; @!regular div.rn.f32 $0,0f3f800000,$1;\n}" + ) + return code, "=f,f", (denominator,) + + +@cute.jit +def sigmoid(value: Float32, fast_math: cutlass.Constexpr = None): + """1 / (1 + exp(-x)); correctly rounded when ``fast_math`` is None.""" + if cutlass.const_expr(fast_math is None): + result = div_rn(Float32(1), Float32(1) + cute.math.exp(-value, fastmath=False)) + else: + denominator = Float32(1) + cute.math.exp(-value, fastmath=True) + result = sigmoid_reciprocal(denominator, fast_math=fast_math) + return result + + +@asm(Float32, pure=True) +def routing_sigmoid(value: Float32): + code = "{ .reg .f32 h,t; mul.ftz.f32 h,$1,0f3f000000; tanh.approx.f32 t,h; " + code += "fma.rn.ftz.f32 $0,t,0f3f000000,0f3f000000; }" + return code, "=f,f", (value,) + + +@asm(Float32, pure=True) +def routing_weight(score: Float32, total: Float32, scale: Float32, *, renorm=True): + code = "mul.ftz.f32 $0,$3,$1;" + if renorm: + code = ( + "{ .reg .f32 n,d; mul.ftz.f32 n,$3,$1; " + "add.ftz.f32 d,0f1e3ce508,$2; div.approx.ftz.f32 $0,n,d; }" + ) + return code, "=f,f,f,f", (score, total, scale) + + +@asm(Float32, pure=True) +def situ(g: Float32, u: Float32, beta_lb: Float32, inv_beta: Float32, inv_lb: Float32): + code = "{ .reg .f32 tg,tu,ex,rc,z; mul.f32 tg,$1,$4; tanh.approx.f32 tg,tg; " + code += "mul.f32 tu,$2,$5; tanh.approx.f32 tu,tu; mul.f32 z,$1,0fbfb8aa3b; " + code += "ex2.approx.ftz.f32 ex,z; add.f32 ex,ex,0f3f800000; " + code += "rcp.approx.ftz.f32 rc,ex; mul.f32 z,$3,tg; mul.f32 z,z,rc; " + code += "mul.f32 $0,z,tu; }" + return code, "=f,f,f,f,f,f", (g, u, beta_lb, inv_beta, inv_lb) + + +@cute.jit +def shared_situ(gate, up, beta, inverse_beta, linear_beta, inverse_linear_beta): + denominator = Float32(1) + cute.math.exp(-gate, fastmath=False) + sigmoid = Float32(0) + if denominator != Float32(float("inf")): + sigmoid = Float32(reciprocal64(Float64(denominator))) + g = (beta * cute.math.tanh(gate * inverse_beta, fastmath=False)) * sigmoid + u = up + if linear_beta > 0: + divided = Float32(Float64(up) * inverse_linear_beta) + u = linear_beta * cute.math.tanh(divided, fastmath=False) + return g * u + + +def bits_float(bits): + return Float32(llvm.bitcast(T.f32(), Uint32(bits).ir_value())) + + +def float_bits(value: Float32) -> Int32: + return Int32(llvm.bitcast(T.i32(), value.ir_value())) + + +def bf16(word, high=False): + """FP32 of the low (or high) BF16 half of ``word``.""" + return bits_float(word & 0xFFFF0000 if high else word << 16) + + +@asm((Float32, Float32), pure=True) +def unpack_bf16(word: Uint32): + return "shl.b32 $0,$2,16; and.b32 $1,$2,0xffff0000;", "=f,=f,r", (word,) + + +@asm((Float32, Float32), pure=True) +def fma2(a, b, c): + """Lanewise a * b + c of FP32 pairs in one packed instruction (IEEE per lane).""" + code = "{ .reg .b64 a,b,c,d; mov.b64 a,{$2,$3}; mov.b64 b,{$4,$5}; " + code += "mov.b64 c,{$6,$7}; fma.rn.f32x2 d,a,b,c; mov.b64 {$0,$1},d; }" + return code, "=f,=f,f,f,f,f,f,f", (*a, *b, *c) + + +@asm((Float32, Float32), pure=True) +def mul2(a, b): + """Lanewise a * b of FP32 pairs in one packed instruction (IEEE per lane).""" + code = "{ .reg .b64 a,b,d; mov.b64 a,{$2,$3}; mov.b64 b,{$4,$5}; " + code += "mul.rn.f32x2 d,a,b; mov.b64 {$0,$1},d; }" + return code, "=f,=f,f,f,f,f", (*a, *b) + + +@asm(Uint16, pure=True) +def gate_bf16_bits(value: Float32): + return "cvt.rn.bf16.f32 $0, $1;", "=h,f", (value,) + + +@asm(Uint32, pure=True) +def pack_pair(low, high, *, policy="plain"): + """BF16x2 of (low, high); ``policy=None`` is the bare conversion.""" + code = "cvt.rn.bf16x2.f32 $0,$2,$1; " + if policy == "stage": + # NaN halves become canonical, so no word is the 0xffffffff staging sentinel. + code += ( + "and.b32 bits,$0,0x7fff; setp.gt.u32 p,bits,0x7f80; " + "@p and.b32 $0,$0,0xffff0000; @p or.b32 $0,$0,0x7fc0; " + "and.b32 bits,$0,0x7fff0000; setp.gt.u32 p,bits,0x7f800000; " + "@p and.b32 $0,$0,0xffff; @p or.b32 $0,$0,0x7fc00000; " + ) + if policy == "latent": + # Negative BF16 zero is the downstream mailbox's unpublished sentinel. + code += ( + "and.b32 bits,$0,0xffff; setp.eq.u32 p,bits,0x8000; " + "@p and.b32 $0,$0,0xffff0000; shr.u32 bits,$0,16; " + "setp.eq.u32 p,bits,0x8000; @p and.b32 $0,$0,0xffff; " + ) + if policy is not None: + code = "{ .reg .b32 bits; .reg .pred p; " + code + "}" + return code, "=r,f,f", (Float32(low), Float32(high)) + + +bf16x2 = partial(pack_pair, policy=None) + + +@cute.jit +def no_neg_zero(word): + """``word`` with each negative-zero BF16 half made +0.""" + result = Uint32(word) + if (result & 0xFFFF) == 0x8000: + result &= Uint32(0xFFFF0000) + if (result >> 16) == 0x8000: + result &= Uint32(0xFFFF) + return result + + +def published(words, ready=True, *, sentinel=0x80000000): + """``ready`` and no word is the mailbox ``sentinel``.""" + for word in words: + ready &= word != Uint32(sentinel) + return ready + + +@cute.jit +def front_dirty(word): + return ((word & 0xFFFF) == 0xFFFF) | ((word >> 16) == 0xFFFF) + + +def fp8_bits(value: Float32): + code = ( + "{ .reg .b16 pair; cvt.rn.satfinite.e4m3x2.f32 pair, 0f00000000, $1; " + "cvt.u32.u16 $0, pair; }" + ) + return Uint8(ptx(code, "=r,f", (value,), Uint32, effects=False)) + + +@asm() +def store_fp8_pair(pointer, low: Float32, high: Float32): + """FP8 E4M3 (satfinite) of ``low`` and ``high`` as the two bytes at ``pointer``.""" + code = "{ .reg .b16 pair; cvt.rn.satfinite.e4m3x2.f32 pair,$2,$1; " + code += "st.global.b16 [$0],pair; }" + return code, "l,f,f,~{memory}", (pointer, low, high) + + +@asm(Uint32, pure=True) +def pack_fp8(values): + code = "{ .reg .b16 lo,hi; cvt.rn.satfinite.e4m3x2.f32 lo,$2,$1; " + code += "cvt.rn.satfinite.e4m3x2.f32 hi,$4,$3; mov.b32 $0,{lo,hi}; }" + return code, "=r,f,f,f,f", values + + +@cute.jit +def mx_exponent(amax: Float32, floor: cutlass.Constexpr[bool] = False): + bits = Uint32(float_bits(amax)) + e = Int32(bits >> 23) + result = Int32(0) + if cutlass.const_expr(floor): + if e == 255: + result = 254 + elif e > 8: + result = e - 8 + else: + if e != 0: + result = min(max(e - 8 + Int32(bits & 8388607 > 6291456), 0), 254) + return result + + +@cute.jit +def mx_rescale(scale: Int32): + e = 254 - scale + return bits_float(Uint32(e) << 23) if e >= 1 else Float32(0) + + +@cute.jit +def warp_sum(value: Float32, lanes: cutlass.Constexpr[int] = 32): + for shift in cutlass.range_constexpr(lanes.bit_length() - 2, -1, -1): + value += cute.arch.shuffle_sync_bfly(value, 1 << shift) + return value + + +@asm(Uint32, pure=True) +def redux(value, op="min"): + return f"redux.sync.{op}.u32 $0,$1,0xffffffff;", "=r,r", (Uint32(value),) + + +def warp_max_abs(value): + return bits_float(redux(float_bits(abs(value)), "max")) + + +@asm((Float32,) * 4, pure=True) +def mma(a, b, accumulator, *, fp8=False, bf16=False): + """m16n8k16 (FP8: k32) mma.sync.""" + kind = "e4m3" if fp8 else "bf16" if bf16 else "f16" + code = f"mma.sync.aligned.m16n8k{32 if fp8 else 16}.row.col.f32.{kind}.{kind}.f32 " + code += "{$0,$1,$2,$3},{$4,$5,$6,$7},{$8,$9},{$10,$11,$12,$13};" + return code, "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", (*a, *b, *accumulator) + + +def mma_into(acc, a, b, **kwargs): + """``acc += a @ b`` with one ``mma``; ``acc`` is a four-element register view.""" + result = mma(a, b, tuple(acc[k] for k in range(4)), **kwargs) + for k in range(4): + acc[k] = result[k] + + +@asm() +def ldmatrix(pointer, *, trans=False, count=4): + """``count`` 8x8 b16 shared matrices; lanes 8i..8i+7 address matrix i's rows.""" + names = ",".join(f"${i}" for i in range(count)) + code = f"ldmatrix.sync.aligned.m8n8.x{count}{'.trans' * trans}.shared.b16 " + code += f"{{{names}}},[${count}];" + constraints = ",".join(["=r"] * count + ["r"]) + return code, constraints, (pointer,), (Uint32,) * count + + +def split(total, i, n): + """(first, count) of part ``i`` of ``total`` items split ``n`` ways.""" + first = total * i // n + return first, total * (i + 1) // n - first + + +def thread(): + """(thread, warp, lane, CTA) indices.""" + tid, _, _ = cute.arch.thread_idx() + cta, _, _ = cute.arch.block_idx() + return tid, tid // 32, tid % 32, cta + + +def shared_address(pointer): + return Uint32(pointer.toint()) + + +def at(tensor, *coord): + """Pointer to ``tensor[coord]``.""" + return tensor.iterator + cute.crd2idx(coord, tensor.layout) + + +def view(base, offset, dtype, shape, stride=None): + """``dtype`` tensor at byte ``offset`` of pointer ``base``.""" + pointer = cute.recast_ptr(base + offset, dtype=dtype) + return cute.make_tensor(pointer, cute.make_layout(shape, stride=stride)) + + +def zeros(shape, dtype): + tensor = cute.make_rmem_tensor(shape, dtype) + tensor.fill(0) + return tensor + + +def alloc(allocator, dtype, shape, align=16): + """Row-major ``shape`` tensor from a ``SmemAllocator``.""" + stride = None + if isinstance(shape, tuple): + stride = tuple(math.prod(shape[i + 1 :]) for i in range(len(shape))) + layout = cute.make_layout(shape, stride=stride) + return allocator.allocate_tensor(dtype, layout, align) + + +@asm() +def load_words( + pointer, + *, + count=2, + shared=False, + evict_first=False, + mutable=False, + volatile=False, + scope=None, + cache_policy="cg", + word_type="b32", + result_type=Uint32, +): + """``count`` words (a tuple); ``scope`` is a relaxed load at that scope. + ``mutable`` (a side-effecting load) for data written during the step (activations, + tables, refills). Only static weights use ``evict_first``.""" + space = "global" + (f".{cache_policy}" if cache_policy else "") + if shared: + space = "shared" + if volatile: + space = "volatile.global" + if scope: + space = f"relaxed.{scope}.global" + policy = hint = tail = "" + if evict_first: + assert not shared + space = "global.nc.L1::no_allocate" + policy = ".reg .b64 pol; createpolicy.fractional.L2::evict_first.b64 pol, 1.0; " + hint, tail = ".L2::cache_hint", ", pol" + vector, registers = _vector(count) + address = f"[${count}]{tail}" + code = f"{{ {policy}ld.{space}{hint}{vector}.{word_type} {registers}, {address}; }}" + constraints = ["=f" if result_type is Float32 else "=r"] * count + constraints += ["r" if shared else "l"] + (["~{memory}"] if scope else []) + # Shared data is written by other threads between reads: never pure. + pure = not (shared or mutable or volatile or scope) + return code, ",".join(constraints), (pointer,), (result_type,) * count, pure + + +smem_words = partial(load_words, count=4, shared=True, mutable=True) +load_volatile = partial(load_words, count=1, volatile=True, word_type="u32") +mailbox_load = partial(load_words, count=1, scope="gpu", word_type="u32") +load_f32x4 = partial(load_words, count=4, result_type=Float32, word_type="f32") +early_weights = partial( + load_words, cache_policy="nc.L1::no_allocate", word_type="u32", mutable=True +) + + +@asm() +def load_bf16x8(pointer, *, shared=False, mutable=False): + opcode = "ld.shared.v4.b32" if shared else "ld.global.cg.v4.b32" + # Shifting BF16 bits into the high FP32 half preserves values exactly. + code = "{ .reg .b32 x<4>; " + opcode + " {x0,x1,x2,x3}, [$8]; " + for j in range(4): + code += f"shl.b32 ${2 * j}, x{j}, 16; and.b32 ${2 * j + 1}, x{j}, 0xffff0000; " + constraints = ",".join(["=f"] * 8 + ["r" if shared else "l"]) + return code + "}", constraints, (pointer,), (Float32,) * 8, not (shared or mutable) + + +@asm() +def load_vector(pointer, *, cache="cg"): + """``cache`` "weight": a streaming L2 evict_first load of static weights; else + ``ld.global.``. Side-effecting in every mode.""" + opcode, policy, tail = "ld.global." + cache, "", "" + if cache == "weight": + policy = ".reg .b64 pol; createpolicy.fractional.L2::evict_first.b64 pol,1.0; " + opcode = "ld.global.L1::no_allocate.L2::cache_hint.L2::256B" + tail = ",pol" + if cache == "volatile": + opcode = "ld.volatile.global" + code = "{ " + policy + opcode + ".v4.b32 {$0,$1,$2,$3},[$4]" + tail + "; }" + clobber = ",~{memory}" if cache == "volatile" else "" + return code, "=r,=r,=r,=r,l" + clobber, (pointer,), (Uint32,) * 4 + + +@asm(Float32, pure=True) +def table_load(pointer): + return "ld.global.nc.f32 $0,[$1];", "=f,l", (pointer,) + + +@asm() +def store_words(pointer, words, *, shared=False, scope=None): + """``scope`` is a relaxed global store at that scope.""" + vector, registers = _vector(len(words), 1) + if shared: + # Other threads read it: a memory side effect, like the global stores. + space, constraints = ( + f"shared{vector}.b32", + "r" + ",r" * len(words) + ",~{memory}", + ) + else: + space = ( + f"relaxed.{scope}.global{vector}.u32" if scope else f"global{vector}.b32" + ) + constraints = "l" + ",r" * len(words) + ",~{memory}" + words = tuple(Uint32(v) for v in words) + return f"st.{space} [$0], {registers};", constraints, (pointer, *words) + + +store_shared_words = partial(store_words, shared=True) +store_relaxed = partial(store_words, scope="sys") +mailbox_store = partial(store_words, scope="gpu") + + +@asm() +def store_f32x4(pointer, values, when=None): + """``when``: a predicate, so callers need no branch.""" + code = "st.global.v4.f32 [$0], {$1, $2, $3, $4};" + if when is None: + return code, "l,f,f,f,f", (pointer, *values) + code = "{ .reg .pred p; setp.ne.b32 p,$5,0; @p " + code + " }" + return code, "l,f,f,f,f,r", (pointer, *values, Int32(when)) + + +@asm() +def store_gate_mailbox(pointer, bits: Uint16): + return "st.relaxed.gpu.global.u16 [$0], $1;", "l,h,~{memory}", (pointer, bits) + + +@asm() +def store_bf16_pair(pointer, a, b): + code = "{ .reg .b32 pair; cvt.rn.bf16x2.f32 pair,$2,$1; " + return code + "st.relaxed.gpu.global.b32 [$0],pair; }", "l,f,f", (pointer, a, b) + + +@asm() +def store_pair(pointer, a, b): + """BF16x2 (a, b) with negative-zero halves made +0.""" + code = "{ .reg .b32 pair,half; .reg .pred p; cvt.rn.bf16x2.f32 pair,$2,$1; " + code += "and.b32 half,pair,0xffff; setp.eq.b32 p,half,0x8000; " + code += "@p and.b32 pair,pair,0xffff0000; shr.u32 half,pair,16; " + code += "setp.eq.b32 p,half,0x8000; @p and.b32 pair,pair,0x0000ffff; " + return code + "st.global.b32 [$0],pair; }", "l,f,f", (pointer, a, b) + + +@asm() +def push_float(pointer, value, barrier, peer, when=None): + """DSMEM store of one FP32 ``value`` (or a 4-tuple) completing on ``barrier`` + (``when``: predicated on it).""" + values = tuple(map(Float32, value if isinstance(value, tuple) else (value,))) + vector, registers = _vector(len(values), 1) + b, p, w = (f"${len(values) + i}" for i in (1, 2, 3)) + code = f"{{ .reg .b32 dst,bar; mapa.shared::cluster.u32 dst,$0,{p}; " + code += f"mapa.shared::cluster.u32 bar,{b},{p}; " + guard = "" if when is None else f".reg .pred w; setp.ne.b32 w,{w},0; @w " + code += ( + guard + f"st.async.shared::cluster.mbarrier::complete_tx::bytes{vector}.b32 " + ) + code += f"[dst],{registers},[bar]; }}" + constraints = "r" + ",f" * len(values) + ",r,r" + ",r" * (when is not None) + operands = (pointer, *values, barrier, Int32(peer)) + operands += () if when is None else (Int32(when),) + return code, constraints + ",~{memory}", operands + + +@asm() +def multicast_words(address, words, *, kind="f32"): + """Relaxed multimem store of FP32 ("f32") or integer ("b32", "bf16x2") words.""" + vector, registers = _vector(len(words), 1) + f32 = kind == "f32" + values = (bits_float(v) if f32 else Uint32(v) for v in words) + constraints = "l" + (",f" if f32 else ",r") * len(words) + ",~{memory}" + code = f"multimem.st.relaxed.sys.global{vector}.{kind} [$0],{registers};" + return code, constraints, (Int64(address), *values) + + +@asm() +def publish_hint(pointer): + return "red.relaxed.gpu.global.add.u32 [$0],1;", "l,~{memory}", (pointer,) + + +@asm() +def copy_g2s(dst, src, size=16, barrier=None, *, evict_first=False): + """16-byte cp.async, or a static/dynamic ``size`` bulk copy on ``barrier``.""" + operands = [dst, src] + policy = hint = tail = "" + if evict_first: + policy = ".reg .b64 pol; createpolicy.fractional.L2::evict_first.b64 pol, 1.0; " + hint, tail = ".L2::cache_hint", ", pol" + if barrier is None: + assert size == 16 + code = f"cp.async.cg.shared.global{hint} [$0], [$1], 16{tail};" + else: + operands.append(barrier) + if not isinstance(size, int): + operands.append(size) + size = "$3" + code = ( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes" + f"{hint} [$0], [$1], {size}, [$2]{tail};" + ) + constraints = ",".join(["r", "l"] + ["r"] * (len(operands) - 2)) + return "{ " + policy + code + " }", constraints, operands + + +copy_evict = partial(copy_g2s, evict_first=True) + + +@asm() +def copy_bulk(dst, src, barrier, size): + code = "{ .reg .b64 p; createpolicy.fractional.L2::evict_first.b64 p,1.0; " + code += "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes" + code += ".L2::cache_hint [$0],[$1],$3,[$2],p; }" + return code, "r,l,r,r,~{memory}", (dst, src, barrier, Int32(size)) + + +@asm() +def copy_s2cluster(dst, src, barrier, peer, *, size): + code = "{ .reg .b32 dst,bar; mapa.shared::cluster.u32 dst,$0,$3; " + code += "mapa.shared::cluster.u32 bar,$2,$3; cp.async.bulk.shared::cluster" + code += f".shared::cta.mbarrier::complete_tx::bytes [dst],[$1],{size},[bar]; }}" + return code, "r,r,r,r", (dst, src, barrier, Int32(peer)) + + +@asm() +def stage_weight(dst, src): + code = "{ .reg .b64 pol; createpolicy.fractional.L2::evict_first.b64 pol,1.0; " + code += "cp.async.cg.shared.global.L2::cache_hint.L2::256B [$0],[$1],16,pol; }" + return code, "r,l,~{memory}", (dst, src) + + +@asm() +def load_tma2(dst, descriptor, barrier, column, row): + code = "{ .reg .b64 pol; createpolicy.fractional.L2::evict_first.b64 pol,1.0; " + code += "cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier" + code += "::complete_tx::bytes.L2::cache_hint [$0],[$1,{$3,$4}],[$2],pol; }" + return code, "r,l,r,r,r", (dst, descriptor, barrier, Int32(column), Int32(row)) + + +@asm() +def load_tma3(dst, descriptor, barrier, row, chunk): + code = "{ .reg .b64 pol; createpolicy.fractional.L2::evict_first.b64 pol,1.0; " + code += "cp.async.bulk.tensor.3d.shared::cluster.global.tile.mbarrier" + code += "::complete_tx::bytes.L2::cache_hint [$0],[$1,{0,$3,$4}],[$2],pol; }" + return code, "r,l,r,r,r", (dst, descriptor, barrier, Int32(row), Int32(chunk)) + + +@asm() +def prefetch_bytes(src, size): + return "cp.async.bulk.prefetch.L2.global [$0], $1;", "l,r", (src, Int32(size)) + + +@dsl_user_op +def prefetch_bulk(address, size, *, loc=None, ip=None): + """``prefetch_bytes`` of a 64-bit address through the NVVM op.""" + pointer = cute.make_ptr(cutlass.Uint8, address, cute.AddressSpace.gmem) + nvvm.cp_async_bulk_prefetch(pointer, size, loc=loc, ip=ip) + + +@asm() +def prefetch_l2(pointer): + return "prefetch.global.L2 [$0];", "l", (pointer,) + + +@asm() +def prefetch_map(descriptor, row, chunk): + code = "cp.async.bulk.prefetch.tensor.3d.L2.global.tile [$0,{0,$1,$2}];" + return code, "l,r,r", (descriptor, Int32(row), Int32(chunk)) + + +def acquire_tensormap(pointer): + generic = cute.make_ptr(cutlass.Uint8, pointer.toint(), cute.AddressSpace.generic) + nvvm.fence_proxy_acquire("sys", generic, 128) + nvvm.prefetch_tensormap(generic) + + +@asm() +def acquire_tensormap_asm(pointer): + """``acquire_tensormap`` as one asm block on the unfolded address.""" + code = "fence.proxy.tensormap::generic.acquire.sys [$0],128; " + return code + "prefetch.tensormap [$0];", "l", (pointer,) + + +@asm(Int32) +def try_wait(barrier, phase=0, *, test=False, scope="cluster"): + """Acquire parity probe at ``scope`` ("cluster", or "cta" when every arrival is + CTA-local); ``test`` is the non-blocking probe of phase 0.""" + op, parity = ("test_wait", "0") if test else ("try_wait", "$2") + code = f"{{ .reg .pred p; mbarrier.{op}.parity.acquire.{scope}.shared::cta.b64 " + code += f"p,[$1],{parity}; selp.u32 $0,1,0,p; }}" + if test: + return code, "=r,r", (barrier,) + return code, "=r,r,r,~{memory}", (barrier, Int32(phase)) + + +@cute.jit +def wait(barrier, phase=0, test: cutlass.Constexpr[bool] = False, scope="cluster"): + done = try_wait(barrier, phase, test=test, scope=scope) + while done == 0: + done = try_wait(barrier, phase, test=test, scope=scope) + + +@asm(Int32) +def try_wait_phase_zero(barrier): + code = "{ .reg .pred ready; mbarrier.try_wait.parity.acquire.cluster" + code += ".shared::cta.b64 ready,[$1],0; selp.u32 $0,1,0,ready; }" + return code, "=r,r,~{memory}", (barrier,) + + +@cute.jit +def wait_phase_zero_cluster(barrier: cute.Pointer): + done = try_wait_phase_zero(barrier) + while done == 0: + done = try_wait_phase_zero(barrier) + + +def init_barriers(barriers, count, *, async_proxy=False): + for i in range(count): + cute.arch.mbarrier_init(barriers + i, 1) + cute.arch.mbarrier_init_fence() + if async_proxy: + cute.arch.fence_proxy("async.shared", space="cta") + + +@asm() +def named_barrier(barrier_id=1, count=128, *, arrive=False, all_groups=False): + if all_groups: + barrier_id, count = 2, 512 + op = "arrive" if arrive else "sync" + return f"bar.{op} {barrier_id},{count};", "~{memory}", () + + +def group_barrier(group): + cute.arch.barrier(barrier_id=1 + 3 * group, number_of_threads=128) + + +@asm() +def sleep_ns(ns): + """``ns`` is an immediate if a Python int, else a register.""" + if isinstance(ns, int): + return f"nanosleep.u32 {ns};", "", () + return "nanosleep.u32 $0;", "r", (Uint32(ns),) + + +@cute.jit +def smem_desc(pointer: cute.Pointer, swizzle: cutlass.Constexpr[bool]): + address = shared_address(pointer) + value = Uint64((address & 262143) >> 4) | Uint64(1) << 46 + if cutlass.const_expr(swizzle): + value |= Uint64(1) << 16 | Uint64(64) << 32 | Uint64(2) << 61 + else: + value |= Uint64(8) << 32 + return value + + +def fence_tmem(*, before=False): + nvvm.tcgen05_fence("before_thread_sync" if before else "after_thread_sync") + + +def wait_tmem(barrier, phase=0, scope="cluster"): + """``wait``, then order later tcgen05 operations after it.""" + wait(barrier, phase, scope=scope) + fence_tmem() + + +@asm() +def free_tmem(address, columns): + code = f"tcgen05.dealloc.cta_group::1.sync.aligned.b32 $0,{columns};" + return code, "r", (Uint32(address),) + + +@cute.jit +def release_tmem(tmem, columns: cutlass.Constexpr[int], warp, owner): + """CTA teardown: warp ``owner`` frees the TMEM after every tcgen05 operation.""" + fence_tmem(before=True) + cute.arch.sync_threads() + if warp == owner: + fence_tmem() + free_tmem(tmem, columns) + + +@asm() +def require_zero_tmem(address): + return "{ .reg .pred p; setp.ne.u32 p,$0,0; @p trap; }", "r", (Uint32(address),) + + +@asm() +def load_tmem(address, count=1): + """``count`` FP32 TMEM columns of this lane, waiting for the load.""" + registers = ",".join(f"${i}" for i in range(count)) + code = f"tcgen05.ld.sync.aligned.32x32b.x{count}.b32 {{{registers}}},[${count}]; " + code += "tcgen05.wait::ld.sync.aligned;" + constraints = ",".join(["=f"] * count + ["r", "~{memory}"]) + result = Float32 if count == 1 else (Float32,) * count + return code, constraints, (Uint32(address),), result + + +@asm() +def copy_scale(column, descriptor): + code = "tcgen05.cp.cta_group::1.32x128b.warpx4 [$0],$1;" + return code, "r,l", (Uint32(column), Uint64(descriptor)) + + +commit = nvvm.tcgen05_commit + + +@asm() +def commit_asm(barrier): + """``commit`` as an asm block on the unfolded address.""" + code = "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [$0];" + return code, "r", (barrier,) + + +@asm() +def mma_bf16_block(d, a, b, instruction, accumulate, empty, release): + """Warp-uniform issue of one 64-K block (K-major 128B-swizzled A and B): the + elected lane runs the four K16 BF16 MMAs into ``d`` (the first accumulates iff + ``accumulate`` is nonzero) and, if ``release``, commits ``empty``.""" + mma = "tcgen05.mma.cta_group::1.kind::f16" + commit = "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64" + code = "{ .reg .pred e,p0,pt,pc; .reg .b64 a1,a2,a3,b1,b2,b3; " + code += "elect.sync _|e,0xffffffff; setp.eq.b32 pt,0,0; setp.ne.b32 p0,$4,0; " + code += "setp.ne.b32 pc,$6,0; and.pred pc,pc,e; " + code += "add.s64 a1,$1,2; add.s64 a2,$1,4; add.s64 a3,$1,6; " + code += "add.s64 b1,$2,2; add.s64 b2,$2,4; add.s64 b3,$2,6; " + code += f"@e {mma} [$0],$1,$2,$3,p0; @e {mma} [$0],a1,b1,$3,pt; " + code += f"@e {mma} [$0],a2,b2,$3,pt; @e {mma} [$0],a3,b3,$3,pt; " + code += f"@pc {commit} [$5]; }}" + operands: tuple = (Uint32(d), Uint64(a), Uint64(b), Uint32(instruction)) + operands += (Uint32(accumulate), empty, Uint32(release)) + return code, "r,l,l,r,r,r,r", operands + + +@asm() +def mma_mx(d, a, b, instruction, sfa, sfb, accumulate): + """One block-scaled MXF8 tcgen05 MMA; ``accumulate`` is a runtime flag.""" + words = (Uint32(v) for v in (instruction, sfa, sfb, accumulate)) + code = "{ .reg .pred p; setp.ne.b32 p,$6,0; tcgen05.mma.cta_group::1" + code += ".kind::mxf8f6f4.block_scale [$0],$1,$2,$3,[$4],[$5],p; }" + return code, "r,l,l,r,r,r,r", (Uint32(d), Uint64(a), Uint64(b), *words) + + +@cute.jit +def mma_steps(d, a, b, sfa, sfb, accumulate, fc2: cutlass.Constexpr[bool]): + for k in cutlass.range_constexpr(6 if fc2 else 8): + a_offset: cutlass.Constexpr = k // 4 * 1024 + k % 4 * 2 + 4 * (fc2 and k >= 4) + instruction: cutlass.Constexpr = 142738048 | (k % 4 << 29) | (k % 4 << 4) + ops: tuple = (d, a + a_offset, b + k // 4 * 64 + k % 4 * 2, instruction) + ops += (sfa + k // 4 * 4, sfb + (k * 4 if fc2 else k // 4 * 4)) + mma_mx(*ops, accumulate if k == 0 else 1) + + +@asm() +def chunk_mma(*args): + """(d, adesc, bdesc, sfa, sfa_src, sfb, sfb_src, copy_sfb, accum, empty_bar, + acc_bar): four K32 MMAs of one K128 chunk, with optional completion arrivals.""" + constraints = "r,l,l,r,l,r,l,r,r,r,r" + types = [Uint32 if kind == "r" else Uint64 for kind in constraints.split(",")] + cp, mma = "tcgen05.cp.cta_group::1.32x128b.warpx4", "tcgen05.mma.cta_group::1" + commit = "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64" + code = ( + "{\n.reg .pred e,p0,pt,pc,pa,pe; .reg .b32 i<4>; .reg .b64 xa<4>,xb<4>;\n" + "elect.sync _|e,0xffffffff; setp.eq.b32 pt,0,0; setp.ne.b32 pc,$7,0;\n" + "and.pred pc,pc,e; setp.ne.b32 pa,$10,0; and.pred pa,pa,e;\n" + "setp.ne.b32 pe,$9,0; and.pred pe,pe,e; setp.ne.b32 p0,$8,0;\n" + "mov.b32 i0,142738048; or.b32 i1,i0,536870928; or.b32 i2,i0,1073741856;\n" + f"or.b32 i3,i0,1610612784; @pc {cp} [$5],$6; @e {cp} [$3],$4;\n" + "add.s64 xa0,$1,0; add.s64 xb0,$2,0; add.s64 xa1,$1,2; add.s64 xb1,$2,2;\n" + "add.s64 xa2,$1,4; add.s64 xb2,$2,4; add.s64 xa3,$1,6; add.s64 xb3,$2,6;\n" + ) + for k in range(4): + code += f"@e {mma}.kind::mxf8f6f4.block_scale [$0],\n" + code += f"xa{k},xb{k},i{k},[$3],[$5],{'pt' if k else 'p0'};\n" + code += f"@pe {commit} [$9]; @pa {commit} [$10];\n}}" + return code, constraints, tuple(t(v) for t, v in zip(types, args)) + + +def launch(kernel, args, grid, block, cluster=None, pdl=True, **options): + """Launch kernel(args[:-1]) on stream args[-1]; an int dimension means (n, 1, 1).""" + dim = lambda d: d if isinstance(d, tuple) else (d, 1, 1) + if cluster is not None: + options["cluster"] = dim(cluster) + kernel(args[:-1]).launch( + grid=dim(grid), block=dim(block), stream=args[-1], use_pdl=pdl, **options + ) + + +@cute.jit +def scatter_reduce(values, lane, count: cutlass.Constexpr[int]): + """Warp-reduce ``count`` values; lane groups end up owning one value each.""" + selected = Int32(0) + for stage in cutlass.range_constexpr(5): + offset = 16 >> stage + remaining = max(count >> stage, 1) + if cutlass.const_expr(remaining > 1): + half = remaining // 2 + upper = (lane & offset) != 0 + for i in cutlass.range_constexpr(half): + send = values[i] if upper else values[i + half] + keep = values[i + half] if upper else values[i] + values[i] = keep + cute.arch.shuffle_sync_bfly(send, offset) + if upper: + selected += half + else: + values[0] += cute.arch.shuffle_sync_bfly(values[0], offset) + return values[0], selected diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/twoshot.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/twoshot.py new file mode 100644 index 0000000000..adb036460a --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/twoshot.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Two-shot Lamport all-reduce of BF16 rows [B, width] over TP ranks. + +Rank ``o`` owns columns [o W, (o + 1) W), W = width / owners. Every buffer is +Lamport-armed and only its local reader re-arms it: + +- ``push``: a rank stores its 8-column chunk of row ``token`` into the owner's + receive slot [rank] (a peer store; the owner's ``recv`` [TP, cap, W]; callers with a + runtime parity pass the whole [parities x TP, cap, W] and slot parity x TP + rank). +- ``reduce``: the owner polls the TP slots of an owned chunk in its own memory, sums + them in FP32 in rank order, rounds once and multicasts the BF16 sum into every + rank's gather row [cap, width]; then re-arms the slots. +- every rank then polls its gather row (one slot) and re-arms it. + +Each GPU moves ~2 x B x width x 2 bytes instead of TP x B x width x 2 (one-shot), and +every rank gets the same bits. ``owners`` = 1 is a one-GPU stand-in: rank 0 owns every +column and its TP local slots play the ranks. The o-projection reduces with ``reduce`` +and re-arms at once (two parities); the MoE tail reduces in WaveMoeTail and re-arms the +previous MoE layer's rows (three parities). +""" + +import cutlass +import cutlass.cute as cute +import torch +import torch.distributed._symmetric_memory as symm_mem +from cutlass import Float32, Int32, Int64, Uint32 + +from .mailbox import rearm +from .primitives import ( + add_rn, + at, + load_vector, + multicast_words, + pack_pair, + published, + store_relaxed, + unpack_bf16, + zeros, +) + +SENTINEL = -0x80000000 + + +class TwoShot: + """Symmetric buffers per layer parity: ``recv`` [TP, cap, W] (this rank's owned + columns, one slot per source rank), ``gather`` [1, cap, width] (multicast) and the + ranks' ``recv`` base addresses [parity, TP].""" + + def __init__(self, group, device, tp, rank, width=7168, cap=64, parities=2): + self.tp, self.rank, self.width = tp, rank, width + shape = (parities, tp, cap, width // tp) + bf16 = dict(dtype=torch.bfloat16, device=device) + self.recv = symm_mem.empty(shape, **bf16) + self.gather = symm_mem.empty((parities, 1, cap, width), **bf16) + handles = [symm_mem.rendezvous(t, group) for t in (self.recv, self.gather)] + size = self.recv[0].numel() * 2 + ptrs = handles[0].buffer_ptrs + bases = [[int(b) + p * size for b in ptrs] for p in range(parities)] + self.peers = torch.tensor(bases, dtype=torch.int64, device=device) + self.gather_mc = int(handles[1].multicast_ptr or 0) + if not self.gather_mc: # views() offsets this address + raise RuntimeError("Kimi K3 two-shot all-reduce needs multicast memory") + for t in (self.recv, self.gather): + t.view(torch.int32).fill_(SENTINEL) + torch.accelerator.synchronize(device) + torch.distributed.barrier(group=group) + + def views(self, parity): + """(recv, peers, gather, gather multicast address) of layer parity.""" + gather_mc = self.gather_mc + parity * self.gather[0].numel() * 2 + return self.recv[parity], self.peers[parity], self.gather[parity], gather_mc + + +@cute.jit +def push(recv, peers, slot, token, column, words): + """Store this rank's BF16 chunk (token, column .. column + 7) in receive slot + ``slot`` of its owner.""" + width: cutlass.Constexpr = recv.shape[2] + owner = column // width + offset = cute.crd2idx((slot, token, column - owner * width), recv.layout) + store_relaxed(Int64(peers[owner]) + Int64(offset) * 2, words) # polled remotely + + +@cute.jit +def reduce(recv, gather, gather_mc, rank, token, column): + """Owner of (token, column): the TP slots summed in FP32 in rank order, rounded + once (no negative zero, the sentinel), multicast to every gather row; slots + re-armed. Returns nothing; callers then poll ``gather`` like a one-slot mailbox.""" + tp: cutlass.Constexpr = recv.shape[0] + local = column - rank * recv.shape[2] + raw = zeros((tp, 4), Uint32) + pending = Int32((1 << tp) - 1) + while pending != 0: + for slot in cutlass.range_constexpr(tp): + if (pending & (1 << slot)) != 0: + words = load_vector(at(recv, slot, token, local), cache="volatile") + for j in cutlass.range_constexpr(4): + raw[slot, j] = words[j] + if published(words): + pending &= ~(1 << slot) + values = cute.make_rmem_tensor(8, Float32) + for j in cutlass.range_constexpr(4): + low, high = unpack_bf16(raw[0, j]) + for slot in cutlass.range_constexpr(1, tp): + a, b = unpack_bf16(raw[slot, j]) + low, high = add_rn(low, a), add_rn(high, b) + values[2 * j], values[2 * j + 1] = low, high + words = tuple( + pack_pair(values[2 * j], values[2 * j + 1], policy="latent") for j in range(4) + ) + offset = cute.crd2idx((0, token, column), gather.layout) + multicast_words(gather_mc + Int64(offset) * 2, words) + for slot in cutlass.range_constexpr(tp): + rearm(at(recv, slot, token, local)) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/__init__.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/projection.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/projection.py new file mode 100644 index 0000000000..f74500a128 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/projection.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Split-K projection on the FlashInfer CuTe core vendored in ``splitk.py`` (with an +evict_first weight TMA; policies map padded rows to (M tile, N tile, split K, +stages); attention owns live slot masks).""" + +from functools import cache + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import make_fake_stream, make_fake_tensor + +from ..common.host import architecture, tensor_maps, to_cute +from ..common.host import stream as current_stream +from .splitk import SplitKGemm, SplitKTactic + + +def matrix_view(tensor, *, output=False): + """Keep packed row pitches and FlashInfer's singleton-safe dynamic major axis.""" + rows, cols = tensor.shape + shape = (cols, rows, 1) if output else (rows, cols, 1) + stride = (1, tensor.stride(0), 0) if output else (tensor.stride(0), 1, 0) + result = to_cute(tensor.as_strided(shape, stride)) + return result.mark_layout_dynamic(leading_dim=0 if output else 1) + + +@cache +def compile_projection(tactic, output_dtype, arch): + """Code-only cache; weights, descriptors and activation pointers stay private.""" + + def fake(leading, dtype=cutlass.BFloat16): + shape = tuple(cute.sym_int() for _ in range(3)) + stride = (1, cute.sym_int(64), 0) if leading == 0 else (cute.sym_int(64), 1, 0) + return make_fake_tensor(dtype, shape, stride, assumed_align=16) + + y = fake(0, cutlass.Float32 if output_dtype == torch.float32 else cutlass.BFloat16) + split = SplitKTactic(*tactic) + kernel = SplitKGemm(tactic=split) + args = (fake(1), fake(1), y, y, cutlass.Int64(0), make_fake_stream()) + return cute.experimental.compile(kernel, *args, options=f"--gpu-arch={arch}") + + +def weight_map(weight, rows): + """Device TMA map of ``weight`` [N, K] (row pitch kept): box {64, rows}, SW128.""" + types = cuda.CUtensorMapDataType + dtype = { + torch.bfloat16: types.CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, + torch.float16: types.CU_TENSOR_MAP_DATA_TYPE_FLOAT16, + }[weight.dtype] + swizzle = cuda.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_128B + l2 = cuda.CUtensorMapL2promotion.CU_TENSOR_MAP_L2_PROMOTION_L2_256B + n, k = weight.shape + pitch = weight.stride(0) * weight.element_size() + spec = (weight.data_ptr(), dtype, (k, n), (pitch,), (64, rows), swizzle, l2) + return tensor_maps([spec], weight.device) + + +class PreparedProjection: + """Split-K projection with a per-batch tactic table and caller-owned outputs.""" + + def __init__(self, weight, policies, *, output_dtype=torch.bfloat16): + if weight.dim() != 2 or weight.dtype not in (torch.bfloat16, torch.float16): + raise ValueError( + f"projection weight must be a BF16/FP16 [N, K] matrix, got " + f"{tuple(weight.shape)} {weight.dtype}" + ) + pitch = weight.stride(0) * weight.element_size() + if weight.stride(1) != 1 or pitch % 16 or weight.data_ptr() % 16: + raise ValueError( + "projection weight rows need unit stride and a 16 B TMA pitch" + ) + if output_dtype not in (torch.bfloat16, torch.float32): + raise ValueError(f"unsupported projection output dtype {output_dtype}") + self.weight, self.device = weight, weight.device + self.output_dtype, self.policies = output_dtype, dict(policies) + self.weight_view = matrix_view(weight) + self.kernels, self.extra, self.maps = {}, {}, {} + + def prepare(self): + arch, dtype = architecture(self.device), self.output_dtype + with torch.accelerator.device_index(self.device.index): + for batch, tactic in self.policies.items(): + code = compile_projection(tactic, dtype, arch) + self.kernels[batch] = code + self.extra[batch] = self.weight_map(tactic[0]) + return self + + def weight_map(self, rows): + """The kernel's extra argument: A's map for ``rows``-row tiles.""" + if rows not in self.maps: + self.maps[rows] = weight_map(self.weight, rows) + return (cutlass.Int64(self.maps[rows].data_ptr()),) + + def run(self, x, out): + y = matrix_view(out, output=True) + b, stream = x.shape[0], current_stream(self.device) + if b not in self.kernels or x.shape[1] != self.weight.shape[1]: + raise ValueError( + f"projection input {tuple(x.shape)}: prepared batches " + f"{sorted(self.kernels)}, K {self.weight.shape[1]}" + ) + code, extra = self.kernels[b], self.extra[b] + code(self.weight_view, matrix_view(x), y, y, *extra, stream) + return out diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/splitk.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/splitk.py new file mode 100644 index 0000000000..c0591ad578 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/splitk.py @@ -0,0 +1,380 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from FlashInfer v0.7.0.post1 (Apache-2.0): +# flashinfer/gemm/kernels/dense_bf16_gemm_sm100_splitk.py +"""FlashInfer 0.7.0.post1 ``SplitKDenseGemmKernel`` (``dense_bf16_gemm_sm100_splitk``) +reduced to our use (PDL, no bias) with the weight (kernel-A) TMA issued through a +device tensor map with an L2 evict_first policy, and its synchronization fixed: the +owner waits for the peers' ``st.async`` partials at ``.acquire.cluster``; the MMA warp +waits for TMEM-pointer phase 0 before its release arrival (else its phase-1 wait can +pass before the epilogue has read TMEM); ``tcgen05.fence::*_thread_sync`` around +every mbarrier hand-off of tcgen05 work. + +Each cluster rank accumulates an exact K slice in FP32; peers publish partials to +rank 0 through DSMEM, which reduces, casts and stores once. Kernel-M carries public N +and kernel-N public M.""" + +import cuda.bindings.driver as _cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass import Int32 +from cutlass._mlir.dialects import llvm +from cutlass.cute import experimental as cute_ext +from cutlass.cute.nvgpu import tcgen05 +from cutlass.cutlass_dsl import dsl_user_op +from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import ( + OWNER_RANK, + SplitKTactic, + _store_shared_remote_v4, +) + +CTA_K, ATOM_K = 128, 64 # K per tile; K per 128-byte swizzle atom (one A box) + +__all__ = ["SplitKTactic", "SplitKGemm"] + + +def _asm(code, constraints, operands, loc, ip): + llvm.inline_asm( + None, operands, code, constraints, has_side_effects=True, + is_align_stack=False, asm_dialect=llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _acquire_tensor_map(tensor_map, *, loc=None, ip=None): + code = "fence.proxy.tensormap::generic.acquire.sys [$0], 128; " + _asm(code + "prefetch.tensormap [$0];", "l", [tensor_map.ir_value()], loc, ip) + + +@dsl_user_op +def _tma_load_evict_first(smem, mbar, tensor_map, k0, row0, *, loc=None, ip=None): + """One {64, rows} box of A at (k0, row0), L2 evict_first.""" + code = "{ .reg .b64 pol; createpolicy.fractional.L2::evict_first.b64 pol, 1.0; " + code += "cp.async.bulk.tensor.2d.shared::cta.global.tile.mbarrier::complete_tx" + code += "::bytes.L2::cache_hint [$0], [$1, {$2, $3}], [$4], pol; }" + operands = [smem.toint().ir_value(), tensor_map.ir_value(), k0.ir_value()] + operands += [row0.ir_value(), mbar.toint().ir_value()] + _asm(code, "r,l,r,r,r", operands, loc, ip) + + +@dsl_user_op +def _wait_cluster(mbar, phase, *, loc=None, ip=None): + """``cute.arch.mbarrier_wait`` at ``.acquire.cluster`` (remote ``st.async``).""" + code = "{ .reg .pred P1; LAB_WAIT: " + code += "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 P1, [$0], $1, " + code += "10000000; @P1 bra.uni DONE; bra.uni LAB_WAIT; DONE: }" + _asm(code, "r,r", [mbar.toint().ir_value(), phase.ir_value()], loc, ip) + + +@dsl_user_op +def _fence_tcgen05(*, before=False, loc=None, ip=None): + side = "before" if before else "after" + _asm(f"tcgen05.fence::{side}_thread_sync;", "", [], loc, ip) + + +def _smem(dtype, size, alignment, layout=None): + layout = cute.make_layout(size) if layout is None else layout + return cute_ext.allocate(dtype, cute.AddressSpace.smem, layout, alignment=alignment) + + +class SplitKGemm: + """BF16 GEMM with a cluster-local split-K reduction; ``a_map`` is the device + address of A's tensor map (box {64, mma_m}, 128-byte swizzle).""" + + def __init__(self, *, tactic: SplitKTactic): + self.acc_dtype = cutlass.Float32 + self.cta_m, self.cta_n, self.cta_k = tactic.mma_m, tactic.mma_n, CTA_K + self.num_ab_stage, self.split_k = tactic.ab_stages, tactic.split_k + self.threads_per_cta, self.epilog_threads = 256, 128 + self.mma_tiler_mn = (tactic.mma_m, tactic.mma_n) + self.cta_group = tcgen05.CtaGroup.ONE + self.cluster_shape = (1, tactic.split_k, 1) + values_per_thread = tactic.mma_m * tactic.mma_n // self.epilog_threads + assert values_per_thread % 4 == 0, "remote stores move 4 values" + self.mailbox_elements = ( + (tactic.split_k - 1) * self.epilog_threads * values_per_thread + ) + self.expected_transaction_bytes = self.mailbox_elements * 4 + + @cute.experimental.jit + def __call__( + self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, bias: cute.Tensor, + a_map: cutlass.Int64, stream: _cuda.CUstream, + ): # fmt: skip + # Grid-y packs output-N tile and cluster rank. + self.kernel(a, b, c, bias, a_map).launch( + grid=( + cute.ceil_div(c.layout.shape[0], self.cta_m), + cute.ceil_div(c.layout.shape[1], self.cta_n) * self.split_k, + c.layout.shape[2], + ), + block=(self.threads_per_cta, 1, 1), + cluster=self.cluster_shape, + smem=cute.Int64(utils.get_smem_capacity_in_bytes("sm_100")), + stream=stream, + use_pdl=True, + ) + + @cute.experimental.kernel + def kernel( + self, mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor, mBias: cute.Tensor, + a_map: cutlass.Int64, + ): # fmt: skip + """A (M, K, L) and B (N, K, L) K-major, C (M, N, L) M-major; bias unused.""" + stages = self.num_ab_stage + ab_dtype = mA.element_type + majors = (utils.LayoutEnum.from_tensor(t).mma_major_mode() for t in (mA, mB)) + tiled_mma = sm100_utils.make_trivial_tiled_mma( + ab_dtype, ab_dtype, *majors, self.acc_dtype, self.cta_group, + self.mma_tiler_mn, + ) # fmt: skip + + mnk_tiler = (self.mma_tiler_mn[0], self.mma_tiler_mn[1], self.cta_k) + block_idx = cute.arch.block_idx() + bidx = block_idx[0] + split_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + n_idx = block_idx[1] // self.split_k + l_idx = block_idx[2] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + shape = (tiled_mma, mnk_tiler, ab_dtype, stages) + sA = _smem(ab_dtype, None, 1024, sm100_utils.make_smem_layout_a(*shape)) + sB = _smem(ab_dtype, None, 1024, sm100_utils.make_smem_layout_b(*shape)) + acc_layout = cute_ext.make_tmem_layout_acc( + tiled_mma, self.mma_tiler_mn, acc_stage=1 + ) + bar_full = _smem(cutlass.Int64, stages, 8).iterator + bar_empty = _smem(cutlass.Int64, stages, 8).iterator + bar_tma_epilog = _smem(cutlass.Int64, 1, 8).iterator # FlashInfer's bias + bar_mma_epilog = _smem(cutlass.Int64, 1, 8).iterator + bar_tmem_alloc = _smem(cutlass.Int64, 1, 8).iterator + tmem_base_ptr = _smem(cutlass.Int32, 1, 4).iterator + mailbox, bar_reduce = sA, bar_mma_epilog # split 1: compile-time dead + if cutlass.const_expr(self.split_k > 1): + mailbox = _smem(cutlass.Float32, self.mailbox_elements, 128) + bar_reduce = _smem(cutlass.Int64, 1, 8).iterator + + if warp_idx == 0: + with cute.arch.elect_one(): + for i in range(stages): + cute.arch.mbarrier_init(bar_full + i, 2) + cute.arch.mbarrier_init(bar_empty + i, 1) + cute.arch.mbarrier_init(bar_tma_epilog, 32) + cute.arch.mbarrier_init(bar_mma_epilog, 1) + cute.arch.mbarrier_init(bar_tmem_alloc, 160) + if cutlass.const_expr(self.split_k > 1): + # Owner arrival plus peer transaction-byte credits. + cute.arch.mbarrier_init(bar_reduce, 1) + + cute.arch.mbarrier_init_fence() + if cutlass.const_expr(self.split_k > 1): + cute.arch.cluster_arrive_relaxed() # peers' barriers before their stores + else: + cute.arch.barrier() + + # K splits evenly over the cluster (tactics are validated for the shapes). + k_tile_count = cute.size(mA, mode=[1]) // self.cta_k // self.split_k + k_tile_start = split_rank * k_tile_count + + if cutlass.const_expr(self.split_k > 1): + cute.arch.cluster_wait() + + # Warp 3 is idle; warps 4-7 run the epilogue. + tiles = (k_tile_start, k_tile_count) + if warp_idx == 0: + self.weight_warp(bar_full, bar_empty, sA, a_map, bidx * self.cta_m, *tiles) + elif warp_idx == 1: + b_tile = cute.local_tile(mB, (self.cta_n, self.cta_k), (n_idx, None, l_idx)) + b_map = cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B") + self.activation_warp(bar_full, bar_empty, b_tile, sB, b_map, *tiles) + elif warp_idx == 2: + k_blocks = self.cta_k // cute.size(tiled_mma.shape_mnk, mode=[2]) + self.mma_warp( + bar_full, bar_empty, bar_mma_epilog, bar_tmem_alloc, tiled_mma, sA, sB, + tmem_base_ptr, acc_layout, k_blocks, k_tile_count, + ) # fmt: skip + elif warp_idx >= 4: + tile = cute.local_tile(mC, (self.cta_m, self.cta_n), (bidx, n_idx, l_idx)) + self.epilog_warp( + bar_mma_epilog, bar_tmem_alloc, tmem_base_ptr, acc_layout, tile, + cute.arch.thread_idx()[0] - 128, mC.element_type, + utils.LayoutEnum.from_tensor(mC), mailbox, bar_reduce, split_rank, + ) # fmt: skip + + @cute.experimental.jit + def weight_warp( + self, bar_full, bar_empty, s_tile: cute.Tensor, a_map: cutlass.Int64, + row0: cutlass.Int32, k_tile_start: cutlass.Int32, k_tile_count: cutlass.Int32, + ): # fmt: skip + """A: two evict_first boxes per K tile, before the PDL wait (weights).""" + stages = self.num_ab_stage + stage_bytes = cute.size_in_bytes( + s_tile.element_type, cute.slice_(s_tile.layout, (None, None, None, 0)) + ) + _acquire_tensor_map(a_map) + empty_phase = cutlass.Int32(1) + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % stages + cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(bar_full + stage, stage_bytes) + with cute.arch.elect_one(): + k0 = (k_tile_start + k_tile) * self.cta_k + tile = s_tile[None, None, None, stage].iterator + for atom in cutlass.range_constexpr(self.cta_k // ATOM_K): + box, mbar = tile + atom * self.cta_m * ATOM_K, bar_full + stage + k = k0 + atom * ATOM_K + _tma_load_evict_first(box, mbar, a_map, k, row0) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + + cute.arch.griddepcontrol_launch_dependents() + self.drain(bar_empty, empty_phase, k_tile_count) + + @cute.experimental.jit + def activation_warp( + self, bar_full, bar_empty, g_tile: cute.Tensor, s_tile: cute.Tensor, + cta_v_map: cute.Layout, k_tile_start: cutlass.Int32, + k_tile_count: cutlass.Int32, + ): # fmt: skip + """B (the activations): after the PDL wait.""" + stages = self.num_ab_stage + stage_bytes = cute.size_in_bytes( + s_tile.element_type, cute.slice_(s_tile.layout, (None, None, None, 0)) + ) + cute.arch.griddepcontrol_wait() + empty_phase = cutlass.Int32(1) + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % stages + cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(bar_full + stage, stage_bytes) + cute_ext.tma_load( + g_tile[None, None, k_tile_start + k_tile], + s_tile[None, None, None, stage], + (bar_full + stage).value, + cta_v_map=cta_v_map, + tma_operation_type=cute_ext.OperationTypeEnum.SM90_TMA_LOAD, + update_expect_tx=False, + ) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + self.drain(bar_empty, empty_phase, k_tile_count) + + @cute.experimental.jit + def drain(self, bar_empty, empty_phase: cutlass.Int32, k_tile_count: cutlass.Int32): + stages = self.num_ab_stage + for tail in cutlass.range(stages, unroll=1): + stage = (tail + k_tile_count) % stages + cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + + @cute.experimental.jit + def mma_warp( + self, bar_full, bar_empty, bar_mma_epilog, bar_tmem_alloc, + tiled_mma: cute.TiledMma, sA: cute.Tensor, sB: cute.Tensor, tmem_base_ptr, + acc_layout: cutlass.Constexpr, mma_inst_tile_k: cutlass.Constexpr, + k_tile_count: cutlass.Int32, + ): # fmt: skip + num_tmem_cols = 256 + cute.arch.alloc_tmem(num_tmem_cols, tmem_base_ptr, is_two_cta=False) + _fence_tcgen05(before=True) + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.relinquish_tmem_alloc_permit(is_two_cta=False) + + tmem_ptr = cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr) + accumulator = cute.make_tensor(tmem_ptr, acc_layout)[None, None, None, 0] + mma_atom = cute.make_mma_atom(tiled_mma.op) + full_phase = cutlass.Int32(0) + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % self.num_ab_stage + cute.arch.mbarrier_wait(bar_full + stage, full_phase) + _fence_tcgen05() + for k_block in range(mma_inst_tile_k): + accumulate = k_tile != 0 if k_block == 0 else True + mma_atom.set(tcgen05.Field.ACCUMULATE, accumulate) + cute_ext.dot( + mma_atom, + cute.append_ones(sA[None, None, k_block, stage], up_to_rank=3), + cute.append_ones(sB[None, None, k_block, stage], up_to_rank=3), + accumulator, + ) + with cute.arch.elect_one(): + tcgen05.commit(bar_empty + stage, None, self.cta_group) + if stage == self.num_ab_stage - 1: + full_phase = full_phase ^ 1 + + with cute.arch.elect_one(): + tcgen05.commit(bar_mma_epilog, None, self.cta_group) + cute.arch.mbarrier_wait(bar_tmem_alloc, 0) # the release lands in phase 1 + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 1) + _fence_tcgen05() + cute.arch.dealloc_tmem(tmem_ptr, num_tmem_cols, is_two_cta=False) + + @cute.experimental.jit + def epilog_warp( + self, bar_mma_epilog, bar_tmem_alloc, tmem_base_ptr, + acc_layout: cutlass.Constexpr, gD_tile: cute.Tensor, epi_tid: cutlass.Int32, + c_dtype: cutlass.Constexpr, d_layout: cutlass.Constexpr, mailbox, bar_reduce, + split_rank: cutlass.Int32, + ): # fmt: skip + # Wait until MMA publishes the TMEM base pointer. + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 0) + _fence_tcgen05() + + acc_view = cute.make_tensor( + cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr), + acc_layout, + )[((None, None), 0, 0, 0)] + epi_tile = (self.cta_m, self.cta_n) + load_op = sm100_utils.get_tmem_load_op( + (self.cta_m, self.cta_n, self.cta_k), d_layout, c_dtype, self.acc_dtype, + epi_tile, False, + ) # fmt: skip + tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy(load_op, acc_view) + gD_epi = cute.flat_divide(gD_tile, epi_tile) + + # Match each epilogue thread's TMEM partition in RMEM. + rmem_layout = cute_ext.make_t2r_rmem_layout(tiled_copy_t2r, gD_epi, epi_tid) + rmem = cute.AddressSpace.rmem + rAcc = cute_ext.allocate(self.acc_dtype, rmem, rmem_layout, alignment=32) + rD = cute_ext.allocate(c_dtype, rmem, rmem_layout, alignment=32) + thr_t2r = tiled_copy_t2r.get_slice(epi_tid) + + cute.arch.mbarrier_wait(bar_mma_epilog, 0) + _fence_tcgen05() + cute_ext.partition_and_copy(thr_t2r, acc_view, rAcc) + # Make tcgen05.ld visible before TMEM release and RMEM use. + cute.arch.fence_view_async_tmem_load() + _fence_tcgen05(before=True) + cute.arch.mbarrier_arrive(bar_tmem_alloc) + + # Peers publish FP32 partials; only rank 0 reduces and stores. + if cutlass.const_expr(self.split_k > 1): + vpt = cutlass.const_expr(cute.size(rmem_layout)) # values per thread + vpp = cutlass.const_expr(self.epilog_threads * vpt) # values per peer + if split_rank != OWNER_RANK: + for i in cutlass.range_constexpr(0, vpt, 4): + words = (rAcc[i], rAcc[i + 1], rAcc[i + 2], rAcc[i + 3]) + slot = mailbox.iterator + (split_rank - Int32(1)) * vpp + slot = slot + epi_tid * vpt + i + _store_shared_remote_v4(*words, slot, bar_reduce, Int32(OWNER_RANK)) + else: + if epi_tid == 0: + expected = self.expected_transaction_bytes + cute.arch.mbarrier_arrive_and_expect_tx(bar_reduce, expected) + _wait_cluster(bar_reduce, Int32(0)) + for peer in cutlass.range_constexpr(self.split_k - 1): + for i in cutlass.range_constexpr(vpt): + rAcc[i] = rAcc[i] + mailbox[peer * vpp + epi_tid * vpt + i] + + if split_rank == OWNER_RANK: + rD.store(rAcc.load().to(c_dtype)) + # Preserve TMEM coordinates; the copy predicates output tails. + cute_ext.partition_and_copy(thr_t2r, rD, gD_epi[None, None, 0, 0]) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/ws_gemm.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/ws_gemm.py new file mode 100644 index 0000000000..260ad2b22e --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/gemm/ws_gemm.py @@ -0,0 +1,293 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Weight-stationary skinny GEMM for the MoE front: y[M, N] = x[M, K] @ W[N, K]^T. + +Each CTA owns ``rt`` 16-row N tiles; its ``rt * kq`` warps split K into ``kq`` ranges. +Weights enter registers and cp.async SMEM stages before the PDL wait. The epilogue +publishes the first ``valid`` of ``split`` rows (router shard) as FP32 words at BF16 +column 3584 + 2 * (``slot`` + row) of a multicast [1, capacity, width] mailbox and the +rest (latent down shard) as sanitized BF16 at column ``base``. + +``tg`` 0 is the single pass: warps load x from global and the cluster leader sums the +K partials over DSMEM. ``tg`` > 0 streams tokens in groups of ``tg``: warp 0 bulk-copies +the CTA's K slice of x once (one mbarrier per group) for all row tiles, and after each +group's MMAs every warp pushes its 8-token x 16-row partial blocks to their owner CTAs +(block % ks) with st.async; owners sum the ks * kq parts in K order and publish one +group behind, so only the last group's blocks follow the MMA loop. +""" + +from typing import Any + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Int64, Uint32, Uint64 + +from ..common.primitives import ( + alloc, + asm, + at, + bf16x2, + bits_float, + copy_g2s, + float_bits, + launch, + load_words, + mma_into, + multicast_words, + push_float, + smem_words, + wait, +) + + +@asm(Uint32, pure=True) +def positive_zero(low, high): + """BF16 (low, high) with -0.0 halves made +0.0, so it never equals the sentinel.""" + code = "{ .reg .b32 t; .reg .pred p,q; and.b32 t,$1,0xFFFF; " + code += "setp.eq.u32 p,t,0x8000; shr.u32 t,$1,16; setp.eq.u32 q,t,0x8000; " + code += "mov.b32 $0,$1; @p and.b32 $0,$0,0xFFFF0000; @q and.b32 $0,$0,0x0000FFFF; }" + return code, "=r,r", (Uint32(bf16x2(low, high)),) + + +@asm(Float32) +def peer_load(pointer, peer): + """FP32 at the same shared-memory offset in cluster CTA ``peer``.""" + code = "{ .reg .b32 a; mapa.shared::cluster.u32 a,$1,$2; " + return code + "ld.shared::cluster.f32 $0,[a]; }", "=f,r,r", (pointer, Int32(peer)) + + +@asm(Uint32, pure=True) +def positive_zero32(value): + """FP32 bits of ``value`` with -0.0 made +0.0, so it never equals the sentinel.""" + code = "{ .reg .pred p; setp.eq.u32 p,$1,0x80000000; selp.b32 $0,0,$1,p; }" + return code, "=r,r", (Uint32(float_bits(value)),) + + +class WsGemm: + def __init__( + self, n, k, m, rt, kq, staged, split, valid, ks=1, tg=0, ch=0 + ): # fmt: skip + self.n, self.k, self.m, self.rt, self.kq = n, k, m, rt, kq + # ks: a cluster of ks CTAs splits K, so each CTA preloads 1/ks of the tile's + # weights; ch: accumulator chains per 8-token tile (0: 4 / tiles per group). + self.split, self.valid, self.ks = split, valid, ks + self.staged, self.kb = staged, k // (32 * kq * ks) + self.rb, self.nt = self.kb - staged, (m + 7) // 8 + self.grouped, self.ng = tg > 0, min(self.nt, tg // 8) if tg else self.nt + if self.grouped and ks <= 1: # token groups push partials to K-split peers + raise ValueError(f"WsGemm token groups (tg={tg}) need a K-split cluster") + self.groups, self.chains = -(-self.nt // self.ng), ch or max(1, 4 // self.ng) + self.grid, self.threads = n // (16 * rt) * ks, 32 * rt * kq + self.blocks, self.parts = self.nt * rt, ks * kq + self.slots = -(-self.blocks // ks) # owned blocks s * ks + kpart + # Bulk x rows: the CTA's K slice plus 64 B, so token pairs hit disjoint banks. + self.slice = k // ks + self.pitch = self.slice // 2 + 16 + + @cute.jit + def __call__(self, *args): + launch(self.kernel, args, self.grid, self.threads, self.ks, min_blocks_per_mp=1) + + @cute.kernel + def kernel(self, args): + x, w = args[0:2] + smem = cutlass.utils.SmemAllocator() + block, _, _ = cute.arch.block_idx() + cta, kpart = block // self.ks, block % self.ks + tid, _, _ = cute.arch.thread_idx() + lane, warp = tid % 32, tid // 32 + group, part = lane // 4, lane % 4 + tile, quarter = warp // self.kq, warp % self.kq + row = (cta * self.rt + tile) * 16 + group + kofs = (kpart * self.kq + quarter) * self.kb * 32 + part * 8 + shape = (self.nt, self.rt, self.kq, 8, 17) + red = smem.allocate_tensor(Float32, cute.make_layout(shape), 16) + stage = alloc(smem, Uint32, (self.rt * self.kq, max(self.staged, 1), 2, 32, 4)) + xs = recv = None + bars: Any = None + owned: Any = None + if cutlass.const_expr(self.grouped): + xs = alloc(smem, Uint32, (self.nt * 8, self.pitch)) + bars = alloc(smem, Uint64, (self.groups,), 8) + recv = alloc(smem, Uint32, (self.slots, self.parts, 32, 4)) + owned = alloc(smem, Uint64, (self.slots,), 8) + if tid == 0: + for g in cutlass.range_constexpr(self.groups): + cute.arch.mbarrier_init(bars.iterator + g, 1) + for s in cutlass.range_constexpr(self.slots): + if s * self.ks + kpart < self.blocks: + cute.arch.mbarrier_init(owned.iterator + s, 1) + nbytes = self.parts * 512 + cute.arch.mbarrier_arrive_and_expect_tx( + owned.iterator + s, nbytes + ) + cute.arch.mbarrier_init_fence() + cute.arch.sync_threads() + cute.arch.cluster_arrive_relaxed() + wp = w.iterator + row * self.k + kofs + for j in cutlass.range_constexpr(self.staged): + for h in cutlass.range_constexpr(2): + dst = at(stage, warp, j, h, lane, 0) + copy_g2s(dst, wp + h * 8 * self.k + (self.rb + j) * 32) + cute.arch.cp_async_commit_group() + wr = {} + for j in cutlass.range_constexpr(self.rb): + for h in cutlass.range_constexpr(2): + src = wp + h * 8 * self.k + j * 32 + wr[j, h] = load_words(src, count=4, evict_first=True) + cute.arch.griddepcontrol_wait() + cute.arch.griddepcontrol_launch_dependents() + if cutlass.const_expr(self.grouped): + if warp == 0: + self.stage_x(x, xs, bars, kpart, lane) + cute.arch.cp_async_wait_group(0) + else: + cute.arch.cp_async_wait_group(0) + cute.arch.sync_threads() + for g in cutlass.range_constexpr(self.groups): + acc = cute.make_rmem_tensor((self.ng, self.chains, 4), Float32) + acc.fill(0) + if cutlass.const_expr(self.grouped): + wait(bars.iterator + g) + for j in cutlass.range_constexpr(self.kb): + if cutlass.const_expr(j < self.rb): + wa, wb = wr[j, 0], wr[j, 1] + else: + sa = at(stage, warp, j - self.rb, 0, lane, 0) + wa, wb = smem_words(sa), smem_words(sa + 128) + for t in cutlass.range_constexpr(self.tiles(g)): + token = (g * self.ng + t) * 8 + group + xv = (Uint32(0), Uint32(0), Uint32(0), Uint32(0)) + if cutlass.const_expr(self.grouped): + col = (quarter * self.kb + j) * 16 + part * 4 + xv = smem_words(at(xs, token, col)) + else: + xg = x.iterator + token * x.stride[0] + kofs + j * 32 + if (g * self.ng + t) * 8 + 7 < self.m or token < self.m: + xv = load_words(xg, count=4, mutable=True) + for h in cutlass.range_constexpr(2): + c = (j * 2 + h) % self.chains + a = (wa[2 * h], wb[2 * h], wa[2 * h + 1], wb[2 * h + 1]) + b = (xv[2 * h], xv[2 * h + 1]) + mma_into(acc[t, c, None], a, b, bf16=True) + if cutlass.const_expr(self.grouped and g == 0): + cute.arch.cluster_wait() # owners' barriers are armed + for t in cutlass.range_constexpr(self.tiles(g)): + values = [] + for v in cutlass.range_constexpr(4): + value = acc[t, 0, v] + for c in cutlass.range_constexpr(1, self.chains): + value = value + acc[t, c, v] + tt, col = g * self.ng + t, group + (v // 2) * 8 + if cutlass.const_expr(self.grouped): + values.append(value) + else: + red[tt, tile, quarter, 2 * part + v % 2, col] = value + if cutlass.const_expr(self.grouped): + blk = (g * self.ng + t) * self.rt + tile + s, peer = blk // self.ks, blk % self.ks + dst = at(recv, s, kpart * self.kq + quarter, lane, 0) + push_float(dst, tuple(values), owned.iterator + s, peer) + if cutlass.const_expr(self.grouped and g > 0): + self.settle(args, red, recv, owned, cta, kpart, warp, lane, g - 1) + if cutlass.const_expr(self.grouped): + self.settle(args, red, recv, owned, cta, kpart, warp, lane, self.groups - 1) + else: + cute.arch.sync_threads() + if cutlass.const_expr(self.ks > 1): + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + self.gather(args, red, cta, kpart, tid, lane) + if cutlass.const_expr(self.ks > 1): + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + + def tiles(self, g): + """8-token tiles in token group ``g``.""" + return min(self.ng, self.nt - g * self.ng) + + @cute.jit + def stage_x(self, x, xs, bars, kpart, lane): + """Warp 0 bulk-copies every token's K slice; one barrier per token group.""" + if lane == 0: + for g in cutlass.range_constexpr(self.groups): + tokens = min(self.m, (g + 1) * self.ng * 8) - g * self.ng * 8 + nbytes = tokens * self.slice * 2 + cute.arch.mbarrier_arrive_and_expect_tx(bars.iterator + g, nbytes) + cute.arch.sync_warp() + for i in cutlass.range_constexpr(-(-self.m // 32)): + token = i * 32 + lane + if i * 32 + 31 < self.m or token < self.m: + src = x.iterator + token * x.stride[0] + kpart * self.slice + bar = bars.iterator + token // (8 * self.ng) + copy_g2s(at(xs, token, 0), src, self.slice * 2, bar) + + @cute.jit + def settle( + self, args, red, recv, owned, cta, kpart, warp, lane, g: cutlass.Constexpr + ): + """Owned blocks of token group ``g``: a warp per block waits for every K part, + sums them in K order, transposes through ``red`` and publishes.""" + for t in cutlass.range_constexpr(self.tiles(g)): + for tile in cutlass.range_constexpr(self.rt): + blk = (g * self.ng + t) * self.rt + tile + s = blk // self.ks + if (kpart == blk % self.ks) & (warp == s % (self.rt * self.kq)): + wait(owned.iterator + s) + total = cute.make_rmem_tensor(4, Float32) + for p in cutlass.range_constexpr(self.parts): + words = smem_words(at(recv, s, p, lane, 0)) + for v in cutlass.range_constexpr(4): + value = bits_float(words[v]) + total[v] = value if p == 0 else total[v] + value + tt = g * self.ng + t + for v in cutlass.range_constexpr(4): + col = lane // 4 + v // 2 * 8 + red[tt, tile, 0, 2 * (lane % 4) + v % 2, col] = total[v] + cute.arch.sync_warp() + token = tt * 8 + lane // 2 + if (lane < 16) & (token < self.m): + half = lane % 2 * 8 + values = [ + red[tt, tile, 0, lane // 2, half + j] for j in range(8) + ] + col = (cta * self.rt + tile) * 16 + half + self.epilogue(args, token, col, values) + + @cute.jit + def gather(self, args, red, cta, kpart, tid, lane): + """Each thread sums one column's K parts in order (the leader also its peers'); + each eight-lane group then gathers the sums for its first lane's 16 B store.""" + total = self.m * self.rt * 16 + for start in range(0, total, self.threads): + item = start + tid + v8, half = item % 8, (item // 8) % 2 + rtile, token = (item // 16) % self.rt, item // (16 * self.rt) + mine = Float32(0) + if item < total: + mine = red[token // 8, rtile, 0, token % 8, half * 8 + v8] + for q in cutlass.range_constexpr(1, self.kq): + mine = mine + red[token // 8, rtile, q, token % 8, half * 8 + v8] + if kpart == 0: + for peer in cutlass.range_constexpr(1, self.ks): + for q in cutlass.range_constexpr(self.kq): + idx = (token // 8, rtile, q, token % 8, half * 8 + v8) + mine = mine + peer_load(at(red, *idx), peer) + values = [cute.arch.shuffle_sync(mine, (lane & 24) + v) for v in range(8)] + if (item < total) & (v8 == 0) & (kpart == 0): + col = (cta * self.rt + rtile) * 16 + half * 8 + self.epilogue(args, token, col, values) + + @cute.jit + def epilogue(self, args, token, col, values): + out, mc, slot, base = args[2:6] + if col < self.valid: + for v in cutlass.range_constexpr(2): + words = tuple(positive_zero32(values[4 * v + i]) for i in range(4)) + offset = token * out.shape[2] + 3584 + 2 * (slot + col + 4 * v) + multicast_words(mc + Int64(offset) * 2, words, kind="bf16x2") + elif col >= self.split: + column = base + col - self.split + words = tuple(positive_zero(*values[i : i + 2]) for i in range(0, 8, 2)) + offset = token * out.shape[2] + column + multicast_words(mc + Int64(offset) * 2, words, kind="bf16x2") diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/__init__.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/kernel.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/kernel.py new file mode 100644 index 0000000000..7caab11e1a --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/kernel.py @@ -0,0 +1,715 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Cluster-split KDA recurrence; an optional fifth warp projects F_B and multicasts the +decays while the main warps run the convolution and query/key normalization.""" + +from functools import partial + +import cutlass +import cutlass.cute as cute +from cutlass import BFloat16, Float32, Int32, Int64, Uint32 +from cutlass.cute.arch import store_async_dsmem + +from ..common.primitives import ( + alloc, + asm, + at, + copy_g2s, + dot, + float_bits, + fma, + launch, + load_f32x4, + load_words, + mma_into, + push_float, + sigmoid, + smem_words, + store_f32x4, + unpack_bf16, + wait, + warp_sum, + zeros, +) + +load_bf16x4 = partial(load_words, count=2, cache_policy="nc", word_type="u32") +load_weight_bf16x8 = partial(load_words, count=4, cache_policy="nc", word_type="u32") +load_activation_bf16x4 = partial(load_words, count=2, word_type="u32", mutable=True) +load_activation_bf16x8 = partial(load_words, count=4, word_type="u32", mutable=True) +load_staged = partial( + load_words, count=4, shared=True, mutable=True, result_type=Float32, word_type="f32" +) + + +@asm((Float32, Int32), pure=True) +def reciprocal_guess(denominator): + """``sigmoid_reciprocal``'s refined fast path (IEEE mode) and whether it holds.""" + code = "{ .reg .f32 r,error,neg_error; .reg .b32 bits; .reg .pred regular; " + code += "mov.b32 bits,$2; add.u32 bits,bits,0x01800000; " + code += "and.b32 bits,bits,0x7f800000; setp.gt.u32 regular,bits,0x01ffffff; " + code += "rcp.approx.ftz.f32 r,$2; fma.rn.f32 error,$2,r,0fbf800000; " + code += "neg.f32 neg_error,error; fma.rn.f32 $0,r,neg_error,r; " + return code + "selp.u32 $1,1,0,regular; }", "=f,=r,f", (denominator,) + + +@asm(Float32, pure=True) +def reciprocal_rn(value): + return "div.rn.f32 $0,0f3f800000,$1;", "=f,f", (value,) + + +@cute.jit +def sigmoids(values): + """``sigmoid(v, False)`` of every value, bit for bit, with all fast paths ahead + of the division fallbacks' branches (verify tokens get ILP).""" + count: cutlass.Constexpr = len(values) + den = [Float32(1) + cute.math.exp(-Float32(v), fastmath=True) for v in values] + result, regular = zeros(count, Float32), zeros(count, Int32) + irregular = Int32(0) + for i in cutlass.range_constexpr(count): + result[i], regular[i] = reciprocal_guess(den[i]) + irregular |= regular[i] ^ 1 + if irregular != 0: # rare: one branch for the whole batch + for i in cutlass.range_constexpr(count): + if regular[i] == 0: + result[i] = reciprocal_rn(den[i]) + return result + + +def advance(base, n): + """``base + n``, or ``base`` itself for a static zero ``n`` (keeps q = 1 code).""" + return base + n if n else base + + +@asm() +def store_evict_f32x4(pointer, values): + code = "{ .reg .b64 pol; createpolicy.fractional.L2::evict_first.b64 pol,1.0; " + code += "st.global.L2::cache_hint.v4.f32 [$0],{$1,$2,$3,$4},pol; }" + return code, "l,f,f,f,f", (pointer, *values) + + +class KdaSplit: + """128 main threads and C value partitions. Conv state [slot, 3 * heads * 128, 3] + and recurrent state [slot, head, value, key] keep their strides; slots <= 0 are + padding (never written) and active rows own distinct slots. Weights must be ready + at launch; each run waits for prior cache restoration and slot publication. + + ``schedule`` = (early, lanes per state row, minimum CTAs per SM, CTAs per head); + zeros keep the batch defaults. Two CTAs per head give each CTA 64 state rows and + each gate lane two decays. Early issues every load the predecessor does not + produce (slot, state and conv rows, conv / F_B / norm weights) before the PDL + wait, converts operands at first use, and syncs the cluster only before each + role's first DSMEM send. Safe because those loads read only weights and data a + graph or non-PDL boundary orders before this launch (states and conv rows of this + layer's previous call or a fresh prompt's non-PDL fill, the host's slot table), + never same-step producer output, which is read after the wait; PDL-chained calls + without such a boundary must not share state slots. Early 2 also bulk-copies + the CTA's (contiguous) state rows to shared memory and updates four passes at a + time, which frees the registers that held them across the input phase; early 3 + (two CTAs per head) also stages the gate warp's F_B rows. + + ``tokens`` q > 1 (speculative verify; ``batch`` counts requests, rows b * q + t) + runs each request's q tokens in order with vLLM's spec cache semantics: slots + [B, >= q] and accepted [B] (an extra last tensor, read early like the slots); the + state comes from slot[b, accepted - 1] and token t's state goes to slot[b, t] if + > 0; conv slot[b, 0] holds [dim, q + 2] with taps at accepted - 1 ..., rewritten to + [c_accepted, c_accepted + 1, x_0, ..., x_q-1]. Token t's F_B row is MMA column t; + CTA t % C normalizes token t.""" + + def __init__(self, heads, batch, gate_mode, has_bias, schedule=(0, 0, 0), tokens=1): + # The gate stage sends decays for tokens 2 * quarter + column (at most 8), while + # the decay barrier expects 512 bytes per token: more tokens would never arrive. + if not 1 <= tokens <= 8: + raise ValueError(f"KdaSplit supports 1-8 tokens per request; got {tokens}") + self.heads, self.gate_mode, self.has_bias = heads, gate_mode, has_bias + early, lanes, self.occupancy, cluster, evict = (*schedule, 0, 0)[:5] + self.early, self.staged = bool(early), early >= 2 + self.gate_staged, self.tokens = early == 3, tokens + assert gate_mode != 2 or not (early or cluster) + assert tokens == 1 or not (self.staged or gate_mode == 2 or cluster == 16) + # TP8 (12 heads) uses fast math; B1 releases dependents at entry. + self.fast_math, self.trigger_at_entry = heads == 12, batch == 1 + self.cluster = cluster or (8 if batch == 1 else 4) + self.lanes_per_row = lanes or (8 if batch <= (4 if heads == 6 else 2) else 16) + self.rows, self.groups = 128 // self.cluster, 128 // self.lanes_per_row + self.passes = (self.rows + self.groups - 1) // self.groups + self.wave = min(4, self.passes) if self.staged else self.passes + self.chunks = 32 // self.lanes_per_row + assert not self.gate_staged or self.rows >= 64 + self.evict = bool(evict) # staged state traffic marked L2 evict-first + assert self.staged or not self.evict + + @cute.jit + def __call__(self, *args): + grid = (self.cluster, self.heads, args[0].shape[0] // self.tokens) + bound = dict(min_blocks_per_mp=self.occupancy) if self.occupancy else {} + launch(self.kernel, args, grid, 160, self.cluster, **bound) + + @cute.kernel + def kernel(self, args): + x, weights, bias, conv_state, raw_gate, beta, a_log, dt_bias, slots = args[:9] + state, output, gate, norm_weight, lower_bound, scale, eps, f_b_weight = args[ + 9:17 + ] + T: cutlass.Constexpr = self.tokens + if cutlass.const_expr(not self.early): + cute.arch.griddepcontrol_wait() + thread, _, _ = cute.arch.thread_idx() + partition, head, token = cute.arch.block_idx() + if cutlass.const_expr(T == 1): + slot = slots[token] + conv_slot, offset, writes = slot, 0, [slot] + else: + request, token = token, token * T + offset = args[17][request] - 1 + slot, conv_slot = slots[request, offset], slots[request, 0] + writes = [slots[request, t] for t in range(T)] + if slot <= 0: + cute.arch.griddepcontrol_wait() + if (partition == 0) & (thread < 128): + for t in cutlass.range_constexpr(T): + output[advance(token, t), head * 128 + thread] = BFloat16(0) + cute.arch.griddepcontrol_launch_dependents() + else: + lane, warp = thread % 32, thread // 32 + group = thread // self.lanes_per_row + group_lane = thread % self.lanes_per_row + dim: cutlass.Constexpr = self.heads * 128 + value_base, qk_channel = partition * self.rows, head * 128 + thread + value_channel = head * 128 + value_base + thread + + smem = cutlass.utils.SmemAllocator() + shared_q, shared_k, shared_decay, shared_v, shared_o, reductions = ( + alloc(smem, Float32, size * T) + for size in (128, 128, 128, self.rows, 128, 8) + ) + barrier, gate_barrier = (smem.allocate(cutlass.Uint64, 8) for _ in range(2)) + owned: cutlass.Constexpr = -(-T // self.cluster) # tokens a CTA normalizes + if cutlass.const_expr(owned > 1): # one output barrier per owned token + barrier = smem.allocate(8 * owned, 8) + barrier = cute.recast_ptr(barrier, dtype=cutlass.Uint64) + if cutlass.const_expr(self.staged): + staged = alloc(smem, Float32, (self.rows, 128), align=128) + state_barrier = smem.allocate(cutlass.Uint64, 8) + gate_rows = None + if cutlass.const_expr(self.gate_staged): # [copy, lane] 16-byte slots + gate_rows = alloc(smem, Uint32, (self.rows // 2, 32, 4), align=128) + if cutlass.const_expr(self.trigger_at_entry): + cute.arch.griddepcontrol_launch_dependents() + if cutlass.const_expr(not self.early): + self.setup(barrier, gate_barrier, partition, thread) + cute.arch.cluster_wait() + if thread >= 128: + if cutlass.const_expr(self.early): + cute.arch.cluster_arrive_relaxed() + coords = (head, token, value_base, lane) + self.gate_warp(args, shared_decay, gate_barrier, coords, gate_rows) + + if thread < 128: + if cutlass.const_expr(self.early): # thread 0 arms; barrier 1 orders it + self.setup(barrier, gate_barrier, partition, thread) + # 64-bit page bases (slot x page stride passes Int32), 32-bit offsets. + page = lambda s: state[Int64(s), None, None, None] + state_page = page(slot) + conv_page = conv_state[Int64(conv_slot), None, None] + state_row = lambda r, c, s=state_page: at(s, head, value_base + r, c) + history = zeros((self.passes, self.chunks, 4), Float32) + if cutlass.const_expr(self.staged): # noqa: SIM102 + if thread == 0: + cute.arch.mbarrier_init(state_barrier, 1) + cute.arch.mbarrier_init_fence() + size = self.rows * 512 + cute.arch.mbarrier_arrive_and_expect_tx(state_barrier, size) + src, policy = state_row(0, 0), dict(evict_first=self.evict) + copy_g2s(staged.iterator, src, size, state_barrier, **policy) + for p in cutlass.range_constexpr(0 if self.staged else self.passes): + row = p * self.groups + group + if row < self.rows: + for j in cutlass.range_constexpr(self.chunks): + column = 4 * (group_lane + self.lanes_per_row * j) + values = load_f32x4(state_row(row, column)) + for e in cutlass.range_constexpr(4): + history[p, j, e] = values[e] + taps = zeros((3, 3), BFloat16) + coeff = zeros((3, 4), weights.element_type if self.early else Float32) + accum = zeros(3 * T, Float32) + for s in cutlass.range_constexpr(3): + channel = value_channel if s == 2 else qk_channel + if (s != 2) | (thread < self.rows): + for t in cutlass.range_constexpr(3): + taps[s, t] = conv_page[s * dim + channel, offset + t] + for t in cutlass.range_constexpr(4): + coeff[s, t] = weights[s, t, channel] + if cutlass.const_expr(self.has_bias): + accum[s] = bias[s * dim + channel] + for t in cutlass.range_constexpr(1, T): + accum[3 * t + s] = accum[s] + nw = norm_weight[thread] + + cute.arch.griddepcontrol_wait() + inputs = cute.make_rmem_tensor(3 * T, BFloat16) + betas, gates = [], [] + for t in cutlass.range_constexpr(T): + row = advance(token, t) + inputs[3 * t], inputs[3 * t + 1] = ( + x[row, qk_channel], + x[row, dim + qk_channel], + ) + inputs[3 * t + 2] = BFloat16(0) + if cutlass.const_expr(t == T - 1): # noqa: SIM102 (all tokens) + if thread < self.rows: + for u in cutlass.range_constexpr(T): + channel = 2 * dim + value_channel + inputs[3 * u + 2] = x[advance(token, u), channel] + if cutlass.const_expr(self.early): # converted at first use + betas.append(beta[row, head]) + gates.append(gate[row, qk_channel]) + else: + betas.append(Float32(beta[row, head])) + gates.append(Float32(gate[row, qk_channel])) + cute.arch.griddepcontrol_launch_dependents() + + vs, squares, qks = [], [], [] + for t in cutlass.range_constexpr(T): # window: 3 taps, then the inputs + for s in cutlass.range_constexpr(3): + for i in cutlass.range_constexpr(4): + w = t + i + term = taps[s, w] if w < 3 else inputs[3 * (w - 3) + s] + term = Float32(term) + cf = Float32(coeff[s, i]) if self.early else coeff[s, i] + acc = accum[3 * t + s] + accum[3 * t + s] = fma(term, cf, acc, ftz=self.fast_math) + act = lambda i, value: sigmoid(value, self.fast_math) + if cutlass.const_expr(T > 1 and not self.fast_math): + raw = [accum[i] for i in range(3 * T)] + betas + gates + act = partial(lambda batch, i, value: batch[i], sigmoids(raw)) + for t in cutlass.range_constexpr(T): + q = accum[3 * t] * act(3 * t, accum[3 * t]) + k = accum[3 * t + 1] * act(3 * t + 1, accum[3 * t + 1]) + v = accum[3 * t + 2] * act(3 * t + 2, accum[3 * t + 2]) + betas[t] = act(3 * T + t, Float32(betas[t])) + gates[t] = act(4 * T + t, Float32(gates[t])) + shared_q[advance(thread, 128 * t)] = q + shared_k[advance(thread, 128 * t)] = k + vs.append(v) + qks += [q, k] + if cutlass.const_expr(t == T - 1): # noqa: SIM102 + if thread < self.rows: + for u in cutlass.range_constexpr(T): + shared_v[advance(thread, self.rows * u)] = vs[u] + squares += [warp_sum(q * q), warp_sum(k * k)] + if cutlass.const_expr(t == T - 1): # noqa: SIM102 + if lane == 0: + for u in cutlass.range_constexpr(T): + reductions[advance(warp, 8 * u)] = squares[2 * u] + reductions[advance(warp + 4, 8 * u)] = squares[ + 2 * u + 1 + ] + cute.arch.barrier(barrier_id=1, number_of_threads=128) + + def factors(t): + r = lambda i: reductions[8 * t + i] + qn = (r(0) + r(2)) + (r(1) + r(3)) + kn = (r(4) + r(6)) + (r(5) + r(7)) + fq = cute.math.rsqrt(qn + 1.0e-6, fastmath=True) * scale + return fq, cute.math.rsqrt(kn + 1.0e-6, fastmath=True) + + if cutlass.const_expr(T > 1): # normalize each token's q, k once + for t in cutlass.range_constexpr(T): + fq, fk = factors(t) + shared_q[128 * t + thread] = qks[2 * t] * fq + shared_k[128 * t + thread] = qks[2 * t + 1] * fk + cute.arch.barrier(barrier_id=1, number_of_threads=128) + + def operands(t): + fq, fk = factors(t) if T == 1 else (None, None) + qr = cute.make_rmem_tensor((self.chunks, 4), Float32) + kr = cute.make_rmem_tensor((self.chunks, 4), Float32) + for j in cutlass.range_constexpr(self.chunks): + column = 4 * (group_lane + self.lanes_per_row * j) + q4 = load_f32x4( + advance(shared_q.iterator, 128 * t) + column, shared=True + ) + k4 = load_f32x4( + advance(shared_k.iterator, 128 * t) + column, shared=True + ) + for e in cutlass.range_constexpr(4): + qr[j, e] = q4[e] * fq if T == 1 else q4[e] + kr[j, e] = k4[e] * fk if T == 1 else k4[e] + if cutlass.const_expr(t == 0): + wait(gate_barrier) + decays = [] + for j in cutlass.range_constexpr(self.chunks): + column = 4 * (group_lane + self.lanes_per_row * j) + decay_ptr = advance(shared_decay.iterator, 128 * t) + column + decays.append(load_f32x4(decay_ptr, shared=True)) + values = shared_v + if cutlass.const_expr(T > 1): + pointer = shared_v.iterator + self.rows * t + values = cute.make_tensor(pointer, cute.make_layout(self.rows)) + return kr, qr, decays, values, betas[t] + + ready = operands(0) + for t in cutlass.range_constexpr(T): + current, ready = ready, operands(t + 1) if t + 1 < T else None + # Butterfly order: interleave rows per level, publish before state. + dots = cute.make_rmem_tensor(self.passes, Float32) + if cutlass.const_expr(self.staged): + wait(state_barrier) + source = staged if self.staged else None + owner, bar = t % self.cluster, advance(barrier, t // self.cluster) + for start in cutlass.range_constexpr(0, self.passes, self.wave): + stop = min(start + self.wave, self.passes) + lanes = (group, group_lane) + self.update( + history, dots, current, lanes, (start, stop), source + ) + if cutlass.const_expr(self.early and start == 0 and t == 0): + cute.arch.cluster_wait() + for p in cutlass.range_constexpr(start, stop): + row = p * self.groups + group + send = (row < self.rows) & (group_lane == 0) + if cutlass.const_expr(T > 1): # predicated: no branches + position = self.output_position(value_base + row) + dst = advance(shared_o.iterator, 128 * t) + position + push_float(dst, dots[p], bar, owner, when=send) + elif send: + position = self.output_position(value_base + row) + dst = shared_o.iterator + position + store_async_dsmem(dst, float_bits(dots[p]), bar, owner) + rows = partial(state_row, s=page(writes[t])) + when = None if T == 1 else writes[t] > 0 + self.store_rows(history, rows, (start, stop), lanes, when) + if thread < self.rows: + channel = 2 * dim + value_channel + window = ( + taps[2, 1], + taps[2, 2], + *(inputs[3 * t + 2] for t in range(T)), + ) + for i, value in enumerate(window): + conv_page[channel, i] = value + + for t in cutlass.range_constexpr(T): + if partition == t % self.cluster: + wait(advance(barrier, t // self.cluster)) + raw_o = shared_o[advance(self.output_position(thread), 128 * t)] + sums = cute.make_rmem_tensor(4, Float32) + for g in cutlass.range_constexpr(4): + base = lane % 8 * 16 + g * 4 + o_row = advance(shared_o.iterator, 128 * t) + o4 = load_f32x4(o_row + base, shared=True) + s0, s1, s2, s3 = (o * o for o in o4) + sums[g] = warp_sum((s0 + s2) + (s1 + s3), 8) + sumsq = (sums[0] + sums[2]) + (sums[1] + sums[3]) + rstd = cute.math.rsqrt(sumsq * (1.0 / 128) + eps, fastmath=True) + result = raw_o * rstd * nw * gates[t] + output[advance(token, t), qk_channel] = BFloat16(result) + if cutlass.const_expr(t == 0): + for section in cutlass.range_constexpr(2): + channel = section * dim + qk_channel + window = (taps[section, 1], taps[section, 2]) + window += tuple( + inputs[3 * u + section] for u in range(T) + ) + for i, value in enumerate(window): + conv_page[channel, i] = value + + @cute.jit + def store_rows(self, history, state_row, span, lanes, when=None): + """Store the passes in ``span`` of the CTA's state rows (``when``: predicated + on it, without branches).""" + (start, stop), (group, group_lane) = span, lanes + for p in cutlass.range_constexpr(start, stop): + row = p * self.groups + group + if cutlass.const_expr(when is not None): + for j in cutlass.range_constexpr(self.chunks): + column = 4 * (group_lane + self.lanes_per_row * j) + values = [history[p, j, e] for e in range(4)] + pointer = state_row(row % self.rows, column) + store_f32x4(pointer, values, when=when & (row < self.rows)) + elif row < self.rows: + for j in cutlass.range_constexpr(self.chunks): + column = 4 * (group_lane + self.lanes_per_row * j) + pointer = state_row(row, column) + row_values = [history[p, j, e] for e in range(4)] + if cutlass.const_expr(self.evict): + store_evict_f32x4(pointer, row_values) + else: + store_f32x4(pointer, row_values) + + @cute.jit + def update(self, history, dots, operands, lanes, span, staged): + """Decay, delta-rule update and query dot of the passes in ``span``, whose rows + ``staged`` (if given) first supplies from shared memory.""" + kr, qr, decays, shared_v, beta = operands + (group, lane), (start, stop) = lanes, span + if cutlass.const_expr(staged is not None): + for p in cutlass.range_constexpr(start, stop): + row = (p * self.groups + group) % self.rows + for j in cutlass.range_constexpr(self.chunks): + column = 4 * (lane + self.lanes_per_row * j) + values = load_staged(staged.iterator + (row * 128 + column)) + for e in cutlass.range_constexpr(4): + history[p, j, e] = values[e] + for p in cutlass.range_constexpr(start, stop): + for j in cutlass.range_constexpr(self.chunks): + for e in cutlass.range_constexpr(4): + history[p, j, e] *= decays[j][e] + dots[p] = self.row_dot(history, kr, p) + for shift in reversed(range(self.lanes_per_row.bit_length() - 1)): + for p in cutlass.range_constexpr(start, stop): + dots[p] += cute.arch.shuffle_sync_bfly(dots[p], 1 << shift) + for p in cutlass.range_constexpr(start, stop): + safe_row = (p * self.groups + group) % self.rows + v_new = (shared_v[safe_row] - dots[p]) * beta + for j in cutlass.range_constexpr(self.chunks): + for e in cutlass.range_constexpr(4): + history[p, j, e] = fma( + kr[j, e], v_new, history[p, j, e], ftz=self.fast_math + ) + dots[p] = self.row_dot(history, qr, p) + for shift in reversed(range(self.lanes_per_row.bit_length() - 1)): + for p in cutlass.range_constexpr(start, stop): + dots[p] += cute.arch.shuffle_sync_bfly(dots[p], 1 << shift) + + @cute.jit + def gate_warp(self, args, shared_decay, barrier, coords, staged_rows): + """Project one value partition and multicast its decays to every CTA. Modes fix + the F_B reduction order: 0 = one chain of eight MMAs, 1 = four chains of two, + 2 = sixteen 8-FMA partials summed in a butterfly. Projection math stays IEEE. + ``staged_rows``: shared [copy, lane, 4] words; each lane cp.asyncs its own + 16-byte F_B operands there before the PDL wait and reads them back after it.""" + f_a, a_log, dt_bias, lower_bound, weight = (args[i] for i in (4, 6, 7, 13, 16)) + head, token, value_base, lane = coords + rows, mode = self.rows, self.gate_mode + per: cutlass.Constexpr = max(1, rows // 32) # decays per lane + a = a_log[head] + if cutlass.const_expr(mode == 2): + feature = lane % rows + row = head * 128 + value_base + feature + # uint4 weights before the wait, two uint2 activation loads per 8-FMA + # partial after it; keep that granularity and the low-then-high FMA order. + wr = [load_weight_bf16x8(at(weight, row, 8 * i)) for i in range(16)] + dt = dt_bias[row] + cute.arch.griddepcontrol_wait() + partial = cute.make_rmem_tensor(16, Float32) + for i in cutlass.range_constexpr(16): + a0 = load_activation_bf16x4(at(f_a, token, 8 * i)) + a1 = load_activation_bf16x4(at(f_a, token, 8 * i + 4)) + av = (*a0, *a1) + acc = Float32(0) + for e in cutlass.range_constexpr(4): + x0, x1 = unpack_bf16(av[e]) + w0, w1 = unpack_bf16(wr[i][e]) + acc = fma(x0, w0, acc) + acc = fma(x1, w1, acc) + partial[i] = acc + for shift in cutlass.range_constexpr(3, -1, -1): + for i in cutlass.range_constexpr(1 << shift): + partial[i] += partial[i | (1 << shift)] + raw_g = [Float32(BFloat16(partial[0]))] + dt = [dt] + else: + tiles: cutlass.Constexpr = rows // 16 + chains: cutlass.Constexpr = 4 if mode == 1 else 1 + group, quarter = lane >> 2, lane & 3 + feature = 16 * (quarter >> 1) + group + 8 * (quarter & 1) + dt = [Float32(0) for _ in range(per)] + if cutlass.const_expr(self.tokens > 1): # token column, its features + token += group % self.tokens + features = [ + 16 * (i // 2) + group + 8 * (i % 2) for i in range(rows // 8) + ] + dt = [dt_bias[head * 128 + value_base + f] for f in features] + for r in cutlass.range_constexpr(per if self.tokens == 1 else 0): + if feature + 32 * r < rows: + dt[r] = dt_bias[head * 128 + value_base + 32 * r + feature] + # Two decays per lane load 16 bytes per row: chunk c then covers columns + # 32 (c // 2) + 8 quarter + 4 (c % 2) + [0, 4) instead of 16 c + 4 quarter. + wide: cutlass.Constexpr = per >= 2 + wr = cute.make_rmem_tensor((tiles, 8, 4), Uint32) + if cutlass.const_expr(staged_rows is not None): + for tile in cutlass.range_constexpr(tiles): + row = head * 128 + value_base + 16 * tile + group + for chunk in cutlass.range_constexpr(4): + col = 32 * chunk + 8 * quarter + for half in cutlass.range_constexpr(2): + slot = staged_rows[ + (tile * 4 + chunk) * 2 + half, lane, None + ] + source = at(weight, row + 8 * half, col) + copy_g2s(slot.iterator, source) + cute.arch.cp_async_commit_group() + else: + for tile in cutlass.range_constexpr(tiles): + row = head * 128 + value_base + 16 * tile + group + for chunk in cutlass.range_constexpr(4 if wide else 8): + if cutlass.const_expr(wide): + col = 32 * chunk + 8 * quarter + low = load_weight_bf16x8(at(weight, row, col)) + high = load_weight_bf16x8(at(weight, row + 8, col)) + self.split_pairs(wr, tile, chunk, low, high) + else: + col = 16 * chunk + 4 * quarter + low = load_bf16x4(at(weight, row, col)) + high = load_bf16x4(at(weight, row + 8, col)) + wr[tile, chunk, 0], wr[tile, chunk, 2] = low + wr[tile, chunk, 1], wr[tile, chunk, 3] = high + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(wide): + pairs = [ + load_activation_bf16x8(at(f_a, token, 32 * c + 8 * quarter)) + for c in range(4) + ] + ar = [pairs[c // 2][2 * (c % 2) : 2 * (c % 2) + 2] for c in range(8)] + else: + ar = [ + load_activation_bf16x4(at(f_a, token, 16 * c + 4 * quarter)) + for c in range(8) + ] + if cutlass.const_expr(staged_rows is not None): + cute.arch.cp_async_wait_group(0) + for tile in cutlass.range_constexpr(tiles): + for chunk in cutlass.range_constexpr(4): + copy = (tile * 4 + chunk) * 2 + low = smem_words(staged_rows[copy, lane, None].iterator) + high = smem_words(staged_rows[copy + 1, lane, None].iterator) + self.split_pairs(wr, tile, chunk, low, high) + acc = zeros((chains, tiles, 4), Float32) + for chunk in cutlass.range_constexpr(8): + chain: cutlass.Constexpr = chunk // 2 if mode == 1 else 0 + for tile in cutlass.range_constexpr(tiles): + wa = tuple(wr[tile, chunk, e] for e in range(4)) + mma_into(acc[chain, tile, None], wa, ar[chunk], bf16=True) + if cutlass.const_expr(self.tokens > 1): + gates = (a, lower_bound, dt, features) + targets = (shared_decay, barrier, value_base) + return self.multicast(acc, gates, targets, lane) + raw_g = [] + for r in cutlass.range_constexpr(per): + partial = cute.make_rmem_tensor(chains, Float32) + for chain in cutlass.range_constexpr(chains): + selected = acc[chain, 2 * r, 0] + if quarter & 1: + selected = acc[chain, 2 * r, 2] + if cutlass.const_expr(tiles > 2 * r + 1): # noqa: SIM102 + if quarter >> 1: + selected = acc[chain, 2 * r + 1, 0] + if quarter & 1: + selected = acc[chain, 2 * r + 1, 2] + partial[chain] = selected + value = partial[0] + if cutlass.const_expr(mode == 1): + value = Float32(0) + ( + ((value + partial[1]) + partial[2]) + partial[3] + ) + raw_g.append(Float32(BFloat16(value))) + + decay = [] + for r in cutlass.range_constexpr(per): + gate = cute.math.exp(a, fastmath=True) * (raw_g[r] + dt[r]) + gate = sigmoid(gate, self.fast_math) + decay.append(cute.math.exp(lower_bound * gate, fastmath=True)) + packets: cutlass.Constexpr = rows // 4 + packet, destination = lane % packets, lane // packets + values = cute.make_rmem_tensor(4, Float32) + for e in cutlass.range_constexpr(4): + source = 4 * packet + e + if cutlass.const_expr(per == 2): + feature = source & 31 + source = (feature & 7) << 2 | (feature >> 4) << 1 | feature >> 3 & 1 + low = cute.arch.shuffle_sync(decay[0], source) + high = cute.arch.shuffle_sync(decay[1], source) + values[e] = low + if packet >= 8: + values[e] = high + elif cutlass.const_expr(per > 2): # one CTA per head: 128 features + feature = source & 31 + source = (feature & 7) << 2 | (feature >> 4) << 1 | feature >> 3 & 1 + shuffled = [cute.arch.shuffle_sync(d, source) for d in decay] + values[e] = shuffled[0] + for r in cutlass.range_constexpr(1, per): + if packet >> 3 == r: + values[e] = shuffled[r] + else: + if cutlass.const_expr(mode != 2): + source = (source & 7) << 2 | (source >> 4) << 1 | source >> 3 & 1 + values[e] = cute.arch.shuffle_sync(decay[0], source) + offset = shared_decay.iterator + value_base + 4 * packet + pointer = cute.make_ptr( + Float32, offset.llvm_ptr, cute.AddressSpace.smem, assumed_align=16 + ) + words = tuple(float_bits(values[e]) for e in range(4)) + if cutlass.const_expr(self.early): + cute.arch.cluster_wait() + store_async_dsmem(pointer, words, barrier, destination) + + @cute.jit + def multicast(self, acc, gates, targets, lane): + """Decays (q > 1) of the lane's D fragments (feature rows, token columns 2 + quarter + {0, 1}) to every CTA's shared [token, 128] decays.""" + (a, lower_bound, dt, features), (shared_decay, barrier, value_base) = ( + gates, + targets, + ) + quarter = lane & 3 + if cutlass.const_expr(self.early): + cute.arch.cluster_wait() + for i in cutlass.range_constexpr(len(features)): + tile, e = i // 2, (i % 2) * 2 + for column in cutlass.range_constexpr(2): + value = acc[0, tile, e + column] + if cutlass.const_expr(self.gate_mode == 1): + parts = [acc[c, tile, e + column] for c in range(4)] + value = Float32(0) + (((value + parts[1]) + parts[2]) + parts[3]) + raw_g = Float32(BFloat16(value)) + gate = cute.math.exp(a, fastmath=True) * (raw_g + dt[i]) + gate = sigmoid(gate, self.fast_math) + decay = float_bits(cute.math.exp(lower_bound * gate, fastmath=True)) + token = 2 * quarter + column + offset = token * 128 + value_base + features[i] + pointer = shared_decay.iterator + offset + if token < self.tokens: + for rank in cutlass.range_constexpr(self.cluster): + store_async_dsmem(pointer, decay, barrier, rank) + + @staticmethod + def split_pairs(wr, tile, chunk, low, high): + """16 bytes of rows r and r + 8 hold the A fragments of chunks 2c and 2c + 1.""" + for h in range(2): + c = 2 * chunk + h + wr[tile, c, 0], wr[tile, c, 2] = low[2 * h : 2 * h + 2] + wr[tile, c, 1], wr[tile, c, 3] = high[2 * h : 2 * h + 2] + + @cute.jit + def setup(self, barrier, gate_barrier, partition, thread): + """Thread 0 arms the output (leader) and decay barriers; the CTA arrives on the + cluster barrier peers wait on before their sends.""" + tokens, cluster = self.tokens, self.cluster + owned = [(barrier, partition == 0, 512)] # token t's outputs: CTA t % cluster + if cutlass.const_expr(tokens > 1): + owned = [ + (barrier + k, partition + k * cluster < tokens, 512) + for k in range(-(-tokens // cluster)) + ] + for bar, cta, size in (*owned, (gate_barrier, True, 512 * tokens)): + if cta & (thread == 0): + cute.arch.mbarrier_init(bar, 1) + cute.arch.mbarrier_arrive_and_expect_tx(bar, size) + cute.arch.mbarrier_init_fence() + cute.arch.cluster_arrive_relaxed() + + @cute.jit + def row_dot(self, history, vector, p): + partial = cute.make_rmem_tensor(self.chunks, Float32) + for j in cutlass.range_constexpr(self.chunks): + row = (history[p, j, e] for e in range(4)) + partial[j] = dot(row, (vector[j, e] for e in range(4)), ftz=self.fast_math) + for shift in cutlass.range_constexpr(self.chunks.bit_length() - 2, -1, -1): + for j in cutlass.range_constexpr(1 << shift): + partial[j] += partial[j + (1 << shift)] + return partial[0] + + @cute.jit + def output_position(self, row: Int32): + return ((row & 7) << 4) | ((row >> 5) << 2) | ((row & 31) >> 3) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/prepared.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/prepared.py new file mode 100644 index 0000000000..4c43cf6834 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/kda/prepared.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Prepared KDA core. One packed projection supplies Q/K/V, the output gate, F_A and +beta; per-batch policies select the gate schedule and the projection row pitch.""" + +from functools import cache + +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import make_fake_stream + +from ..common.host import fake, layout, stream, to_cute +from ..common.policy import batch_table +from .kernel import KdaSplit + +# KdaSplit schedule "early.lanes.occupancy.cluster.evict" per batch by heads (TP8 12, +# TP16 6), from GB300 sweeps: early predecessor-independent loads (MMA gate modes only; +# 2 also stages the state rows in shared memory, 3 also the F_B operands), lanes per +# state row, minimum CTAs per SM, CTAs per head (0: defaults) and L2 evict-first state +# traffic (staged only). +SCHEDULE = {12: "16-32:3.16.3.2.1,64:3.16.4.2.1", 6: "16:1.0.0,32-64:3.16.3.2"} +# Verify (q > 1) schedules by request count B, from GB300 +# sweeps over q 4-8: no staging; eight-CTA clusters, then fewer as B grows. +SPEC_SCHEDULE = { + 12: "1:1.8.2.8,2-4:1.8.3.8,5-12:1.8.3.4,13-16:1.8.3.2", + 6: "1:1.8.2.8,2:1.8.0.8,3-12:1.8.3.8,13-16:1.8.3.4", +} + + +@cache +def executable(specs, heads, gate_mode, has_bias, device, tokens=1): + args = tuple(fake(shape, stride, dtype) for shape, stride, dtype in specs) + batch = specs[0][0][0] + table = SCHEDULE if tokens == 1 else SPEC_SCHEDULE + schedule = batch_table(table.get(heads, "")) + schedule = schedule.get(batch // tokens, (0, 0, 0)) if gate_mode != 2 else (0,) * 3 + if schedule[0] >= 2 and specs[9][1][-2:] != (128, 1): # staging needs whole rows + schedule = (1, *schedule[1:4]) # no evict hint without staging + requests = batch // tokens + kernel = KdaSplit(heads, requests, gate_mode, has_bias, schedule, tokens) + scalars = (cutlass.Float32(v) for v in (-5, 128**-0.5, 1e-6)) + with torch.accelerator.device_index(device): + return cute.compile( + kernel, *args[:13], *scalars, *args[13:], make_fake_stream() + ) + + +def decode_conv_weight(weight): + """The [3, width, D] (Q, K, V) decode layout of the [3 * D, 1, width] conv + weight.""" + width = weight.shape[-1] + return weight.detach().view(3, -1, width).transpose(1, 2).contiguous() + + +class PreparedKda: + def __init__(self, attention, policies, conv_dim_first): + bias, dt = attention.conv1d.bias, attention.dt_bias + self.heads = attention.A_log.numel() + self.dim, self.has_bias = self.heads * 128, bias is not None + conv_w = decode_conv_weight(attention.conv1d.weight) + bias, fb = dt if bias is None else bias, attention.f_b_proj.weight + norm = attention.o_norm.weight.detach().float() # BF16 slows the fused kernel + self.weights = (conv_w, bias, fb, attention.A_log, dt, norm) + self.device, self.policies = dt.device, policies + scalars = (attention.gate_lower_bound, 128**-0.5, attention.o_norm.eps) + self.scalars = tuple(map(cutlass.Float32, scalars)) + self.conv_dim_first, self.kernels = conv_dim_first, {} + + def _operands(self, caches, projected, slots, out, accepted=None): + """Kernel tensors in launch order: 13 before the scalars, then F_B (and the + verify step's accepted counts).""" + conv, state = caches[:2] + conv = conv if self.conv_dim_first else conv.transpose(-1, -2) + w, bias, fb, a_log, dt, norm = self.weights + d, h = self.dim, self.heads + p = lambda start, width: projected[:, start : start + width] + ops = (p(0, 3 * d), w, bias, conv, p(4 * d, 128), p(4 * d + 128, h), a_log, dt) + tail = (fb,) if accepted is None else (fb, accepted) + return ops + (slots, state, out, p(3 * d, d), norm) + tail + + def prepare_caches(self, caches): + """Policy keys: a batch, or (requests, tokens) for verify steps.""" + for key, (pitch, mode) in self.policies.items(): + requests, tokens = key if isinstance(key, tuple) else (key, 1) + batch, meta = requests * tokens, dict(device="meta") + projected = torch.empty((batch, pitch), dtype=torch.bfloat16, **meta) + out = torch.empty(batch, self.dim, dtype=torch.bfloat16, **meta) + accepted = torch.empty(requests, dtype=torch.int32, **meta) + shape = (requests, tokens) if tokens > 1 else (batch,) + slots = torch.empty(shape, dtype=torch.int32, **meta) + accepted = accepted if tokens > 1 else None + operands = self._operands(caches, projected, slots, out, accepted) + self.kernels[key] = executable( + tuple(map(layout, operands)), + *(self.heads, mode, self.has_bias, self.device.index, tokens), + ) + return self + + def run(self, projected, out, *, caches, slots, accepted=None): + """``accepted`` [>= B] (verify steps): ``slots`` is [>= B, q] contiguous.""" + batch = projected.shape[0] + key, rows = batch, batch + if accepted is not None: + rows = batch // slots.shape[1] + key, accepted = (rows, slots.shape[1]), accepted[:rows] + if slots.stride() != (slots.shape[1], 1): # e.g. a block-table slice + slots = slots.contiguous() + tensors = self._operands(caches, projected, slots[:rows], out, accepted) + views = tuple(to_cute(t, 16) for t in tensors) + self.kernels[key](*views[:13], *self.scalars, *views[13:], stream(self.device)) + return out diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/__init__.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/dcp.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/dcp.py new file mode 100644 index 0000000000..4cd94b4261 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/dcp.py @@ -0,0 +1,378 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MLA decode under decode context parallelism (DCP = TP, rank d owns head block d) +without collectives, three kernels per layer. Each rank's KV cache holds 1 / DCP of +every context. + + MlaQuery FP8 query rows [B, H, 576] (+ the KV insert where this rank owns the + slot); each CTA multicasts its bytes into every rank's Lamport q + mailbox [B, DCP*H, 576], then polls its share until all landed + tokenspeed all DCP*H heads over the local KV shard -> BF16 rows + LSE (log2) + DcpValue re-arms the q mailbox; rows pushed to their head owner's Lamport + mailbox (peer stores); the owner LSE-merges its H heads over the DCP + sources, then W_UV and the gate (padding from the global lengths) +""" + +import math +from functools import cache + +import cutlass +import cutlass.cute as cute +import torch +import torch.distributed._symmetric_memory as symm_mem +from cutlass import BFloat16, Float32, Int32, Int64, Uint32 +from cutlass.cute.runtime import make_fake_stream + +from ..common.host import fake +from ..common.primitives import ( + alloc, + at, + bf16x2, + bits_float, + float_bits, + launch, + load_volatile, + load_words, + mma_into, + no_neg_zero, + pack_pair, + published, + sigmoid, + store_relaxed, + store_shared_words, + store_words, + unpack_bf16, + warp_sum, + zeros, +) + +LSE_SENT = 0xFFFFFFFF # a NaN pattern the merge never sends +# The LSE of an empty or invalid source: finite, so the merge never forms inf - inf. +EMPTY_LSE = -3.0e38 +Q_SENT = 0xFFFFFFFF # four FP8 0xFF bytes: never written (satfinite, canonical NaN) +PAIR_SENT = 0x80000000 # BF16 pair (+0, -0); sent pairs have no negative zero + + +class DcpBox: + """Model-owned DCP state over ``group``, Lamport mailboxes (sentinel-filled, + re-armed by their reader): q [2, CAP, DCP*H, 576] FP8 bytes (multicast stores), + partials [2, DCP, CAP, H, 512] BF16 and LSE [2, DCP, CAP, H] FP32 (peer stores), + and the local tokenspeed rows / LSE. Parity alternates by layer. ``interleave``: + the KV cache's DCP interleave (global masking of spec verify rows).""" + + def __init__(self, group, device, size, rank, heads, cap=64, interleave=1): + self.size, self.rank, self.heads, self.cap = size, rank, heads, cap + self.interleave = interleave + empty = symm_mem.empty + self.q = empty((2, cap, size * heads, 576), dtype=torch.uint8, device=device) + shape = (2, size, cap, heads) + self.partial = empty((*shape, 512), dtype=torch.bfloat16, device=device) + self.lse = empty(shape, dtype=torch.float32, device=device) + tensors = (self.q, self.partial, self.lse) + handles = [symm_mem.rendezvous(t, group) for t in tensors] + self.q_mc = int(handles[0].multicast_ptr or 0) + if not self.q_mc: + raise RuntimeError("Kimi K3 specialized MLA DCP needs multicast memory") + i64 = dict(dtype=torch.int64, device=device) + self.peers = torch.tensor(handles[1].buffer_ptrs, **i64) + self.lse_peers = torch.tensor(handles[2].buffer_ptrs, **i64) + self.q.fill_(0xFF) + self.partial.view(torch.int32).fill_(-0x80000000) + self.lse.view(torch.int32).fill_(-1) + rows = (cap, 1, size * heads) + self.rows = torch.empty((*rows, 512), dtype=torch.bfloat16, device=device) + self.rows_lse = torch.empty(rows, dtype=torch.float32, device=device) + torch.accelerator.synchronize(device) + torch.distributed.barrier(group=group) + + def mc(self, parity): + """The q mailbox multicast address of ``parity``.""" + return Int64(self.q_mc + parity * self.q[0].numel()) + + +class DcpValue: + """The DCP tail of one MLA layer, one 512-thread CTA per (head, 16 requests). + Every rank re-arms this parity's q mailbox (tokenspeed has read it) and pushes all + its local tokenspeed rows (+ LSE) to their head owner's Lamport mailbox (empty + local shards: zeros, LSE -inf). Warp r then LSE-weights request t0 + r of this + head over the DCP sources (exp2(lse - max) / sum in FP32), the CTA merges the + polled rows in FP32 + (one BF16 rounding, re-arming what it read) into shared memory, and W_UV (BF16 + MMA, FP32 accumulation) and the sigmoid gate follow as in MlaValue: warp w owns + outputs 32 (w // 4) + 8 (w % 4) .. + 7. Rows are tokens, ``q`` per request.""" + + def __init__(self, heads, size, batch, sms=148, q=1): + self.heads, self.size, self.batch = heads, size, batch + self.q = q + # Push-only CTAs (z > 0) spread the NVLink stores over more SMs: up to one + # wave, at least 4 16 B chunks per thread. + merge = heads * ((batch + 15) // 16) + chunks = batch * size * heads * 64 + pushers = min(max(1, sms // merge), -(-chunks // (merge * 512 * 4))) + self.grid = (heads, (batch + 15) // 16, pushers) + + @cute.jit + def __call__(self, *args): + launch(self.kernel, args, self.grid, 512) + + @cute.kernel + def kernel(self, args): + rows, rows_lse, lens, box, lbox, peers, lpeers, mailbox = args[0:8] + w_uv, gate, totals, output, parity, rank = args[8:14] + H, D, B, cap = self.heads, self.size, self.batch, box.shape[1] + smem = cutlass.utils.SmemAllocator() + latent = alloc(smem, BFloat16, (16, 544), 128) + scales = alloc(smem, Float32, (16, D)) + tid, _, _ = cute.arch.thread_idx() + head, tile, pusher = cute.arch.block_idx() + lane, warp = tid % 32, tid // 32 + group, quarter = lane // 4, lane % 4 + v = 32 * (warp // 4) + 8 * (warp % 4) + group + # B fragments: 8 latent values per 32-wide chunk; A uses the same permutation. + wb = [ + load_words(at(w_uv, head, 32 * c + 8 * quarter, v), count=4) + for c in range(16) + ] + cute.arch.griddepcontrol_wait() + threads = H * self.grid[1] * self.grid[2] * 512 + me = ((pusher * self.grid[1] + tile) * H + head) * 512 + tid + for c in range(me, B * D * H * 36, threads): # 16 B q mailbox chunks + store_words(mailbox.iterator + c * 16, (Uint32(Q_SENT),) * 4) + for i in range(me, B * D * H * 64, threads): + r, k = i // 64, i % 64 # r = b * D * H + global head + b, gh = r // (D * H), r % (D * H) + words = (Uint32(0),) * 4 + if lens[b // self.q] > 0: + words = load_words(at(rows, b, 0, gh, 8 * k), count=4, mutable=True) + words = tuple(no_neg_zero(w) for w in words) + slot = ((parity * D + rank) * cap + b) * H + gh % H + store_relaxed(peers[gh // H] + (Int64(slot) * 512 + 8 * k) * 2, words) + for i in range(me, B * D * H, threads): + b, gh = i // (D * H), i % (D * H) + value, n = rows_lse[b, 0, gh], lens[b // self.q] + bad = (value != value) | (value <= Float32(EMPTY_LSE)) # NaN, -inf + if (n <= 0) | bad | (value >= Float32(-EMPTY_LSE)): # or +inf + value = Float32(EMPTY_LSE) + slot = ((parity * D + rank) * cap + b) * H + gh % H + store_relaxed(lpeers[gh // H] + Int64(slot) * 4, (float_bits(value),)) + cute.arch.griddepcontrol_launch_dependents() + if pusher == 0: + self.merge(args, latent, scales, wb, tid) + + @cute.jit + def merge(self, args, latent, scales, wb, tid): + """The owner's merge of request rows t0.. of this head, then W_UV + gate.""" + box, lbox, gate, totals, output = args[3], args[4], args[9], args[10], args[11] + D, B = self.size, self.batch + head, tile, _ = cute.arch.block_idx() + lane, warp = tid % 32, tid // 32 + group, quarter = lane // 4, lane % 4 + t0 = tile * 16 + col = head * 128 + 32 * (warp // 4) + 8 * (warp % 4) + 2 * quarter + gates = [] + for r in cutlass.range_constexpr(2): + word = Uint32(0) + if t0 + group + 8 * r < B: + word = load_words( + at(gate, t0 + group + 8 * r, col), count=1, mutable=True + )[0] + gates.append(unpack_bf16(word)) + if t0 + warp < B: # warp r: the source weights of request t0 + r + lse = Float32(EMPTY_LSE) + if lane < D: + source = at(lbox, lane, t0 + warp, head) + word = load_volatile(source)[0] + while word == Uint32(LSE_SENT): + word = load_volatile(source)[0] + store_words(source, (Uint32(LSE_SENT),)) + lse = bits_float(word) + top = lse + for shift in cutlass.range_constexpr(4, -1, -1): + top = max(top, cute.arch.shuffle_sync_bfly(top, 1 << shift)) + weight = Float32(0) + if lse > Float32(EMPTY_LSE): + weight = cute.math.exp2(lse - top, fastmath=True) + total = warp_sum(weight) + scale = Float32(0) + if total > 0: + scale = weight / total + if lane < D: + scales[warp, lane] = scale + cute.arch.sync_threads() + chunk = tid % 64 + for half in cutlass.range_constexpr(2): # rows r and r + 8, 16 B chunk + r = tid // 64 + 8 * half + merged = zeros(8, Float32) + if t0 + r < B: + for s0 in cutlass.range_constexpr(0, D, 8): # 8 polls in flight + sources = [ + at(box, s, t0 + r, head, 8 * chunk) + for s in range(s0, min(D, s0 + 8)) + ] + loads = [load_volatile(source, count=4) for source in sources] + for i in cutlass.range_constexpr(len(sources)): + words = loads[i] + while not published(words): + words = load_volatile(sources[i], count=4) + store_words(sources[i], (Uint32(PAIR_SENT),) * 4) + w = scales[r, s0 + i] + for e in cutlass.range_constexpr(4): + lo, hi = unpack_bf16(words[e]) + merged[2 * e] += w * lo + merged[2 * e + 1] += w * hi + pairs = tuple( + pack_pair(merged[2 * e], merged[2 * e + 1], policy=None) + for e in range(4) + ) + store_shared_words(latent.iterator + r * 544 + 8 * chunk, pairs) + cute.arch.sync_threads() + parts = zeros((2, 4), Float32) # two accumulation chains + for c in cutlass.range_constexpr(16): + x, y = ( + load_words( + latent.iterator + r * 544 + 32 * c + 8 * quarter, + count=4, + shared=True, + ) + for r in (group, group + 8) + ) + for s in cutlass.range_constexpr(2): + a = (x[2 * s], y[2 * s], x[2 * s + 1], y[2 * s + 1]) + b = (wb[c][2 * s], wb[c][2 * s + 1]) + mma_into(parts[s, None], a, b, bf16=True) + acc = [parts[0, k] + parts[1, k] for k in range(4)] + for r in cutlass.range_constexpr(2): + t = t0 + group + 8 * r + if t < B: + # Padded rows may hold a NaN gate; keep inactive rows zero. + lo, hi = Float32(0), Float32(0) + if totals[t // self.q] > 0: + lo, hi = ( + Float32(BFloat16(acc[2 * r + e])) * sigmoid(gates[r][e], False) + for e in range(2) + ) + store_words(at(output, t, col), (bf16x2(lo, hi),)) + + +@cache +def compile_value(heads, size, batch, cap, uv, pitch, device, sms=148, q=1): + """DcpValue of one batch (``q`` rows per request); ``uv``: the (shape, stride, + dtype) of W_UV, ``pitch``: the projected rows' (the gate follows the 2112 query / + KV columns).""" + f32, i32 = torch.float32, torch.int32 + w_uv, gate = fake(*uv), fake((batch, heads * 128), (pitch, 1)) + rows = fake((batch, 1, size * heads, 512)) + rows_lse = fake((batch, 1, size * heads), None, f32, 4) + lens = fake((batch // q,), None, i32, 4) + box = fake((size, cap, heads, 512)) + lbox = fake((size, cap, heads), None, f32, 4) + peers = fake((size,), None, torch.int64, 8) + mailbox = fake((cap, size * heads, 576), None, torch.uint8) + output = fake((batch, heads * 128), (heads * 128, 1)) + args: tuple = (rows, rows_lse, lens, box, lbox, peers, peers, mailbox, w_uv) + args += (gate, lens, output, Int32(0), Int32(0), make_fake_stream()) + with torch.accelerator.device_index(device): + return cute.compile(DcpValue(heads, size, batch, sms, q), *args) + + +_WORKSPACE: dict[tuple[torch.device, int], torch.Tensor] = {} + + +def stock_decode( + query, + cache, + block_table, + lengths, + max_seq, + scales, + out, + lse=None, + min_split=1, + pdl=True, + cp=(1, 0, 1), + totals=None, + pack=False, +): + """tokenspeed_mla_decode's FP8 path with caller-owned ``out`` and ``lse`` (log2) + buffers and no per-call allocation: the split count follows from (B, max_seq, + min_split); ``lengths`` bound the work, empty splits are skipped by its reducer. + The q rows of a request mask causally; under DCP in global positions (``cp``: + world, rank, interleave; ``totals``: the global lengths). ``pack``: M128 tiles hold + consecutive (query row, head) rows across the q rows.""" + import tvm_ffi + from tokenspeed_mla import mla_decode as md + + batch, q_len, heads, _ = query.shape + fp8, device = torch.float8_e4m3fn, query.device + arch = torch.cuda.get_device_capability(device) + tiler, _ = md.select_mla_decode_tilers( + heads, q_len, is_fp8=True, compute_capability=arch + ) + pack = pack and tiler[0] == 128 + fold = 1 if pack else md.get_mla_decode_fold_sq_factor(heads, q_len, tiler[0]) + rows = (q_len // fold, heads * fold) + if pack: + rows = (md.compute_q_tile_layout(heads, q_len, 128)[1], 128) + sms = md.get_num_sm(device) + split, size = md._get_split_kv_and_workspace_size( + batch, *rows, 512, sms, max_seq, fp8, tiler, min_split + ) + # One split-KV workspace per (device, heads), shared by every MLA layer: the calls + # are stream ordered. PreparedMla.prepare_cache runs each prepared (batch, q) before + # any capture, so the buffer reaches its final size first; growing it during a + # capture would leave the graphs captured before with the freed buffer. + key = device, heads + if key not in _WORKSPACE or _WORKSPACE[key].numel() < size: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "the MLA split-KV workspace must be sized before capture" + ) + bound = max(size, sms * heads * 513 * 4) + _WORKSPACE[key] = torch.empty(bound, dtype=torch.int8, device=device) + kernel = md._get_compiled_mla_kernel( + torch_dtype=fp8, + page_size=cache.shape[1], + kv_lora_rank=512, + qk_rope_head_dim=64, + is_persistent=False, + is_var_seq=True, + is_var_split_kv=False, + compute_capability=arch, + has_local_visible_lens=False, + skip_correction_threshold=8.0, + is_workspace_size_zero=size == 0, + fold_sq_factor=fold, + causal_mask=True, + num_heads=heads, + seq_len_q=q_len, + window_left=-1, + cp_world=cp[0], + cp_rank=cp[1], + cp_interleave_size=cp[2], + use_pdl=pdl, + return_lse=lse is not None, + reducer_d_tiles=md._get_reducer_d_tiles(batch, q_len, heads, sms, split), + reducer_max_splits=md._get_reducer_max_splits(split), + pack_q=pack, + partial_fp16=False, + ) + cache = cache.view(fp8) + softmax, output_scale = scales + args = [query[..., :512], query[..., 512:], cache[..., :512], cache[..., 512:]] + args += [block_table, out, lse, _WORKSPACE[key][:size] if size else None] + args += [Int32(split), lengths, lengths if totals is None else totals, None] + args += [Float32(softmax), Float32(output_scale)] + with tvm_ffi.use_torch_stream(): + kernel(*args) + + +def fill_splits(batch, device): + """Splits filling the SMs (one CTA pair each); tokenspeed caps itself at 32.""" + from tokenspeed_mla import get_num_sm + + return min(256, max(1, get_num_sm(device) // (2 * batch))) + + +def local_max(max_context, size): + """Longest local shard of a ``max_context`` sequence (block-interleaved).""" + return math.ceil(max_context / (size * 64)) * 64 diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/prepared.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/prepared.py new file mode 100644 index 0000000000..758f21d28f --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/prepared.py @@ -0,0 +1,255 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Prepared MLA core: the packed input projection feeds the query kernel, then either +tokenspeed_mla_decode (core "tokenspeed") or FlashInfer's trtllm-gen MLA decode (core +"trtllm"; FP8 Q and KV as tokenspeed, FP32 softmax and accumulation) into the BF16 +latent, and MlaValue (W_UV / gate). Under DCP (``dcp``: a DcpBox) the tokenspeed core +runs over the local KV shard for every DCP rank's heads between the query's q +publication and ``DcpValue``'s merge (see mla.dcp). Spec verify: ``q`` query tokens +per request (rows grouped by request, causal among them).""" + +from functools import cache + +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import make_fake_stream + +from ..common.host import fake, layout, prepare_shared_carveout, stream, to_cute +from .dcp import ( + compile_value, + fill_splits, + local_max, + stock_decode, +) +from .query import MlaQuery +from .value import MlaValue + +# tokenspeed (SM-filling splits) streams long contexts fast; trtllm-gen wins from B16. +TRTLLM_CORE_BATCH, CORES = 16, ("tokenspeed", "trtllm") +_FLASHINFER: dict[tuple[torch.device, int], tuple[torch.Tensor, torch.Tensor]] = {} + + +@cache +def compile_core(batch, heads, pitch, specs, device, dcp=0, sms=148, q=1): + """CuTe kernels (query, value) of one batch of rows (``q`` per request); under DCP + (``dcp``: its size) the query publishes and the value is DcpValue.""" + qnorm, kvnorm, qb, uk, uv, qinv, kinv, kv = ( + fake(shape, stride, dtype, 4 if dtype == torch.float32 else 16) + for shape, stride, dtype in specs + ) + query = fake((batch, heads, 576), (heads * 576, 576, 1), torch.uint8) + projected = fake((batch, 2112), (pitch, 1)) + slots = fake((batch,), (1,), torch.int64, 8) + lengths = fake((batch // q,), (1,), torch.int32, 4) + gate = fake((batch, heads * 128), (pitch, 1)) + output = fake((batch, heads * 128), (heads * 128, 1)) + latent = fake((batch, heads, 512)) + f32, s = cutlass.Float32, make_fake_stream() + cp_async = heads == 12 and batch == 1 + mla_query = MlaQuery(heads, batch, q_b_cp_async=cp_async, dcp=dcp) + q_args: tuple = (projected, qnorm, kvnorm, qb, uk, kv, slots, qinv, kinv) + q_args += (f32(1e-6), query) + if dcp: # this parity's q mailbox, its multicast address, rank + mailbox = fake((64, dcp * heads, 576), None, torch.uint8) + q_args += (mailbox, cutlass.Int64(0), cutlass.Int32(0)) + with torch.accelerator.device_index(device): + kernels = [cute.compile(mla_query, *q_args, s)] + if batch > 1: + kernels = [prepare_shared_carveout(kernels[0], device, 100)] + if dcp: + uv_spec = specs[4] + kernels.append( + compile_value(heads, dcp, batch, 64, uv_spec, pitch, device, sms, q) + ) + else: + v_args = (latent, uv, gate, lengths, output, s) + value = MlaValue(heads, batch, q) + kernels.append(cute.compile(value, *v_args)) + return tuple(kernels) + + +def core_for(batch, core="auto", q=1): + """``core`` unless "auto": trtllm-gen from B16, else tokenspeed (spec verify, + q > 1: tokenspeed, whose SM-filling splits stream long contexts far faster).""" + if core == "auto": + core = "trtllm" if batch >= TRTLLM_CORE_BATCH and q == 1 else "tokenspeed" + if core not in CORES: + raise ValueError( + f"Kimi K3 MLA core {core!r}: expected auto, tokenspeed or trtllm" + ) + return core + + +def pad_requests(batch, block_table, lengths, totals=None): + """Request metadata for ``batch`` rows: PIECEWISE graphs pad the tokens but not + the requests, so missing rows become empty requests (length 0 on page 0).""" + rows = min(lengths.shape[0], block_table.shape[0]) + if rows >= batch: + return block_table, lengths, totals + pad = torch.nn.functional.pad + block_table = pad(block_table[:rows], (0, 0, 0, batch - rows)) + lengths = pad(lengths[:rows], (0, batch - rows)) + totals = None if totals is None else pad(totals[:rows], (0, batch - rows)) + return block_table, lengths, totals + + +def trtllm_decode(query, cache, block_table, lengths, max_seq, scales, out): + """FlashInfer trtllm-gen MLA decode of FP8 ``query`` [B, q, H, 576] (causal among + the q rows) over the FP8 paged cache into BF16 ``out`` [B, q, H, 512], PDL chained; + zeroed workspace and multi-CTA counter (self-resetting) owned here, one per (device, + heads).""" + from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla + from flashinfer.utils import get_trtllm_gen_multi_ctas_kv_counter_bytes + + device, heads = query.device, query.shape[2] + if (device, heads) not in _FLASHINFER: + sms = torch.cuda.get_device_properties(device).multi_processor_count + count = get_trtllm_gen_multi_ctas_kv_counter_bytes(64, heads, sms) + workspace = torch.zeros(128 << 20, dtype=torch.int8, device=device) + counter = torch.zeros(count, dtype=torch.uint8, device=device) + _FLASHINFER[device, heads] = workspace, counter + workspace, counter = _FLASHINFER[device, heads] + softmax, output_scale = scales + trtllm_batch_decode_with_kv_cache_mla( + query=query, + kv_cache=cache.view(torch.float8_e4m3fn).unsqueeze(1), + workspace_buffer=workspace, + qk_nope_head_dim=128, + kv_lora_rank=512, + qk_rope_head_dim=64, + block_tables=block_table, + seq_lens=lengths, + max_seq_len=max_seq, + out=out, + bmm1_scale=softmax, + bmm2_scale=output_scale, + enable_pdl=True, + backend="trtllm-gen", + multi_ctas_kv_counter_buffer=counter, + ) + + +class PreparedMla: + """``max_context`` (the deployment's longest sequence) bounds the split count; + ``core`` forces "tokenspeed" or "trtllm". ``dcp``: the model's DcpBox, with this + layer's mailbox ``parity``. ``queries``: the query lengths q (spec verify: 1 + the + draft length) prepared for every batch of rows divisible by q.""" + + def __init__( + self, + attention, + pitches, + query_scratch, + *, + max_context=None, + core="auto", + dcp=None, + parity=0, + queries=(1,), + ): + self.heads = attention.W_UK_T.shape[0] + norms = attention.q_a_layernorm.weight, attention.kv_a_layernorm.weight + self.weights: tuple = (*norms, attention.q_b_proj.weight, attention.W_UK_T) + self.weights += (attention.W_UV, attention._q_scale_inv, attention._k_scale_inv) + self.device, self.pitches = attention.q_b_proj.weight.device, pitches + k_scale = attention._k_scale_float + softmax = attention.scale * attention._q_scale_float * k_scale + self.eps = cutlass.Float32(attention.rms_norm_eps) + self.stock_scales = softmax, k_scale + self.views = tuple( + to_cute(w, 4 if w.dtype == torch.float32 else 16) for w in self.weights + ) + self.query_scratch, self.kernels = query_scratch, {} + shapes = [(b, q) for b in pitches for q in queries if b % q == 0] + self.cores = {(b, q): core_for(b, core, q) for b, q in shapes} + # The split count only depends on this bound; seq_lens bound the work. + self.max_seq = max_context or 1 << 17 + self.dcp, self.parity = dcp, parity + if dcp is not None: + self.cores = dict.fromkeys(shapes, "tokenspeed") + self.max_seq = local_max(self.max_seq, dcp.size) + shape, bf16 = (max(pitches, default=0), self.heads, 512), torch.bfloat16 + self.latent = torch.empty(shape, dtype=bf16, device=self.device) + + def prepare_cache(self, kv): + specs = tuple(layout(w) for w in (*self.weights, kv.view(torch.uint8))) + for batch, q in self.cores: + device = self.device.index + dcp = self.dcp.size if self.dcp is not None else 0 + sms = torch.cuda.get_device_properties(self.device).multi_processor_count + pitch = self.pitches[batch] + args = batch, self.heads, pitch, specs, device, dcp, sms, q + query = to_cute(self.query_scratch[:batch]) + latent = to_cute(self.latent[:batch]) + self.kernels[batch, q] = (*compile_core(*args), query, latent) + if self.dcp is not None: + self.prepare_dcp(batch) + if kv.is_cuda: # compile before graph capture + table = torch.zeros(batch, 1, dtype=torch.int32, device=kv.device) + ones = torch.full((batch,), q, dtype=torch.int32, device=kv.device) + self.stock_core(kv, table, ones, batch, q, ones) + return self + + def prepare_dcp(self, batch): + """(MlaQuery's publish arguments, DcpValue's mailbox arguments, parity and + rank) of one batch.""" + box, p = self.dcp, self.parity + mailbox, ids = to_cute(box.q[p]), (cutlass.Int32(p), cutlass.Int32(box.rank)) + rows = to_cute(box.rows[:batch]), to_cute(box.rows_lse[:batch], 4) + boxes = to_cute(box.partial[p]), to_cute(box.lse[p], 4) + peers = to_cute(box.peers, 8), to_cute(box.lse_peers, 8) + self.dcp_args = getattr(self, "dcp_args", {}) + self.dcp_args[batch] = (mailbox, box.mc(p), ids[1]), (rows, boxes, peers), ids + + def stock_core(self, cache, block_table, lengths, batch, q=1, totals=None): + """tokenspeed_mla_decode or trtllm-gen into the BF16 latent (DCP: tokenspeed + for every rank's heads over the local shard into the DcpBox rows + LSE; q > 1 + masks causally in global positions, ``totals``). tokenspeed: a fixed split + count per batch (one CTA pair per split filling the SMs); both PDL chained + between our kernels.""" + box, reqs = self.dcp, batch // q + query = self.query_scratch[:batch].view(torch.float8_e4m3fn) + out, lse, cp = self.latent[:batch], None, {} + if box is not None: + query, out = box.q[self.parity, :batch].view(query.dtype), box.rows[:batch] + lse = box.rows_lse[:batch].view(reqs, q, -1) + if q > 1: # D*H heads: packed (row, head) tiles stream the KV once per tile + cp = dict(cp=(box.size, box.rank, box.interleave), totals=totals[:reqs]) + cp.update(pack=True) + query, out = query.view(reqs, q, *query.shape[1:]), out.view(reqs, q, -1, 512) + args = cache, block_table[:reqs], lengths[:reqs], self.max_seq + if self.cores[batch, q] == "trtllm": + trtllm_decode(query, *args, self.stock_scales, out) + return + split = fill_splits(batch, self.device) + stock_decode(query, *args, self.stock_scales, out, lse, split, **cp) + + def run( + self, projected, out, *, cache, slots, block_table, lengths, totals=None, q=1 + ): + """``lengths``/``block_table``: this rank's KV (the local shard under DCP); + ``totals``: under DCP, the requests' full lengths. ``q`` rows per request (spec + verify; ``slots`` per row). Rows of ``projected`` beyond the requests run as + empty requests (zero output).""" + batch, s = projected.shape[0], stream(self.device) + reqs = batch // q + block_table, lengths, totals = pad_requests(reqs, block_table, lengths, totals) + kv = to_cute(cache.view(torch.uint8)) + mq, v, query, latent = self.kernels[batch, q] + qnorm, kvnorm, qb, uk, uv, qinv, kinv = self.views + q_slots, x = to_cute(slots[:batch], 8), to_cute(projected[:, :2112]) + q_args = (x, qnorm, kvnorm, qb, uk, kv, q_slots, qinv, kinv, self.eps, query) + out_view, gate = to_cute(out), to_cute(projected[:, 2112:]) + lens = to_cute(lengths[:reqs], 4) + if self.dcp is not None: # DcpValue zeroes padded requests (global lengths) + publish, ((rows, rows_lse), boxes, peers), ids = self.dcp_args[batch] + mq(*q_args, *publish, s) + self.stock_core(cache, block_table, lengths, batch, q, totals) + v_args = (rows, rows_lse, lens, *boxes, *peers, publish[0], uv, gate) + v(*v_args, to_cute(totals[:reqs], 4), out_view, *ids, s) + return out + mq(*q_args, s) + self.stock_core(cache, block_table, lengths, batch, q) + v(latent, uv, gate, lens, out_view, s) + return out diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/query.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/query.py new file mode 100644 index 0000000000..45fff905b6 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/query.py @@ -0,0 +1,474 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Cluster MLA query RMSNorm, Q_B projection, W_UK absorption and paged FP8 KV insert. +Static weights load before the PDL wait; current-token data is consumed after it.""" + +import cutlass +import cutlass.cute as cute +from cutlass import BFloat16, Float32, Int64 + +from ..common.primitives import ( + alloc, + at, + bf16x2, + copy_evict, + copy_g2s, + copy_s2cluster, + float_bits, + fma, + fma2, + fp8_bits, + launch, + ldmatrix, + load_bf16x8, + load_volatile, + load_words, + mma_into, + mul2, + multicast_words, + prefetch_l2, + published, + query_norm_rsqrt, + smem_words, + store_fp8_pair, + store_shared_words, + unpack_bf16, + view, + wait, + warp_sum, + zeros, +) + +KP = 1536 + 32 # padded Q_B / x row (elements): 3136 B, conflict-free 16 B fragments +# The smallest batch on the tiled kernel (per head count): +# its 16-request tiles beat FP32 FMAs from B4 at 12 heads and from B6 at 6 heads. +TILED_FROM = {12: 4, 6: 6} +Q_SENT = 0xFFFFFFFF # q mailbox sentinel: four FP8 0xFF bytes, never written + + +class MlaQuery: + """A cluster splits one head's 192 Q_B rows; each CTA then absorbs its share of the + 512 latent columns. Small batches (below TILED_FROM) are served by one CTA per + cluster part with FP32 FMAs; wider batches tile 16-64 requests per CTA (one wave + of 96 CTAs at B64 for both head counts) and use BF16 MMAs (FP32 accumulation). + + ``dcp`` (decode context parallel size; args then end with this parity's local q + mailbox, its multicast address and the rank): each CTA also multicasts the query + bytes it wrote into head block ``rank`` of every rank's Lamport q mailbox + (mla.dcp.DcpBox), then polls its share of the gathered mailbox until no word is + the sentinel, so the next kernel reads the gathered query.""" + + def __init__( + self, + heads: int, + batch: int, + *, + q_b_cp_async: bool = False, + dcp: int = 0, + ): + tiled = TILED_FROM[heads] + self.heads, self.batch, self.wide = heads, batch, batch >= tiled + self.dcp = dcp + self.cluster = 16 if heads == 6 and not self.wide else 8 + self.rows, self.latent = 192 // self.cluster, 512 // self.cluster + self.groups, self.p_per_group = 256 // self.latent, self.latent // 2 + self.tokens = 1 if batch == 2 else min(batch, 8) + if self.wide: + self.tokens = 16 if batch <= 16 else 32 if batch <= 32 or heads < 8 else 64 + self.grid = (self.cluster, heads, (batch + self.tokens - 1) // self.tokens) + carveout = 100 if batch > 1 else None + self.options = dict(cluster=self.cluster, preferred_smem_carveout=carveout) + self.q_b_cp_async = q_b_cp_async + + @cute.jit + def __call__(self, *args): + kernel = self.tiled if self.wide else self.kernel + launch(kernel, args, self.grid, 256, **self.options) + + @cute.jit + def live(self, token): + return cutlass.const_expr(not self.wide) or token < self.batch + + @cute.jit + def publish(self, args, tok0, tid): + """DCP: this CTA's query bytes (latent columns L*part.. as 16 B, the rope rows + of Q_B rows R*part.. as 4 B multicast stores) -> every rank's q mailbox (FP8 + bytes are never 0xFF: no word is the sentinel).""" + output, (box_mc, rank) = args[10], args[12:14] + part, head, _ = cute.arch.block_idx() + R, L, D = self.rows, self.latent, self.dcp + chunks: cutlass.Constexpr = L // 16 + items: cutlass.Constexpr = chunks + R // 4 + cute.arch.sync_threads() + for i in range(tid, self.tokens * items, 256): + t, w = i // items, i % items + p = R * part + 4 * (w - chunks) # Q_B row of a rope word + row = ((tok0 + t) * D + rank) * self.heads + head + if w < chunks: + byte = L * part + 16 * w + if self.live(tok0 + t): + words = load_words( + at(output, tok0 + t, head, byte), count=4, mutable=True + ) + multicast_words(box_mc + Int64(row) * 576 + byte, words) + elif self.live(tok0 + t) & (p >= 128): + byte = 512 + p - 128 + word = load_words( + at(output, tok0 + t, head, byte), count=1, mutable=True + ) + multicast_words(box_mc + Int64(row) * 576 + byte, word, kind="b32") + + @cute.jit + def gathered(self, args, tid): + """DCP: poll this CTA's share of the gathered mailbox (16 B chunks of the + B x DCP*H x 576 bytes) until every rank's bytes have landed.""" + mailbox = args[11] + part, head, tile = cute.arch.block_idx() + threads: cutlass.Constexpr = self.grid[0] * self.grid[1] * self.grid[2] * 256 + chunks: cutlass.Constexpr = self.batch * self.dcp * self.heads * 36 + first = ((tile * self.heads + head) * self.cluster + part) * 256 + tid + last = chunks - 1 # threads past the end re-poll it + for i0 in cutlass.range_constexpr(0, -(-chunks // threads), 8): + n: cutlass.Constexpr = min(8, -(-chunks // threads) - i0) + sources = [ + mailbox.iterator + min(first + (i0 + i) * threads, last) * 16 + for i in range(n) + ] + loads = [load_volatile(source, count=4) for source in sources] + for i in cutlass.range_constexpr(n): # the loads above are in flight + words = loads[i] + while not published(words, sentinel=Q_SENT): + words = load_volatile(sources[i], count=4) + + @cute.kernel + def kernel(self, args): + qkv, w_qa, w_kva, w_qb, w_uk, cache, slots, q_scale_inv = args[0:8] + k_scale_inv, eps, output = args[8:11] + tid, _, _ = cute.arch.thread_idx() + part, head, tok0 = cute.arch.block_idx() + lane, warp = tid % 32, tid // 32 + B, R, L, G = self.tokens, self.rows, self.latent, self.groups + smem = cutlass.utils.SmemAllocator() + qb = alloc(smem, BFloat16, (R, 1536), 128) + uk = alloc(smem, BFloat16, (128, L)) + qa = alloc(smem, BFloat16, 1536) + xq = alloc(smem, BFloat16, (B, 1536)) + qrow = alloc(smem, Float32, (R, B)) + qn = alloc(smem, Float32, (B, 128)) + red = alloc(smem, Float32, (G, B, L)) + ss = alloc(smem, Float32, (8, max(B, 2))) + if cutlass.const_expr(not self.q_b_cp_async): + barw = smem.allocate(cutlass.Uint64, 8) + barq = smem.allocate(cutlass.Uint64, 8) + if tid == 0: + if cutlass.const_expr(not self.q_b_cp_async): + cute.arch.mbarrier_init(barw, 1) + cute.arch.mbarrier_init(barq, 1) + cute.arch.mbarrier_init_fence() + if cutlass.const_expr(not self.q_b_cp_async): + cute.arch.mbarrier_arrive_and_expect_tx(barw, R * 1536 * 2) + copy_evict( + qb.iterator, + w_qb.iterator + (head * 192 + part * R) * 1536, + size=R * 1536 * 2, + barrier=barw, + ) + cute.arch.mbarrier_arrive_and_expect_tx(barq, B * 128 * 4) + cute.arch.cluster_arrive_relaxed() + for c in range(tid, 128 * L // 8, 256): + p, j = c // (L // 8), c % (L // 8) + copy_evict(uk.iterator + p * L + j * 8, at(w_uk, head, p, part * L + j * 8)) + if tid < 192: + copy_evict(qa.iterator + tid * 8, w_qa.iterator + tid * 8) + cute.arch.cp_async_commit_group() + if cutlass.const_expr(self.q_b_cp_async): + for c in range(tid, R * 1536 // 8, 256): + dst = qb.iterator + c * 8 + src = w_qb.iterator + (head * 192 + part * R) * 1536 + c * 8 + copy_evict(dst, src) + cute.arch.cp_async_commit_group() + cute.arch.griddepcontrol_wait() + cute.arch.griddepcontrol_launch_dependents() + qsi = q_scale_inv[0] + xv = cute.make_rmem_tensor((B, 8), Float32) + for b in cutlass.range_constexpr(B): + total = Float32(0) + if (tid < 192) & self.live(tok0 + b): + values = load_bf16x8(at(qkv, tok0 + b, tid * 8), mutable=True) + for i in cutlass.range_constexpr(8): + xv[b, i] = values[i] + total = fma(values[i], values[i], total) + total = warp_sum(total) + if lane == 0: + ss[warp, b] = total + cute.arch.cp_async_wait_group(1 if self.q_b_cp_async else 0) + cute.arch.sync_threads() + if tid < 192: + weights = load_bf16x8(qa.iterator + tid * 8, shared=True) + for b in cutlass.range_constexpr(B): + total = Float32(0) + for w in cutlass.range_constexpr(8): + total += ss[w, b] + rr = query_norm_rsqrt(total, eps) + for i in cutlass.range_constexpr(8): + xq[b, tid * 8 + i] = BFloat16(xv[b, i] * rr * weights[i]) + cute.arch.sync_threads() + if cutlass.const_expr(self.q_b_cp_async): + cute.arch.cp_async_wait_group(0) + cute.arch.sync_threads() + else: + wait(barw) + for row in range(warp, R, 8): + acc = zeros((B, 2), Float32) + for chunk in cutlass.range_constexpr(6): + k = (chunk * 32 + lane) * 8 + weights = load_bf16x8(qb.iterator + row * 1536 + k, shared=True) + for b in cutlass.range_constexpr(B): + values = load_bf16x8(xq.iterator + b * 1536 + k, shared=True) + for i in cutlass.range_constexpr(8): + acc[b, i % 2] = fma(weights[i], values[i], acc[b, i % 2]) + for b in cutlass.range_constexpr(B): + value = warp_sum(acc[b, 0] + acc[b, 1]) + if lane == 0: + qrow[row, b] = Float32(BFloat16(value)) + cute.arch.sync_threads() + cute.arch.cluster_wait() + nrows = min(R, max(0, 128 - R * part)) + for i in range(tid, nrows * B * self.cluster, 256): + dst, jb = i % self.cluster, i // self.cluster + j, b = jb % nrows, jb // nrows + cute.arch.store_async_dsmem( + qn.iterator + b * 128 + R * part + j, float_bits(qrow[j, b]), barq, dst + ) + for i in range(tid, R * B, 256): + j, b = i % R, i // R + p = R * part + j + if (p >= 128) & self.live(tok0 + b): + output[tok0 + b, head, 512 + p - 128] = fp8_bits(qrow[j, b] * qsi) + wait(barq) + latent_col, group = tid % L, tid // L + acc = zeros((B, 2), Float32) + for i in cutlass.range_constexpr(self.p_per_group): + p = group * self.p_per_group + i + w = Float32(uk[p, latent_col]) + for b in cutlass.range_constexpr(B): + acc[b, i % 2] = fma(qn[b, p], w, acc[b, i % 2]) + for b in cutlass.range_constexpr(B): + red[group, b, latent_col] = acc[b, 0] + acc[b, 1] + cute.arch.sync_threads() + for i in range(tid, L * B, 256): + b, ll = i // L, i % L + value = Float32(0) + for g in cutlass.range_constexpr(G): + value += red[g, b, ll] + if self.live(tok0 + b): + rounded = Float32(BFloat16(value)) + output[tok0 + b, head, L * part + ll] = fp8_bits(rounded * qsi) + if cutlass.const_expr(self.dcp): + self.publish(args, tok0, tid) + valid = (head == 0) & (part < B) & self.live(tok0 + part) + self.insert(args, ss, tok0, part, valid, tid) + if cutlass.const_expr(self.dcp): + self.gathered(args, tid) + + @cute.jit + def insert(self, args, ss, tok0, j, valid, tid): + """RMSNorm the current latent KV row of request tok0 + j and write it, with + k_pe, to its FP8 cache slot.""" + qkv, w_kva, (cache, slots), (k_scale_inv, eps) = ( + args[0], + args[2], + args[5:7], + args[8:10], + ) + lane, warp = tid % 32, tid // 32 + slot = cutlass.Int64(-1) + if valid: + slot = slots[tok0 + j] + if slot >= 0: + ksi = k_scale_inv[0] + xx = cute.make_rmem_tensor(8, Float32) + total = Float32(0) + if tid < 64: + values = load_bf16x8(at(qkv, tok0 + j, 1536 + tid * 8), mutable=True) + for i in cutlass.range_constexpr(8): + xx[i] = values[i] + total = fma(values[i], values[i], total) + total = warp_sum(total) + cute.arch.sync_threads() + if (lane == 0) & (warp < 2): + ss[warp, 0] = total + cute.arch.sync_threads() + page, offset = slot // cache.shape[1], slot % cache.shape[1] + if tid < 64: + mean = fma(ss[0, 0] + ss[1, 0], Float32(1 / 512), eps) + rr = cute.math.rsqrt(mean, fastmath=True) + weights = load_bf16x8(w_kva.iterator + tid * 8) + for i in cutlass.range_constexpr(8): + normed = Float32(BFloat16(xx[i] * rr * weights[i])) + cache[page, offset, tid * 8 + i] = fp8_bits(normed * ksi) + elif tid < 72: + c = tid - 64 + values = load_bf16x8(at(qkv, tok0 + j, 2048 + c * 8), mutable=True) + for i in cutlass.range_constexpr(8): + cache[page, offset, 512 + c * 8 + i] = fp8_bits(values[i] * ksi) + + @cute.kernel + def tiled(self, args): + """B>8: T requests per CTA. Q_B fragments load to registers before the PDL + wait; in-place RMSNorm of the staged rows, split-K Q_B MMAs over the warps, + nope rows gathered from the six producing CTAs by bulk DSMEM copies (into the + dead x rows), then W_UK MMAs (warp w owns latent columns 8w..8w+7).""" + qkv, w_qa, w_qb, w_uk, q_scale_inv = (args[i] for i in (0, 1, 3, 4, 7)) + eps, output = args[9], args[10] + tid, _, _ = cute.arch.thread_idx() + part, head, tok0 = cute.arch.block_idx() + tok0 *= self.tokens + smem = cutlass.utils.SmemAllocator() + T, R, L = self.tokens, self.rows, self.latent + MT, senders = T // 16, (128 + R - 1) // R + lane, warp = tid % 32, tid // 32 + group, quarter = lane // 4, lane % 4 + xq = alloc(smem, BFloat16, (T, KP), 128) + uk = alloc(smem, BFloat16, (128, L + 8), 128) # 144 B rows for ldmatrix + qa = alloc(smem, BFloat16, 1536) + ss = alloc(smem, Float32, (8, 2)) + barq = smem.allocate(cutlass.Uint64, 8) + # After the Q_B MMAs the x rows hold the split-K partials, then our nope rows + # and the six producers' (peers write them only after our cluster arrival). + red = view(xq.iterator, 0, Float32, (8, T * R // 32, 32), (T * R, 32, 1)) + qblk = view(xq.iterator, 8 * T * R * 4, BFloat16, (T, R), (R, 1)) + qn = view(xq.iterator, 9 * T * R * 4, BFloat16, (senders, T, R), (T * R, R, 1)) + if tid == 0: + cute.arch.mbarrier_init(barq, 1) + cute.arch.mbarrier_init_fence() + cute.arch.mbarrier_arrive_and_expect_tx(barq, senders * T * R * 2) + # Split-K Q_B: warp w owns k = 192w..192w+191 (8 values per lane and 32-wide + # chunk; the x fragments use the same permutation). L2 now, registers later. + rows = head * 192 + part * R + group + qb_at = [ + [ + at(w_qb, rows + 8 * n, (6 * warp + cc) * 32 + 8 * quarter) + for n in range(3) + ] + for cc in range(6) + ] + for chunk in qb_at: + for address in chunk: + prefetch_l2(address) + for c in range(tid, 128 * L // 8, 256): + p, j = c // (L // 8), c % (L // 8) + dst = uk.iterator + p * (L + 8) + j * 8 + copy_evict(dst, at(w_uk, head, p, part * L + j * 8)) + if tid < 192: + copy_evict(qa.iterator + tid * 8, w_qa.iterator + tid * 8) + cute.arch.cp_async_commit_group() + cute.arch.griddepcontrol_wait() + cute.arch.griddepcontrol_launch_dependents() + for c in range(tid, T * 192, 256): + t, j = c // 192, c % 192 + if self.live(tok0 + t): + copy_g2s(xq.iterator + t * KP + j * 8, at(qkv, tok0 + t, j * 8)) + cute.arch.cp_async_commit_group() + qsi = q_scale_inv[0] + # The KV insert overlaps the staging of x (CTA j of the tile owns request j). + j = head * self.cluster + part + self.insert(args, ss, tok0, j, (j < T) & self.live(tok0 + j), tid) + cute.arch.cp_async_wait_group(0) + cute.arch.sync_threads() + # Packed FP32 pairs: the same IEEE ops in the same order as the scalar form. + gamma = [] + for i in cutlass.range_constexpr(6): + gamma += list(load_bf16x8(qa.iterator + 256 * i + 8 * lane, shared=True)) + for pair in cutlass.range_constexpr(T // 16): # requests warp + 8k, two at once + xrows = [xq.iterator + (warp + 16 * pair + 8 * h) * KP for h in range(2)] + values: list = [] + rr: list = [] + for h in cutlass.range_constexpr(2): + values.append([]) + for i in cutlass.range_constexpr(6): + for word in smem_words(xrows[h] + 256 * i + 8 * lane): + values[h] += list(unpack_bf16(word)) + for h in cutlass.range_constexpr(2): + zero = Float32(0), Float32(0) + chains = [zero, zero] # (0, 1), (2, 3): element e feeds chain e % 4 + for e in cutlass.range_constexpr(0, 48, 2): + x = values[h][e], values[h][e + 1] + chains[e % 4 // 2] = fma2(x, x, chains[e % 4 // 2]) + (c0, c1), (c2, c3) = chains + rr.append(query_norm_rsqrt(warp_sum((c0 + c1) + (c2 + c3)), eps)) + for h in cutlass.range_constexpr(2): + for i in cutlass.range_constexpr(6): + words = [] + for e in cutlass.range_constexpr(8 * i, 8 * i + 8, 2): + x = values[h][e], values[h][e + 1] + y = mul2(mul2(x, (rr[h], rr[h])), (gamma[e], gamma[e + 1])) + words.append(bf16x2(*y)) + store_shared_words(xrows[h] + 256 * i + 8 * lane, tuple(words)) + qbr = [[load_words(address, count=4) for address in chunk] for chunk in qb_at] + cute.arch.sync_threads() + acc = zeros((MT * 3, 4), Float32) + for cc in cutlass.range_constexpr(6): + k = (6 * warp + cc) * 32 + 8 * quarter + z = qbr[cc] + for m in cutlass.range_constexpr(MT): + x, y = ( + load_words(xq.iterator + r * KP + k, count=4, shared=True) + for r in (16 * m + group, 16 * m + group + 8) + ) + for n in cutlass.range_constexpr(3): + for s in cutlass.range_constexpr(2): + a = (x[2 * s], y[2 * s], x[2 * s + 1], y[2 * s + 1]) + b = (z[n][2 * s], z[n][2 * s + 1]) + mma_into(acc[m * 3 + n, None], a, b, bf16=True) + cute.arch.sync_threads() + cute.arch.cluster_arrive_relaxed() # x is consumed: peers may fill qn + for i in cutlass.range_constexpr(MT * 12): + red[warp, i, lane] = acc[i // 4, i % 4] + cute.arch.sync_threads() + for e in range(tid, T * R, 256): + i, ln = e // 32, e % 32 + value = Float32(0) + for w in cutlass.range_constexpr(8): + value += red[w, i, ln] + value = Float32(BFloat16(value)) + tile, c = i // 4, i % 4 + t = 16 * (tile // 3) + ln // 4 + 8 * (c // 2) + r = 8 * (tile % 3) + 2 * (ln % 4) + c % 2 + p = R * part + r + if p < 128: + qblk[t, r] = BFloat16(value) + elif self.live(tok0 + t): + output[tok0 + t, head, 512 + p - 128] = fp8_bits(value * qsi) + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_threads() + cute.arch.cluster_wait() + if (part < senders) & (tid < self.cluster): + copy_s2cluster( + qn.iterator + part * T * R, qblk.iterator, barq, tid, size=T * R * 2 + ) + wait(barq) + cute.arch.cluster_arrive_relaxed() # peers may exit once qblk is delivered + b_frags = [ + ldmatrix(uk.iterator + (32 * kp + lane) * (L + 8) + 8 * warp, trans=True) + for kp in range(4) + ] + for m in cutlass.range_constexpr(MT): + out = zeros(4, Float32) + t = 16 * m + lane % 8 + 8 * ((lane // 8) % 2) + for ks in cutlass.range_constexpr(8): + p = 16 * ks + 8 * (lane // 16) + a = ldmatrix(qn.iterator + ((p // R) * T + t) * R + p % R) + b = b_frags[ks // 2][2 * (ks % 2) : 2 * (ks % 2) + 2] + mma_into(out, a, b, bf16=True) + for h in cutlass.range_constexpr(2): # rows group, group + 8: 2 bytes + row = 16 * m + group + 8 * h + if self.live(tok0 + row): + lo, hi = (Float32(BFloat16(out[2 * h + e])) * qsi for e in range(2)) + col = L * part + 8 * warp + 2 * quarter + store_fp8_pair(at(output, tok0 + row, head, col), lo, hi) + if cutlass.const_expr(self.dcp): + self.publish(args, tok0, tid) + cute.arch.cluster_wait() + if cutlass.const_expr(self.dcp): + self.gathered(args, tid) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/value.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/value.py new file mode 100644 index 0000000000..e074498400 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/mla/value.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MLA value projection (W_UV) and sigmoid gate of the BF16 attention latent.""" + +import cutlass +import cutlass.cute as cute +from cutlass import BFloat16, Float32, Uint32 + +from ..common.primitives import ( + alloc, + at, + bf16x2, + copy_g2s, + launch, + load_words, + mma_into, + sigmoid, + store_words, + unpack_bf16, + zeros, +) + + +class MlaValue: + """W_UV and the sigmoid gate, [B, H, 512] BF16 latent -> [B, H*128]. One + 128-thread CTA per (32 outputs, head, 16 rows); warp w owns outputs + 8w..8w+7 over the whole latent (BF16 MMA, FP32 accumulation). W_UV must have a + contiguous latent axis (the absorbed kv_b_proj view). Rows are tokens, ``q`` per + request (``lengths`` per request).""" + + def __init__(self, heads: int, batch: int, q: int = 1): + self.heads, self.batch, self.q = heads, batch, q + self.grid = (4, heads, (batch + 15) // 16) + + @cute.jit + def __call__(self, *args): + launch(self.kernel, args, self.grid, 128) + + @cute.kernel + def kernel(self, args): + latent, w_uv, gate, lengths, output = args + rows = alloc(cutlass.utils.SmemAllocator(), BFloat16, (16, 544), 128) + tid, _, _ = cute.arch.thread_idx() + chunk, head, tile = cute.arch.block_idx() + lane, warp = tid % 32, tid // 32 + group, quarter = lane // 4, lane % 4 + t0, v = tile * 16, chunk * 32 + warp * 8 + group + # B fragments: 8 latent values per 32-wide chunk; A uses the same permutation. + wb = [ + load_words(at(w_uv, head, 32 * c + 8 * quarter, v), count=4) + for c in range(16) + ] + cute.arch.griddepcontrol_wait() + for c in range(tid, 16 * 64, 128): + r, j = c // 64, c % 64 + if t0 + r < self.batch: + src = at(latent, t0 + r, head, j * 8) + copy_g2s(rows.iterator + r * 544 + j * 8, src) + cute.arch.cp_async_commit_group() + col = head * 128 + chunk * 32 + warp * 8 + 2 * quarter + gates = [] + for r in cutlass.range_constexpr(2): + word = Uint32(0) + if t0 + group + 8 * r < self.batch: + word = load_words( + at(gate, t0 + group + 8 * r, col), count=1, mutable=True + )[0] + gates.append(unpack_bf16(word)) + cute.arch.griddepcontrol_launch_dependents() + cute.arch.cp_async_wait_group(0) + cute.arch.sync_threads() + parts = zeros((2, 4), Float32) # two accumulation chains + for c in cutlass.range_constexpr(16): + x, y = ( + load_words( + rows.iterator + r * 544 + 32 * c + 8 * quarter, count=4, shared=True + ) + for r in (group, group + 8) + ) + for s in cutlass.range_constexpr(2): + a = (x[2 * s], y[2 * s], x[2 * s + 1], y[2 * s + 1]) + b = (wb[c][2 * s], wb[c][2 * s + 1]) + mma_into(parts[s, None], a, b, bf16=True) + acc = [parts[0, k] + parts[1, k] for k in range(4)] + for r in cutlass.range_constexpr(2): + t = t0 + group + 8 * r + if t < self.batch: + # Padded rows may hold a NaN gate; keep inactive rows zero. + lo, hi = Float32(0), Float32(0) + if lengths[t // self.q] > 0: + lo, hi = ( + Float32(BFloat16(acc[2 * r + e])) * sigmoid(gates[r][e], False) + for e in range(2) + ) + store_words(at(output, t, col), (bf16x2(lo, hi),)) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/__init__.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/dense.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/dense.py new file mode 100644 index 0000000000..b95720ca1d --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/dense.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Prepared first dense MLP, sharing projection and SiTU decode arithmetic.""" + +from functools import cache + +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import make_fake_stream + +from ..common.host import architecture, dynamic_fake, dynamic_view, stream +from ..common.primitives import launch +from ..gemm.projection import PreparedProjection +from .shared import situ + + +class DenseSitu: + def __init__(self, hidden, padded): + self.hidden, self.padded = hidden, padded + + @cute.jit + def __call__(self, *args): + grid = (cute.ceil_div(self.padded, 256), args[0].shape[0], 1) + launch(self.kernel, args, grid, 256) + + @cute.kernel + def kernel(self, args): + x, out, beta, linear = args[0:4] + tile, row, _ = cute.arch.block_idx() + tid, _, _ = cute.arch.thread_idx() + col = tile * 256 + tid + cute.arch.griddepcontrol_wait() + if col < self.padded: + value = cutlass.Float32(0) + if col < self.hidden: + gate = cutlass.Float32(x[row, col]) + up = cutlass.Float32(x[row, self.hidden + col]) + value = situ(gate, up, beta, linear, native=True) + out[row, col] = cutlass.BFloat16(value) + cute.arch.griddepcontrol_launch_dependents() + + +@cache +def _compile_activation(hidden, padded, arch): + args = (dynamic_fake(2), dynamic_fake(2), *map(cutlass.Float32, (4, 25))) + kernel = DenseSitu(hidden, padded) + return cute.compile(kernel, *args, make_fake_stream(), options=f"--gpu-arch={arch}") + + +class PreparedDenseMLP: + """Dense MLP; TP16 pads down K 2112 -> 2176 and the activation zero-fills it.""" + + def __init__(self, mlp, batches, reduce=None): + gate_up, down = mlp.gate_up_proj.weight, mlp.down_proj.weight + self.device, self.reduce = gate_up.device, reduce + self.hidden = hidden = gate_up.shape[0] // 2 + self.padded = padded = (hidden + 127) // 128 * 128 + linear = -1 if mlp.act_fn.linear_beta is None else mlp.act_fn.linear_beta + self.scalars = tuple(map(cutlass.Float32, (mlp.act_fn.beta, linear))) + if padded != hidden: + down = torch.nn.functional.pad(down.detach(), (0, padded - hidden)) + self.gates = gate_up.new_empty((64, 2 * hidden)) + self.hidden_rows = gate_up.new_empty((64, padded)) + tile = {b: next(n for n in (8, 16, 32, 64) if b <= n) for b in batches} + up_tiles = {8: (64, 8, 4, 8), 16: (64, 16, 2, 8), 32: (64, 32, 2, 8)} + up_tiles[64] = (64, 64, 2, 5) + down_tiles = {8: (128, 8, 1, 6), 16: (128, 16, 1, 6), 32: (128, 32, 1, 5)} + down_tiles[64] = (64, 64, 1, 5) + policy = lambda tiles: {b: tiles[n] for b, n in tile.items()} + self.gate = PreparedProjection(gate_up, policy(up_tiles)) + self.down = PreparedProjection(down, policy(down_tiles)) + self.gate.prepare(), self.down.prepare() + self.activation = _compile_activation(hidden, padded, architecture(self.device)) + + def run(self, x, out): + gate, hidden = self.gates[: x.shape[0]], self.hidden_rows[: x.shape[0]] + self.gate.run(x, gate) + views = (dynamic_view(gate), dynamic_view(hidden)) + self.activation(*views, *self.scalars, stream(self.device)) + self.down.run(hidden, out) + if self.reduce is not None: + reduced = self.reduce(out) + if reduced.data_ptr() != out.data_ptr(): + out.copy_(reduced) + return out diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts.py new file mode 100644 index 0000000000..80aa8bcb2f --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts.py @@ -0,0 +1,1579 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Persistent routed experts: TMA producer, MXFP8 activations, FP4 x MXFP8 FC1/FC2 MMAs, +SiTU hidden exchange and output epilogue. Capacity <= 8 keeps the batch's activations in +SMEM; capacity > 8 (grouped) runs (expert, <= 8 tokens) units whose FC1 stages gather +8-row slices of the grid-quantized batch.""" + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int8, Int16, Int32, Int64, Uint8, Uint32, Uint64 + +from ..common.primitives import ( + acquire_tensormap, + alloc, + asm, + bf16, + bf16x2, + bits_float, + commit, + copy_evict, + copy_g2s, + fence_tmem, + float_bits, + fma, + fp8_bits, + group_barrier, + launch, + load_tma3, + load_tmem, + load_volatile, + load_words, + multicast_words, + mx_exponent, + mx_rescale, + pack_fp8, + pack_pair, + release_tmem, + shared_address, + situ, + smem_desc, + smem_words, + store_shared_words, + store_words, + thread, + unpack_bf16, + view, + wait, + wait_tmem, +) +from ..common.twoshot import push + +shuffle = cute.arch.shuffle_sync + + +class ExpertsGeometry: + """Static plain-input geometry. No cluster-local or multicast storage. The workspace + is capacity-independent; capacity > 8 swaps the batch's SMEM activation image for + per-slot 8-row images, which frees a sixth weight stage, and keeps per-CTA windows + of the unit tables (expert, 8 pair rows) in the misc block.""" + + def __init__(self, tp, capacity=8): + self.tp = tp + self.tiles, self.w13_rows, self.h_chunks, self.fc2_stages = ( + (6, 768, 3, 2) if tp == 8 else (3, 512, 2, 1) + ) + self.grouped = capacity > 8 + self.stages = stages = 6 if self.grouped else 5 + self.units = 16 * capacity if self.grouped else 128 + self.w13_scale_bytes = self.w13_rows * 112 + self.w2_scale_columns = 4 * self.h_chunks + self.w2_scale_bytes = 3584 * self.w2_scale_columns + self.w2_tile_scale_bytes = 128 * self.w2_scale_columns + self.h_image_bytes = self.h_chunks * 1024 + self.h_scale_bytes = self.h_chunks * 128 + 384 + self.h_compact_bytes = self.h_image_bytes + 128 + self.off_sfa = stages * 32768 + self.off_x = self.off_sfa + stages * 1024 + self.off_h = self.off_x + (stages * 2048 if self.grouped else 28672) + self.off_xsf = self.off_h + 2 * self.h_image_bytes + self.off_hsf = self.off_xsf + (stages * 256 if self.grouped else 28 * 128) + 384 + self.off_misc = self.off_hsf + 2 * self.h_scale_bytes + self.smem_bytes = self.off_misc + 5120 + # Misc barriers (byte offsets): weight full/empty [stages], accumulator + # full/empty [16], hidden full/empty [2], x ready [14]. Grouped stage metadata + # and unit windows (4 x 16 units) use the B<=8 routing tables' space. + self.bar_full, self.bar_empty, self.bar_accfull = (0, 8 * stages, 16 * stages) + acc = self.bar_accfull + self.bar_accempty, self.bar_hfull = (acc + 128, acc + 256) + self.bar_hempty, self.bar_xready = (acc + 272, acc + 288) + self.off_meta = self.off_misc + (1152 if self.grouped else 520) + self.off_window = self.off_misc + 1280 + # Stealing: FC2 unit queue (producer -> hidden loader) [8], full/empty [8, 8]. + self.off_queue, self.bar_queue = (self.off_misc + 4608, 4640) + self.ws_partial = 2 * 128 * self.h_compact_bytes + self.ws_counter = self.ws_partial + 2 * 160 * 8 * 128 * 4 + self.ws_xflag = self.ws_counter + 160 * 4 + self.ws_stage = (self.ws_xflag + 160 * 4 + 255) // 256 * 256 + # Publish stage rows [64 tokens x 16 slots, 3584] BF16, then grouped: call + # counters [160] and unit counts [2], hidden images [2, 1024]. + self.ws_calls = self.ws_stage + 1024 * 3584 * 2 + self.ws_units = self.ws_calls + 1024 + # Pregrouped unit table (``ExpertsGroupKernel``): header (units, call, grouped + # call), experts [1024], rows [1024, 8]. + self.ws_group = self.ws_units + 2 * 1024 * self.h_compact_bytes + # FC2(B) tail claim counters [2 parities, 160] (``ExpertsKernel.steal``). + self.ws_steal = self.ws_group + 16 + 20 * 1024 + self.workspace_bytes = self.ws_steal + 2 * 160 * 4 + + def image(self, workspace, parity, unit): + """Compact hidden image of ``unit`` in the ``parity`` half of the workspace.""" + if self.grouped: + base = workspace.iterator + self.ws_units + return base + (Int64(parity) * 1024 + unit) * self.h_compact_bytes + return workspace.iterator + (Int64(parity) * 128 + unit) * self.h_compact_bytes + + +@cute.jit +def swizzled(chunk, row, base=0): + """Byte offset of 16-byte ``chunk`` of ``row`` in a 128B-swizzled 8-row image.""" + return base + chunk // 8 * 1024 + row * 128 + (chunk % 8 ^ row) * 16 + + +@cute.jit +def has_ff_byte(word: Uint32): + value = ~word + return value - Uint32(16843009) & ~value & Uint32(2155905152) != 0 + + +@cute.jit +def mx_block(words): + """MXFP8 of one 32-BF16 block (16 words): eight e4m3x4 words and the E8M0 scale.""" + values = cute.make_rmem_tensor(32, Float32) + for i in cutlass.range_constexpr(16): + pair = unpack_bf16(words[i]) + values[2 * i], values[2 * i + 1] = (pair[0], pair[1]) + maxima = cute.make_rmem_tensor(16, Float32) + for j in cutlass.range_constexpr(16): + maxima[j] = cute.arch.fmax(abs(values[j]), abs(values[j + 16])) + for width in (8, 4, 2, 1): + for j in cutlass.range_constexpr(width): + maxima[j] = cute.arch.fmax(maxima[j], maxima[j + width]) + scale = mx_exponent(maxima[0]) + reciprocal = mx_rescale(scale) + packed = [ + pack_fp8(tuple(values[4 * j + i] * reciprocal for i in range(4))) + for j in range(8) + ] + return packed, scale + + +# Grouped expert plan, written before the experts run (by the routing kernel's copy +# role, or ExpertsPlanKernel): MXFP8 rows [64, 3584] then scale words [64, 28, 16 B] +# (bytes 0-3 used), the layout of the experts' per-slot SMEM images. +PLAN_XSF = 64 * 3584 +PLAN_BYTES = PLAN_XSF + 64 * 448 + + +@cute.jit +def plan_chunk(plan, token, chunk, words): + """Write 8-element ``chunk`` (four BF16x2 ``words``) of ``token`` as MXFP8 into + the plan; lanes 4k..4k+3 hold one 32-element block, 16 lanes one scale word. + Bit-identical to ``mx_block`` (an exact max, the same scaling and packing).""" + lane, values = (thread()[2], cute.make_rmem_tensor(8, Float32)) + for i in cutlass.range_constexpr(4): + pair = unpack_bf16(words[i]) + values[2 * i], values[2 * i + 1] = (pair[0], pair[1]) + top = abs(values[0]) + for i in cutlass.range_constexpr(1, 8): + top = cute.arch.fmax(top, abs(values[i])) + for width in (1, 2): + top = cute.arch.fmax(top, cute.arch.shuffle_sync_bfly(top, width)) + scale = mx_exponent(top) + reciprocal = mx_rescale(scale) + packed = [ + pack_fp8(tuple(values[4 * j + i] * reciprocal for i in range(4))) + for j in range(2) + ] + store_words(plan.iterator + token * 3584 + chunk * 8, packed) + word = Uint32(scale) + for s in cutlass.range_constexpr(1, 4): + word |= shuffle(Uint32(scale), lane + 4 * s) << 8 * s + if chunk % 16 == 0: + dst = plan.iterator + PLAN_XSF + token * 448 + chunk // 16 * 16 + store_words(dst, (word, Uint32(0), Uint32(0), Uint32(0))) + + +class ExpertsPlanKernel: + """Standalone plan writer (no routing kernel in front): one CTA per token.""" + + def __init__(self, batch): + self.batch = batch + + @cute.jit + def __call__(self, *args): + launch(self.kernel, args, self.batch, 256) + + @cute.kernel + def kernel(self, args): + x, plan = args + tid, _, _, token = thread() + cute.arch.griddepcontrol_wait() + for i in cutlass.range_constexpr(2): + chunk = tid + 256 * i + if chunk < 448: # warp-uniform: chunks 256-447 are warps 0-5 + words = load_words( + x.iterator + token * 3584 + chunk * 8, count=4, mutable=True + ) + plan_chunk(plan, token, chunk, words) + cute.arch.griddepcontrol_launch_dependents() + + +@asm() +def stage(mma, weight_scales, activation_scales, release, done, *, kind): + """One FC1/FC2 stage. ``mma``: (TMEM accumulator, weight desc, activation desc); + scales: (TMEM column, desc[, activation copy flag]); ``release``: (FC1 accumulate + or FC2 prior-H release, weight empty barrier); ``done``: accumulator barrier.""" + half = kind in ("fc2b", "fc2p") + steps = 6 if kind == "fc2" else 4 if half else 8 + cp = "tcgen05.cp.cta_group::1.32x128b.warpx4" + arrive = "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64" + code = [ + "{ .reg .pred e,p0,pt,pc,pa,ph; .reg .b32 i<4>; .reg .b64 as1,bs1,bs2;", + f".reg .b32 af1,bf1,bf2; .reg .b64 xa<{steps}>,xb<{steps}>;", + "elect.sync _|e,0xffffffff; setp.eq.b32 pt,0,0; setp.ne.b32 pc,$7,0;", + "and.pred pc,pc,e; setp.ne.b32 pa,$10,0; and.pred pa,pa,e;", + ] + if kind == "fc1": + code.append("setp.ne.b32 p0,$8,0;") + elif half: + code.append("setp.eq.b32 p0,0,0;") + else: + code.append("setp.eq.b32 p0,0,1; setp.ne.b32 ph,$8,0; and.pred ph,ph,e;") + code.append(f"@ph {arrive} [$8];") + code += [f"mov.b32 i{k},{142738048 | k << 29 | k << 4};" for k in range(4)] + if not half: + code.append("add.s64 as1,$4,32; add.s64 bs1,$6,8;") + code.append("add.u32 af1,$3,4; add.u32 bf1,$5,4;") + if kind == "fc2a": + code.append("add.s64 bs2,$6,16; add.u32 bf2,$5,8;") + code += [f"@pc {cp} [$5],$6;", f"@pc {cp} [bf1],bs1;"] + if kind == "fc2a": + code.append(f"@pc {cp} [bf2],bs2;") + code.append(f"@e {cp} [$3],$4;") + if not half: + code.append(f"@e {cp} [af1],as1;") + for k in range(steps): + a = k // 4 * 1024 + k % 4 * 2 + (4 if kind == "fc2" and k >= 4 else 0) + b = k // 4 * 64 + k % 4 * 2 + code.append(f"add.s64 xa{k},$1,{a}; add.s64 xb{k},$2,{b};") + for k in range(steps): + af, bf = ("$3", "$5") if k < 4 else ("af1", "bf1") + code.append( + "@e tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale " + f"[$0],xa{k},xb{k},i{k % 4},[{af}],[{bf}],{'p0' if k == 0 else 'pt'};" + ) + if kind != "fc2p": + code.append(f"@e {arrive} [$9];") + if kind != "fc2a": + code.append(f"@pa {arrive} [$10];") + code.append("}") + kinds = "rllrlrlrrrr" + operands = (*mma, *weight_scales, *activation_scales, *release, done) + values = [(Uint64 if k == "l" else Uint32)(v) for k, v in zip(kinds, operands)] + return "\n".join(code), ",".join(kinds) + ",~{memory}", values + + +@asm(Uint32) +def claim(pointer, value): + """Relaxed GPU-scope fetch-add (FC2 tail claims): the counter before the add.""" + code = "atom.relaxed.gpu.global.add.u32 $0,[$1],$2;" + return code, "=r,l,r,~{memory}", (pointer, Uint32(value)) + + +def fc2_pairs(tp, capacity, streamed): + """TP8 FC2 jobs are tile pairs on the incremental producer from capacity 4 (the + grouped engine always streams): chunks 0-1 of tiles t and t + 1, then both tiles' + chunk 2 through one 256-row box (map 5), so every stage is 16 KB.""" + return tp == 8 and streamed and capacity >= 4 + + +class ExpertsKernel(ExpertsGeometry): + """The engine for one capacity and persistent grid. CTAs wait on other clusters' + partials (``load_hidden``, ``partials``, ``finalize``), so + ``PersistentExperts.prepare`` checks the grid is co-resident. It is not launched + cooperatively: that forfeits the PDL prologue overlap (1-2% decode at B4-B8). + ``publish`` stages FC2 rows for ``finalize``: routing weights replace the output + arg, and (latent multicast, parity, rank) follow the SiTU scalars; ``"twoshot"``: + (receive slots [3 x TP, 64, 3584/TP], their owners' bases, parity x TP + rank) of + ``common.twoshot`` instead. Capacity > 8 (grouped) takes the plan (``PLAN_BYTES``: + MXFP8 x) in place of x. ``stream`` runs capacity <= 8 on the incremental producer + and CTA-scoped waits (grouped always). ``pregrouped`` (grouped) reads the unit table + ``ExpertsGroupKernel`` left in the workspace instead of grouping the routes itself. + ``steal`` (grouped) balances the end: each CTA's last ``zone`` FC2(B) grains (TP8 a + tile pair, TP16 three tiles) are claimed by the owner from the front and by the + ``victims`` CTAs before it from the back, through one counter per CTA re-armed by + call parity; a stolen grain is computed exactly as by its owner.""" + + def __init__( + self, tp, capacity, experts, grid, publish=False, + stream=True, pregrouped=False, steal=False, + ): # fmt: skip + super().__init__(tp, capacity) + self.pregrouped = pregrouped and self.grouped + self.steal = steal and self.grouped + self.streamed = stream or self.grouped + self.paired = fc2_pairs(tp, capacity, self.streamed) + self.jobs = 14 if self.paired else 28 # FC2 jobs per unit + self.grain, self.zone, self.victims = (1 if self.paired else 3, 8, 3) + self.threads = 384 + self.capacity, self.experts, self.grid = (capacity, experts, grid) + self.publish = publish + assert not self.grouped or grid >= capacity <= 64 + self.width = min(capacity, 8) + self.dense = capacity <= 4 + self.rounds = (112 * capacity + 319) // 320 if self.dense else 3 + + def sync(self, barrier, phase=0, tmem=False): + """Parity wait (then tcgen05 ordering if ``tmem``). Every barrier here is + CTA-local, so grouped and streamed waits use CTA scope.""" + scope = "cta" if self.streamed else "cluster" + (wait_tmem if tmem else wait)(barrier, phase, scope=scope) + + @cute.jit + def route(self, sm, ids, counters, batch: Int32): + """Dedupe after the dependency wait. Capacity <= 2: one work item per routing + pair; else first occurrences across tokens in order, keeping token slots.""" + tid, warp, lane, cta = thread() + misc = self.off_misc + table = view(sm.iterator, 0, Int32, self.experts) + expert, uof = (view(sm.iterator, misc + o, Int32, 128) for o in (1140, 1652)) + slots = view(sm.iterator, misc + 2176, Int8, 1024) + slot_words = view(sm.iterator, misc + 2176, Uint32, 256) + counts = view(sm.iterator, misc + 868, Int32, 4) + distinct, parity = (view(sm.iterator, misc + o, Int32, 1) for o in (860, 864)) + for i in range(tid, self.experts, 384): + table[i] = Int32(2147483647) + myid, pairs = (Int32(-1), batch * 16) + if tid < 128: + lo, hi = (Uint32(4294967295), Uint32(4294967295)) + if tid < pairs: + myid = ids[tid] + if cutlass.const_expr(self.capacity <= 2): + expert[tid] = myid + if tid < 16: + lo = lo & Uint32(4294967040) | Uint32(tid % 16) + else: + lo = lo & Uint32(4294902015) | Uint32(tid % 16) << 8 + slot_words[2 * tid], slot_words[2 * tid + 1] = (lo, hi) + if tid == 160: + parity[0] = counters[cta] + if cutlass.const_expr(self.capacity <= 2) and tid == 0: + distinct[0] = pairs + cute.arch.sync_threads() + if cutlass.const_expr(self.capacity > 2): + if tid < pairs: + cute.arch.atomic_min( + table.iterator + myid, tid, sem="relaxed", scope="cta" + ) + cute.arch.sync_threads() + if tid < 128: + firstk = Int32(-1) + if tid < pairs: + firstk = table[myid] + first = (tid < pairs) & (firstk == tid) + ballot = Uint32(cute.arch.vote_ballot_sync(first)) + if lane == 0: + counts[warp] = Int32(cute.arch.popc(ballot)) + cute.arch.barrier(barrier_id=1, number_of_threads=128) + base = Int32(0) + for w in cutlass.range_constexpr(3): + if warp > w: + base += counts[w] + rank = base + Int32(cute.arch.popc(ballot & (Uint32(1) << lane) - 1)) + if first: + uof[tid] = rank + expert[rank] = myid + if tid == 0: + distinct[0] = counts[0] + counts[1] + counts[2] + counts[3] + cute.arch.barrier(barrier_id=1, number_of_threads=128) + if tid < pairs: + slots[uof[firstk] * 8 + tid // 16] = Int8(tid % 16) + # Retire every thread's hash accesses before async weight-ring reuse. + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_threads() + return (distinct[0], parity[0], expert) + + def tables(self, sm, window=True): + """Grouped unit tables: (expert, eight pair rows (Int16, -1 padded), the rows + as four words per unit) of every unit in the weight-ring scratch during + routing, or of this CTA's windows. Views all start from the byte pointer.""" + base, units = (self.off_window, 64) if window else (12288, self.units) + experts = view(sm.iterator, base, Int32, units) + rows = view(sm.iterator, base + 4 * units, Int16, 8 * units) + return experts, rows, view(sm.iterator, base + 4 * units, Uint32, 4 * units) + + @cute.jit + def first_unit(self, ranges, fc2: cutlass.Constexpr, group): + """First unit of this CTA's FC1 or FC2 range in ``group``.""" + first = ranges[group, 4] // self.jobs + if cutlass.const_expr(not fc2): + first = ranges[group, 0] + ranges[group, 2] // 14 // self.tiles + return first + + @cute.jit + def local(self, ranges, unit, fc2: cutlass.Constexpr, group): + """Window index of ``unit`` (windows: FC1 A, FC1 B, FC2 A, FC2 B).""" + return (2 * fc2 + group) * 16 + unit - self.first_unit(ranges, fc2, group) + + @cute.jit + def windows(self, sm, ranges): + """Copy this CTA's four unit windows (<= 16 units each) out of the scratch.""" + tid, _, _, cta = thread() + experts, _, words = self.tables(sm, False) + window_experts, _, window_words = self.tables(sm) + for w in cutlass.range_constexpr(4): + fc2, group = (w // 2, w % 2) + lo, hi = (ranges[group, 2 + 2 * fc2], ranges[group, 3 + 2 * fc2]) + if cutlass.const_expr(self.steal and w == 3): + hi = self.stolen_end(ranges, cta) + last = (hi - 1) // self.jobs + if cutlass.const_expr(not fc2): + last = ranges[group, 0] + (hi - 1) // 14 // self.tiles + unit = self.first_unit(ranges, fc2, group) + tid % 16 + if (tid // 16 == w) & (lo < hi) & (unit <= last): + window_experts[tid] = experts[unit] + for i in cutlass.range_constexpr(4): + window_words[4 * tid + i] = words[4 * unit + i] + + @cute.jit + def stolen_end(self, ranges, cta): + """End of the FC2(B) jobs of CTAs up to ``cta + victims`` (window 3's reach).""" + stop = cta + self.victims + 1 + if stop > self.grid: + stop = Int32(self.grid) + count = ranges[1, 1] // (self.tiles * 14) + return self.jobs * ranges[1, 0] + stop * self.jobs * count // self.grid + + @cute.jit + def group(self, sm, ids, workspace, batch: Int32, scalars=None): + """Grouped routing: per-expert token masks; expert e's + tokens, in token order, fill ceil(n / 8) units numbered in expert order. Unit + rows hold pair indices (token * 16 + slot), -1 when padded. Misc scalars: + units, call, grouped call (``scalars`` given: the units only).""" + tid, warp, lane, cta = thread() + masks = view(sm.iterator, 0, Uint32, 2 * self.experts) + base = view(sm.iterator, 8 * self.experts, Int32, self.experts) + totals = view(sm.iterator, 12 * self.experts, Int32, self.threads // 32) + experts, rows, row_words = self.tables(sm, False) + counted = scalars is None + if cutlass.const_expr(counted): + scalars = view(sm.iterator, self.off_misc + 1140, Int32, 3) + threads: cutlass.Constexpr = self.threads + rounds: cutlass.Constexpr = (16 * self.capacity + threads - 1) // threads + mine = cute.make_rmem_tensor(rounds, Int32) + for k in cutlass.range_constexpr(rounds): + mine[k] = Int32(0) + if tid + self.threads * k < batch * 16: + mine[k] = ids[tid + self.threads * k] + if cutlass.const_expr(counted): + if tid == 0: + scalars[1] = view(workspace.iterator, self.ws_counter, Int32, 160)[cta] + if tid == 1: + scalars[2] = view(workspace.iterator, self.ws_calls, Int32, 160)[cta] + for i in range(tid, 2 * self.experts, self.threads): + masks[i] = Uint32(0) + for i in range(tid, 4 * self.units, self.threads): + row_words[i] = Uint32(4294967295) + cute.arch.sync_threads() + for k in cutlass.range_constexpr(rounds): + pair = tid + self.threads * k + if pair < batch * 16: + token = pair >> 4 + bit = Uint32(1) << (token & 31) + word = masks.iterator + 2 * mine[k] + (token >> 5) + cute.arch.atomic_or(word, bit, sem="relaxed", scope="cta") + cute.arch.sync_threads() + span: cutlass.Constexpr = (self.experts + self.threads - 1) // self.threads + units = cute.make_rmem_tensor(span, Int32) + count = Int32(0) + for j in cutlass.range_constexpr(span): + units[j], e = (Int32(0), span * tid + j) + if e < self.experts: + n = cute.arch.popc(masks[2 * e]) + cute.arch.popc(masks[2 * e + 1]) + units[j] = Int32(n + 7) >> 3 + count += units[j] + total = count + for shift in cutlass.range_constexpr(5): + other = cute.arch.shuffle_sync_up(total, 1 << shift) + if lane >= 1 << shift: + total += other + if lane == 31: + totals[warp] = total + cute.arch.sync_threads() + start = total - count + for w in cutlass.range_constexpr(self.threads // 32 - 1): + if warp > w: + start += totals[w] + if tid == self.threads - 1: + scalars[0] = start + count + for j in cutlass.range_constexpr(span): + e = span * tid + j + if e < self.experts: + base[e] = start + for v in range(units[j]): + experts[start + v] = e + start += units[j] + cute.arch.sync_threads() + for k in cutlass.range_constexpr(rounds): + pair = tid + self.threads * k + if pair < batch * 16: + e, token = (mine[k], pair >> 4) + low, word = (masks[2 * e], masks[2 * e + 1]) + rank = Int32(0) + if token < 32: + word = low + else: + rank = Int32(cute.arch.popc(low)) + below = (Uint32(1) << (token & 31)) - 1 + rank += Int32(cute.arch.popc(word & below)) + rows[(base[e] + (rank >> 3)) * 8 + (rank & 7)] = Int16(pair) + cute.arch.sync_threads() + + @cute.jit + def fetch_windows(self, sm, workspace): + """Pregrouped: the misc scalars from the unit table's header (every CTA's call + counters are equal between calls), then thread w * 16 + i copies unit i of + window w (FC1 A, FC1 B, FC2 A, FC2 B as ``partition``) out of the table. One + compact block (this code runs cold at kernel start).""" + tid, _, _, cta = thread() + scalars = view(sm.iterator, self.off_misc + 1140, Int32, 3) + table = workspace.iterator + self.ws_group + if tid == 0: + header = load_words(table, count=4, mutable=True) + for i in cutlass.range_constexpr(3): + scalars[i] = Int32(header[i]) + cute.arch.sync_threads() + if tid < 64: + units, fc2, group = (scalars[0], tid >= 32, tid // 16 % 2 == 1) + middle = units + if units > 40: + middle = (units * 50 + 99) // 100 + start, count = (Int32(0), middle) + if group: + start, count = (middle, units - middle) + per = self.tiles * 14 + lo = cta * per * count // self.grid + hi = (cta + 1) * per * count // self.grid + first, last = (start + lo // per, start + (hi - 1) // per) + if fc2: + lo = cta * self.jobs * count // self.grid + hi = (cta + 1) * self.jobs * count // self.grid + if cutlass.const_expr(self.steal): + stop = cta + self.victims + 1 + if stop > self.grid: + stop = Int32(self.grid) + if group: + hi = stop * self.jobs * count // self.grid + first, last = (start + lo // self.jobs, start + (hi - 1) // self.jobs) + unit = first + tid % 16 + if (lo < hi) & (unit <= last): + expert = load_words(table + 16 + 4 * unit, count=1, mutable=True) + row = load_words(table + 16 + 4096 + 16 * unit, count=4, mutable=True) + window_experts, _, window_words = self.tables(sm) + window_experts[tid] = Int32(expert[0]) + for i in cutlass.range_constexpr(4): + window_words[4 * tid + i] = row[i] + + @cute.jit + def gather_x(self, sm, plan, ranges, batch): + """Warp 3 (grouped): cp.async each FC1 stage's 8-row MXFP8 slice and scale + rows out of the plan (complete before the dependency wait returns) into its + slot; full[slot] counts the 32 lanes besides the producer (which stands in for + them on FC2 stages and the end). Padded rows repeat the unit's first token + (their outputs are dropped).""" + lane, ws = (thread()[2], plan.iterator) + stages, rows = (self.stages, self.tables(sm)[1]) + full, empty = ( + view(sm.iterator, self.off_misc + o, Uint64, stages) + for o in (self.bar_full, self.bar_empty) + ) + chunk, lanes = (lane & 15, lane // 16) + image = [self.off_x + swizzled(chunk, 2 * j + lanes) for j in range(4)] + scale = self.off_xsf + lane // 8 * 128 + (lane & 7) * 16 + sources = cute.make_rmem_tensor(5, Int32) + slot, phase, issued = (Int32(0), Int32(0), Int32(0)) + for group in cutlass.range_constexpr(2): + lo, hi = (ranges[group, 2], ranges[group, 3]) + ks, r, index = (lo % 14, lo // 14 % self.tiles, Int32(16 * group)) + if lo < hi: + self.row_sources(rows, index, sources, batch) + for s in cutlass.range(lo, hi, unroll=1): + if issued >= stages: + self.sync(empty.iterator + slot, phase ^ 1) + for j in cutlass.range_constexpr(4): + src = sources[j] + ks * 256 + chunk * 16 + copy_g2s(sm.iterator + image[j] + slot * 2048, ws + src) + if lane < 16: + src = PLAN_XSF + sources[4] + (2 * ks + lane // 8) * 16 + copy_g2s(sm.iterator + scale + slot * 256, ws + src) + cute.arch.cp_async_mbarrier_arrive_noinc(full.iterator + slot) + slot, issued = (slot + 1, issued + 1) + if slot == stages: + slot, phase = (Int32(0), phase ^ 1) + ks += 1 + if ks == 14: + ks, r = (Int32(0), r + 1) + if r == self.tiles: + r, index = (Int32(0), index + 1) + if s + 1 < hi: + self.row_sources(rows, index, sources, batch) + + @cute.jit + def row_sources(self, rows, index, sources, batch): + """Byte offsets of this lane's four MXFP8 rows and scale row of window unit + ``index`` (padded rows repeat its first token).""" + lane = thread()[2] + for j in cutlass.range_constexpr(5): + row = 2 * j + lane // 16 if j < 4 else lane & 7 + pair = Int32(rows[index * 8 + row]) + if pair < 0: + pair = Int32(rows[index * 8]) + sources[j] = (pair >> 4) * (3584 if j < 4 else 448) + + @cute.jit + def partition(self, distinct: Int32, cta: Int32): + """FC1(A), FC1(B), FC2(A), FC2(B) work partitions.""" + middle = distinct + if distinct > 40: + middle = (distinct * 50 + 99) // 100 + ranges = cute.make_rmem_tensor((2, 6), Int32) + for group in cutlass.range_constexpr(2): + start, end = (Int32(0), middle) if group == 0 else (middle, distinct) + count = end - start + size = self.tiles * 14 * count + ranges[group, 0], ranges[group, 1] = (start, size) + ranges[group, 2] = cta * size // self.grid + ranges[group, 3] = (cta + 1) * size // self.grid + n = self.jobs + ranges[group, 4] = n * start + cta * n * count // self.grid + ranges[group, 5] = n * start + (cta + 1) * n * count // self.grid + return ranges + + @cute.jit + def schedule(self, ranges, experts, n: Int32): + """Returns the stage's (kind, a, b, c) metadata and its TMA arguments.""" + n1a = ranges[0, 3] - ranges[0, 2] + n1b = ranges[1, 3] - ranges[1, 2] + n2a = ranges[0, 5] - ranges[0, 4] + kind, a, b, c = (Int32(0), Int32(0), Int32(0), Int32(0)) + row, chunk, scale_bytes, scale_offset = (Int32(0), Int32(0), Int32(0), Int64(0)) + if n < n1a + n1b: + group = Int32(n >= n1a) + s = ranges[group, 2] + (n - n1a if group == 1 else n) + tile = self.tiles * ranges[group, 0] + s // 14 + ks = s % 14 + expert, r = (experts[tile // self.tiles], tile % self.tiles) + a, b = (tile, ks) + opens = Int32((s == ranges[group, 2]) | (ks == 0)) + c = opens | Int32((s == ranges[group, 3] - 1) | (ks == 13)) << 1 + row, chunk = (expert * self.w13_rows + r * 128, 2 * ks) + scale_offset = Int64(expert) * self.w13_scale_bytes + r * 14336 + ks * 1024 + scale_bytes = Int32(1024) + else: + step = n - n1a - n1b + unit, half = (step // self.fc2_stages, step % self.fc2_stages) + group = Int32(unit >= n2a) + job = ranges[group, 4] + (unit - n2a if group == 1 else unit) + u, tile = (job // 28, job % 28) + expert = experts[u] + kind = Int32(1) if half == 0 else Int32(3) + a, b = (u, tile) + row, chunk = (expert * 3584 + tile * 128, 2 * half) + scale_offset = Int64(expert) * self.w2_scale_bytes + scale_offset += tile * self.w2_tile_scale_bytes + scale_offset += half * 1024 + scale_bytes = Int32(1024 if half == 0 else self.w2_tile_scale_bytes - 1024) + return (kind, a, b, c), (row, chunk, scale_offset, scale_bytes) + + @cute.jit + def poll_hidden(self, sm, source, buffer: Int32): + """Poll real K data and scale bytes, preserving TP16 padding semantics.""" + lane = thread()[2] + chunks: cutlass.Constexpr = self.h_image_bytes // 512 + values, scale = (cute.make_rmem_tensor((chunks, 4), Uint32), Uint32(0)) + valid_bytes = min(4, max(0, self.tiles * 2 - 4 * (lane & 3))) + mask = Uint32(4294967295) + if valid_bytes < 4: + mask = (Uint32(1) << 8 * valid_bytes) - 1 + ready = False + while not ready: + bad = False + for v in cutlass.range_constexpr(chunks): + index = lane + 32 * v + words = load_volatile(source + index * 16, count=4) + for w in cutlass.range_constexpr(4): + values[v, w] = words[w] + atom, token, position = (index >> 6, index >> 3 & 7, index & 7) + if 8 * atom + (position ^ token) < self.tiles * 4: + for w in cutlass.range_constexpr(4): + bad |= has_ff_byte(words[w]) + scale = load_volatile(source + self.h_image_bytes + lane * 4)[0] + if valid_bytes > 0: + bad |= has_ff_byte(scale & mask) + ready = cute.arch.vote_all_sync(not bad) + image = sm.iterator + self.off_h + buffer * self.h_image_bytes + for v in cutlass.range_constexpr(chunks): + words = tuple(values[v, w] for w in range(4)) + store_shared_words(image + (lane + 32 * v) * 16, words) + word = cute.arch.shuffle_sync(scale, 4 * (lane & 7) + (lane >> 3)) + if lane < 8 * self.h_chunks: + sf = sm.iterator + self.off_hsf + buffer * self.h_scale_bytes + store_shared_words(sf + (lane >> 3) * 128 + (lane & 7) * 16, (word,) * 4) + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_warp() + + @cute.jit + def load_hidden(self, sm, workspace, ranges, parity): + lane = thread()[2] + full = view(sm.iterator, self.off_misc + self.bar_hfull, Uint64, 2) + empty = view(sm.iterator, self.off_misc + self.bar_hempty, Uint64, 2) + ordinal = Int32(-1) + if cutlass.const_expr(self.steal): + unit = self.pop_unit(sm, Int32(0)) + while unit >= 0: + ordinal += 1 + buffer = ordinal & 1 + if ordinal >= 2: + self.sync(empty.iterator + buffer, (ordinal >> 1) - 1 & 1) + cute.arch.fence_proxy("async.shared", space="cta") # MMA read it + image = self.image(workspace, parity, unit) + self.poll_hidden(sm, image, buffer) + if lane == 0: + cute.arch.mbarrier_arrive(full.iterator + buffer) + unit = self.pop_unit(sm, ordinal + 1) + for g in cutlass.range_constexpr(2 * (not self.steal)): + if ranges[g, 5] > ranges[g, 4]: + lo, hi = (ranges[g, 4] // self.jobs, (ranges[g, 5] - 1) // self.jobs) + for u in range(lo, hi + 1): + ordinal += 1 + buffer = ordinal & 1 + if ordinal >= 2: + self.sync(empty.iterator + buffer, (ordinal >> 1) - 1 & 1) + cute.arch.fence_proxy("async.shared", space="cta") + image = self.image(workspace, parity, u) + self.poll_hidden(sm, image, buffer) + if lane == 0: + cute.arch.mbarrier_arrive(full.iterator + buffer) + + @cute.jit + def queue(self, sm): + """The unit queue: (units [8], full [8], empty [8] barriers).""" + bars = self.off_misc + self.bar_queue + units = view(sm.iterator, self.off_queue, Int32, 8) + full = view(sm.iterator, bars, Uint64, 8) + return units, full, view(sm.iterator, bars + 64, Uint64, 8) + + @cute.jit + def push_unit(self, sm, pushed, unit): + """Producer: queue ``unit`` (-1 ends) for the hidden loader.""" + units, full, empty = self.queue(sm) + slot = pushed % 8 + if pushed >= 8: + self.sync(empty.iterator + slot, (pushed // 8 - 1) & 1) + units[slot] = unit + cute.arch.mbarrier_arrive(full.iterator + slot) + return pushed + 1 + + @cute.jit + def pop_unit(self, sm, popped): + """Hidden loader warp: the ``popped``-th queued unit.""" + units, full, empty = self.queue(sm) + slot = popped % 8 + self.sync(full.iterator + slot, popped // 8 & 1) + unit = units[slot] + cute.arch.sync_warp() + if thread()[2] == 0: + cute.arch.mbarrier_arrive(empty.iterator + slot) + return unit + + @cute.jit + def fc2_job(self, ring, state, maps, experts, ranges, j, group: cutlass.Constexpr): + """FC2 job ``j`` (absolute) of window ``group``: queue a new unit, then issue + its stages (TP8 a tile pair, TP16 a tile) exactly as the static walk does.""" + slot, phase, issued, pushed, current = state + unit, job = (j // self.jobs, j % self.jobs) + if unit != current: + pushed = self.push_unit(ring[0], pushed, unit) + current = unit + index = 32 + 16 * group + unit - self.first_unit(ranges, 1, group) + expert = experts[index] + if cutlass.const_expr(self.paired): + work = (maps.iterator, expert, unit, 2 * job) + slot, phase, issued = self.fc2_pair(ring, slot, phase, issued, *work) + else: + row = expert * 3584 + job * 128 + offset = Int64(expert) * self.w2_scale_bytes + offset += job * self.w2_tile_scale_bytes + tma = (maps.iterator + 128, ring[3], row, Int32(0), offset, Int32(1024)) + meta = (Int32(1), unit, job, Int32(0)) + slot, phase, issued = self.issue(ring, slot, phase, issued, meta, tma, True) + return slot, phase, issued, pushed, current + + @cute.jit + def fc2_jobs(self, ring, state, maps, experts, ranges, lo, hi, group): + for j in cutlass.range(lo, hi, unroll=1): + state = self.fc2_job(ring, state, maps, experts, ranges, j, group) + return state + + @cute.jit + def claims(self, ring, state, maps, experts, ranges, counter, base, grains, back): + """Grains of one zone (``grains`` from job ``base``) claimed through + ``counter`` (baseline -1: low half front claims, high half back claims) from + the front or the ``back``; the next claim is in flight while one issues. + Returns the state and the grains taken.""" + step = (1 << 16) if back else 1 + pending = claim(counter, step) + active, taken = (grains > 0, Int32(0)) + while active: + value = pending + 1 + front, rear = (Int32(value & 65535), Int32(value >> 16)) + active = front + rear < grains + if active: + ahead = claim(counter, step) + index = front + if back: + index = grains - 1 - rear + lo = base + index * self.grain + hi = lo + self.grain + state = self.fc2_jobs(ring, state, maps, experts, ranges, lo, hi, 1) + pending, taken = (ahead, taken + 1) + return state, taken + + @cute.jit + def fc2_steal(self, ring, state, maps, experts, ranges, workspace, parity): + """FC2 with tail stealing: FC2(A) and the static head of FC2(B), the owner's + zone from the front, then the zones of CTAs cta + 1 .. cta + victims from the + back while their units fit window 3; the queue ends with -1.""" + cta = thread()[3] + n, g = (self.jobs, self.grain) + head = (ranges[0, 4], ranges[0, 5], 0) + state = self.fc2_jobs(ring, state, maps, experts, ranges, *head) + counters = workspace.iterator + self.ws_steal + 640 * parity + lo, hi = (ranges[1, 4], ranges[1, 5]) + grains = (hi - lo) // g + if grains > self.zone: + grains = Int32(self.zone) + base = hi - grains * g + state = self.fc2_jobs(ring, state, maps, experts, ranges, lo, base, 1) + own = (counters + 4 * cta, base, grains, False) + state, _ = self.claims(ring, state, maps, experts, ranges, *own) + start, count = (ranges[1, 0], ranges[1, 1] // (self.tiles * 14)) + first = ranges[1, 4] // n + victim = cta + 1 + while (victim <= cta + self.victims) & (victim < self.grid): + lo = n * start + victim * n * count // self.grid + hi = n * start + (victim + 1) * n * count // self.grid + grains = (hi - lo) // g + if grains > self.zone: + grains = Int32(self.zone) + if (hi - 1) // n - first > 15: + victim = cta + self.victims + 1 # beyond window 3: stop + else: + counter = counters + 4 * victim + work = (counter, hi - grains * g, grains, True) + state, taken = self.claims(ring, state, maps, experts, ranges, *work) + victim = victim + 1 + slot, phase, issued, pushed, current = state + pushed = self.push_unit(ring[0], pushed, Int32(-1)) + return slot, phase, issued + + @cute.jit + def fetch(self, sm, desc, scales, full, slot, row, chunk, offset, size, more=None): + """Weight box and scale bytes (``more``: a second ``size``-byte scale run).""" + load_tma3(sm.iterator + slot * 32768, desc, full.iterator + slot, row, chunk) + dst = sm.iterator + self.off_sfa + slot * 1024 + copy_evict(dst, scales.iterator + offset, size, full.iterator + slot) + if cutlass.const_expr(more is not None): + copy_evict(dst + 512, scales.iterator + more, size, full.iterator + slot) + + @cute.jit + def produce(self, sm, maps, w13_scales, w2_scales, ranges, experts): + """Producer lane; dynamic scale bytes: 1024, or 512 for the TP8 FC2 tail.""" + stages = self.stages + barriers = view(sm.iterator, self.off_misc, Uint64, 2 * stages) + metadata = view(sm.iterator, self.off_meta, Int32, 4 * stages) + fc1 = ranges[0, 3] - ranges[0, 2] + ranges[1, 3] - ranges[1, 2] + fc2 = ranges[0, 5] - ranges[0, 4] + ranges[1, 5] - ranges[1, 4] + total = fc1 + self.fc2_stages * fc2 + for n in cutlass.range(total, unroll=1): + slot, phase = (n % stages, n // stages % 2) + meta, tma = self.schedule(ranges, experts, n) + if n >= stages: + empty = barriers.iterator + stages + slot + self.sync(empty, phase ^ 1) + for i in cutlass.range_constexpr(4): + metadata[4 * slot + i] = meta[i] + cute.arch.mbarrier_arrive_and_expect_tx( + barriers.iterator + slot, 16384 + tma[3] + ) + if meta[0] == 0: + self.fetch(sm, maps.iterator, w13_scales, barriers, slot, *tma) + else: + self.fetch(sm, maps.iterator + 128, w2_scales, barriers, slot, *tma) + slot, phase = (total % stages, total // stages % 2) + if total >= stages: + self.sync(barriers.iterator + stages + slot, phase ^ 1) + for i in cutlass.range_constexpr(4): + metadata[4 * slot + i] = Int32(2 if i == 0 else 0) + cute.arch.mbarrier_arrive(barriers.iterator + slot) + + @cute.jit + def stream(self, sm, maps, w13_scales, w2_scales, ranges, experts, steal=None): + """Producer lane walking its FC1 then FC2 stages incrementally (no per-stage + divisions; the expert is re-read only when the unit changes). Grouped expert + tables are this CTA's windows; else they hold every distinct expert.""" + tiles, w13, w2 = (self.tiles, maps.iterator, maps.iterator + 128) + barriers = view(sm.iterator, self.off_misc, Uint64, 2 * self.stages) + metadata = view(sm.iterator, self.off_meta, Int32, 4 * self.stages) + ring = (sm, barriers, metadata, w2_scales) + slot, phase, issued = (Int32(0), Int32(0), Int32(0)) + for group in cutlass.range_constexpr(2): + lo, hi = (ranges[group, 2], ranges[group, 3]) + ks, tile = (lo % 14, tiles * ranges[group, 0] + lo // 14) + r, index = (tile % tiles, self.first_unit(ranges, 0, group)) + if cutlass.const_expr(self.grouped): + index = Int32(16 * group) + expert = experts[index] + for s in cutlass.range(lo, hi, unroll=1): + opens, closes = ((s == lo) | (ks == 0), (s == hi - 1) | (ks == 13)) + c = Int32(opens) | Int32(closes) << 1 + row = expert * self.w13_rows + r * 128 + offset = Int64(expert) * self.w13_scale_bytes + r * 14336 + ks * 1024 + tma = (w13, w13_scales, row, 2 * ks, offset, Int32(1024)) + slot, phase, issued = self.issue( + ring, slot, phase, issued, (Int32(0), tile, ks, c), tma, False + ) + ks += 1 + if ks == 14: + ks, tile, r = (Int32(0), tile + 1, r + 1) + if r == tiles: + r, index = (Int32(0), index + 1) + expert = experts[index] + if cutlass.const_expr(self.steal): + state = (slot, phase, issued, Int32(0), Int32(-1)) + work = (ring, state, maps, experts, ranges, *steal) + slot, phase, issued = self.fc2_steal(*work) + for group in cutlass.range_constexpr(2 * (not self.steal)): + lo, hi = (ranges[group, 4], ranges[group, 5]) + unit, tile, index = (lo // self.jobs, lo % self.jobs, lo // self.jobs) + if cutlass.const_expr(self.grouped): + index = Int32(32 + 16 * group) + expert = experts[index] + for _ in cutlass.range(lo, hi, unroll=1): + if cutlass.const_expr(self.paired): + job = (maps.iterator, expert, unit, 2 * tile) + slot, phase, issued = self.fc2_pair(ring, slot, phase, issued, *job) + else: + for half in cutlass.range_constexpr(self.fc2_stages): + size = 1024 if half == 0 else self.w2_tile_scale_bytes - 1024 + row = expert * 3584 + tile * 128 + offset = Int64(expert) * self.w2_scale_bytes + half * 1024 + offset += tile * self.w2_tile_scale_bytes + tma = (w2, w2_scales, row, 2 * half, offset, Int32(size)) + meta = (Int32(1 if half == 0 else 3), unit, tile, Int32(0)) + slot, phase, issued = self.issue( + ring, slot, phase, issued, meta, tma, True + ) + tile += 1 + if tile == self.jobs: + unit, tile, index = (unit + 1, Int32(0), index + 1) + expert = experts[index] + if issued >= self.stages: + self.sync(barriers.iterator + self.stages + slot, phase ^ 1) + for i in cutlass.range_constexpr(4): + metadata[4 * slot + i] = Int32(2 if i == 0 else 0) + arrivals = 33 if self.grouped else 1 + cute.arch.mbarrier_arrive(barriers.iterator + slot, arrive_count=arrivals) + + @cute.jit + def fc2_pair(self, ring, slot, phase, issued, maps, expert, unit, tile): + """The three 16 KB stages of FC2 tiles ``tile`` and ``tile + 1``: chunks 0-1 of + each (1 KB scales), then chunk 2 of both (map 5, two 512 B scale runs).""" + w2, scales, step = (maps + 128, ring[3], self.w2_tile_scale_bytes) + row = expert * 3584 + tile * 128 + offset = Int64(expert) * self.w2_scale_bytes + tile * step + for t in cutlass.range_constexpr(2): + tma = (w2, scales, row + 128 * t, Int32(0), offset + t * step, Int32(1024)) + meta = (Int32(1), unit, tile + t, Int32(0)) + slot, phase, issued = self.issue(ring, slot, phase, issued, meta, tma, True) + tma = (maps + 640, scales, row, Int32(2), offset + 1024, Int32(512)) + meta = (Int32(3), unit, tile, Int32(0)) + more = offset + step + 1024 + return self.issue(ring, slot, phase, issued, meta, tma, True, more) + + @cute.jit + def issue( + self, ring, slot, phase, issued, meta, tma, fc2: cutlass.Constexpr, more=None + ): + """One grouped stage: wait for its slot, publish the metadata, expect and issue + the weight box and scales (FC2: standing in for the gather warp's 32 + arrivals); returns the advanced (slot, phase, issued).""" + sm, barriers, metadata = ring[:3] + desc, scales, row, chunk, offset, size = tma + if issued >= self.stages: + empty = barriers.iterator + self.stages + slot + self.sync(empty, phase ^ 1) + for i in cutlass.range_constexpr(4): + metadata[4 * slot + i] = meta[i] + full = barriers.iterator + slot + total = 16384 + size if more is None else 16384 + 2 * size + cute.arch.mbarrier_arrive_and_expect_tx(full, total) + self.fetch(sm, desc, scales, barriers, slot, row, chunk, offset, size, more) + if cutlass.const_expr(fc2 and self.grouped): + cute.arch.mbarrier_arrive(full, arrive_count=32) + slot += 1 + if slot == self.stages: + slot, phase = (Int32(0), phase ^ 1) + return slot, phase, issued + 1 + + @cute.jit + def ownership(self, index: cutlass.Constexpr[int], batch): + xtid = thread()[0] - 64 + if cutlass.const_expr(self.dense): + block = xtid + 320 * index + q = block // 8 + stage = q // batch + token = q - stage * batch + kblock = block & 7 + valid = block < 112 * batch + else: + stage = 5 * index + xtid // 64 + token = (xtid & 63) // 8 + kblock = xtid & 7 + valid = (stage < 14) & (token < batch) + return (stage, token, kblock, valid) + + @cute.jit + def load_x(self, x, raw, batch): + """Warps 2..11 preload their 32-BF16 x blocks before the dependent launch.""" + raw.fill(0) + for index in cutlass.range_constexpr(self.rounds): + stage, token, kblock, valid = self.ownership(index, batch) + if valid: + source = x.iterator + token * 3584 + stage * 256 + kblock * 32 + for fragment in cutlass.range_constexpr(4): + words = load_words( + source + fragment * 8, count=4, cache_policy="", mutable=True + ) + for j in cutlass.range_constexpr(4): + raw[index, fragment, j] = words[j] + + @cute.jit + def quantize_x(self, sm, raw, barriers, batch): + """Each thread quantizes its block (two uint4 fragments, one scale byte) and + arrives once at xready[stage]: eight arrivals per active token, no DSMEM.""" + for index in cutlass.range_constexpr(self.rounds): + stage, token, kblock, valid = self.ownership(index, batch) + if valid: + packed, scale = mx_block([raw[index, i // 4, i % 4] for i in range(16)]) + for half, words in enumerate((packed[:4], packed[4:])): + offset = swizzled(16 * stage + 2 * kblock + half, token, self.off_x) + store_shared_words(sm.iterator + offset, words) + kb = 8 * stage + kblock + sm[self.off_xsf + kb // 4 * 128 + token * 16 + kb % 4] = Uint8(scale) + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.mbarrier_arrive(barriers + stage) + + @cute.jit + def acquire(self, empty, acc: Int32, uses: Int32): + if uses >= 16: + self.sync(empty.iterator + acc, uses // 16 - 1 & 1, True) + + @cute.jit + def signal(self, metadata, full, acc, values, arrivals: cutlass.Constexpr = 1): + if thread()[2] == 0: + for i in cutlass.range_constexpr(4): + metadata[4 * acc + i] = values[i] + for _ in cutlass.range_constexpr(arrivals): + cute.arch.mbarrier_arrive(full.iterator + acc) + + @cute.jit + def consume(self, sm, tmem): + lane = thread()[2] + stages = self.stages + bars: tuple = (self.bar_full, self.bar_empty, self.bar_accfull) + bars += (self.bar_accempty, self.bar_hfull, self.bar_hempty, self.bar_xready) + full, empty, accfull, accempty, hfull, hempty, xready = ( + view(sm.iterator, self.off_misc + o, Uint64, n) + for o, n in zip(bars, (stages, stages, 16, 16, 2, 2, 14)) + ) + meta = view(sm.iterator, self.off_meta, Int32, 4 * stages) + accmeta = view(sm.iterator, self.off_misc + 600, Int32, 64) + offsets = (0, self.off_sfa, self.off_x, self.off_xsf, self.off_h, self.off_hsf) + a0, asf0, x0, xsf0, h0, hsf0 = ( + smem_desc(sm.iterator + o, i % 2 == 0) for i, o in enumerate(offsets) + ) + xmask, ended = (Uint32(0), False) + slot, uses, acc, count = (Int32(0), Int32(0), Int32(0), Int32(0)) + current, ordinal, hbuffer = (Int32(-1), Int32(-1), Int32(0)) + while not ended: + slot = count % stages + ready = (full.iterator + slot, count // stages & 1, True) + self.sync(*ready) + kind, a, b, c = (meta[4 * slot + i] for i in range(4)) + if kind == 2: + if (lane == 0) & (current >= 0) & (hbuffer < 2): + commit(hempty.iterator + hbuffer) + for _ in cutlass.range(2, unroll=1): + self.acquire(accempty, acc, uses) + self.signal(accmeta, accfull, acc, (kind, a, b, c), 2) + uses += 1 + acc = uses % 16 + ended = True + else: + publish = False + if kind == 0: + first, last = (c & 1 != 0, c & 2 != 0) + if first: + self.acquire(accempty, acc, uses) + copy_scales = Uint32(0) + if cutlass.const_expr(self.grouped): + # The slot's own image (its cp.async completed full[slot]). + cute.arch.fence_proxy("async.shared", space="cta") + copy_scales = Uint32(1) + elif xmask >> b & 1 == 0: + self.sync(xready.iterator + b, 0, True) + xmask |= Uint32(1) << b + copy_scales = Uint32(1) + completion = Uint32(0) + if last: + completion = shared_address(accfull.iterator + acc) + xi = slot if self.grouped else b + stage( + (tmem + acc * 8, a0 + slot * 2048, x0 + xi * 128), + (tmem + 128 + slot * 8, asf0 + slot * 64), + (tmem + 256 + xi * 8, xsf0 + xi * 16, copy_scales), + (Uint32(not first), shared_address(empty.iterator + slot)), + completion, + kind="fc1", + ) + cute.arch.sync_warp() + publish = last + else: + copy_scales, release = (Uint32(0), Uint32(0)) + if a != current: + if (current >= 0) & (hbuffer < 2): + release = shared_address(hempty.iterator + hbuffer) + current = a + ordinal += 1 + hbuffer = ordinal & 1 + ready = (hfull.iterator + hbuffer, ordinal >> 1 & 1, True) + self.sync(*ready) + copy_scales = Uint32(1) + hd = h0 + hbuffer * (self.h_image_bytes // 16) + hs = hsf0 + hbuffer * (self.h_scale_bytes // 16) + sfbc = tmem + 128 + 8 * stages + hbuffer * self.w2_scale_columns + tp16 = cutlass.const_expr(self.tp == 16) + if cutlass.const_expr(self.paired): + weights = (tmem + 128 + slot * 8, asf0 + slot * 64) + released = shared_address(empty.iterator + slot) + if kind == 1: # chunks 0-1 of the pair's tile b + target, owner = (acc, uses) + if b % 2 == 1: + target, owner = ((uses + 1) % 16, uses + 1) + self.acquire(accempty, target, owner) + stage( + (tmem + target * 8, a0 + slot * 2048, hd), + weights, + (sfbc, hs, copy_scales), + (release, released), + 0, + kind="fc2a", + ) + else: # chunk 2 of tiles b (rows 0-127) and b + 1 (128-255) + for t in cutlass.range_constexpr(2): + rows = a0 + slot * 2048 + 1024 * t + stage( + (tmem + acc * 8, rows, hd + 128), + (weights[0] + 4 * t, weights[1] + 32 * t), + (sfbc + 8, 0, 0), + (0, released), + shared_address(accfull.iterator + acc), + kind="fc2b" if t else "fc2p", + ) + if cutlass.const_expr(t == 0): + cute.arch.sync_warp() + first = (Int32(1), a, b, c) + self.signal(accmeta, accfull, acc, first) + uses += 1 + acc = uses % 16 + cute.arch.sync_warp() + publish, kind, b = (True, Int32(1), b + 1) + else: + if tp16 or kind == 1: + self.acquire(accempty, acc, uses) + stage( + (tmem + acc * 8, a0 + slot * 2048, hd), + (tmem + 128 + slot * 8, asf0 + slot * 64), + (sfbc, hs, copy_scales), + (release, shared_address(empty.iterator + slot)), + shared_address(accfull.iterator + acc) if tp16 else 0, + kind="fc2" if tp16 else "fc2a", + ) + publish = tp16 + else: + stage( + (tmem + acc * 8, a0 + slot * 2048, hd + 128), + (tmem + 128 + slot * 8, asf0 + slot * 64), + (sfbc + 8, 0, 0), + (0, shared_address(empty.iterator + slot)), + shared_address(accfull.iterator + acc), + kind="fc2b", + ) + publish = True + kind = Int32(1) + cute.arch.sync_warp() + if publish: + self.signal(accmeta, accfull, acc, (kind, a, b, c)) + uses += 1 + acc = uses % 16 + count += 1 + + @cute.jit + def partials(self, workspace, values, tile, ranges, batch): + """Publish non-owner fragments, or reduce the owner's remaining pieces.""" + ws, live = (workspace.iterator, 8 if self.grouped else batch) + _, warp, lane, cta = thread() + row = 32 * (warp & 3) + lane + group = Int32(tile >= self.tiles * ranges[1, 0]) + start = (tile - self.tiles * ranges[group, 0]) * 14 + owner = start >= ranges[group, 2] + if not owner: + offset = self.ws_partial + ((Int64(group) * 160 + cta) * 1024 + row) * 4 + for t in cutlass.range_constexpr(self.width): + if t < live: + store_words(ws + offset + t * 512, (float_bits(values[t]),)) + else: + peer, size = (cta + 1, ranges[group, 1]) + while (peer < self.grid) & (peer * size // self.grid < start + 14): + if (peer + 1) * size // self.grid != peer * size // self.grid: + source = ws + self.ws_partial + source += ((Int64(group) * 160 + peer) * 1024 + row) * 4 + words = cute.make_rmem_tensor(self.width, Uint32) + ready = False + while not ready: + bad = False + for t in cutlass.range_constexpr(self.width): + words[t] = Uint32(0) + if t < live: + words[t] = load_volatile(source + t * 512)[0] + bad |= words[t] == Uint32(4294967295) + ready = cute.arch.vote_all_sync(not bad) + for t in cutlass.range_constexpr(self.width): + if t < live: + values[t] += bits_float(words[t]) + store_words(source + t * 512, (Uint32(4294967295),)) + peer += 1 + return owner + + @cute.jit + def store_hidden(self, sm, workspace, values, tile, parity, scalars): + _, warp, lane, _ = thread() + ew, eg = (warp & 3, (warp - 4) // 4) + expert, chunk = (tile // self.tiles, tile % self.tiles) + up = lane & 8 == 0 + h, maxima = (cute.make_rmem_tensor(self.width, Float32) for _ in range(2)) + shared_max = view(sm.iterator, self.off_misc + 884, Float32, 64) + for t in cutlass.range_constexpr(self.width): + gate = cute.arch.shuffle_sync_bfly(values[t], 8) + h[t] = situ(gate, values[t], *scalars) if up else Float32(0) + maxima[t] = abs(h[t]) + for t in cutlass.range_constexpr(self.width): + for off in cutlass.range_constexpr(5): + other = cute.arch.shuffle_sync_bfly(maxima[t], 1 << off) + maxima[t] = cute.arch.fmax(maxima[t], other) + if lane == 0: + shared_max[(eg * 4 + ew) * 8 + t] = maxima[t] + group_barrier(eg) + stage_base, scale_token = (self.off_misc + 3200 + (eg * 4 + ew) * 128, Int32(0)) + for t in cutlass.range_constexpr(self.width): + peer = shared_max[(eg * 4 + (ew ^ 1)) * 8 + t] + scale = mx_exponent(cute.arch.fmax(maxima[t], peer), floor=self.tp == 8) + if lane == t: + scale_token = scale + if up: + bits = fp8_bits(h[t] * mx_rescale(scale)) + sm[stage_base + t * 16 + 2 * (lane % 8) + lane // 16] = bits + cute.arch.sync_warp() + if lane < 8: + kc, data = (4 * chunk + ew, (Uint32(0), Uint32(0), Uint32(0), Uint32(0))) + if lane < self.width: + data = smem_words(sm.iterator + stage_base + lane * 16) + image = self.image(workspace, parity, expert) + store_words(image + swizzled(kc, lane), data) + if ew & 1 == 0: + byte = view(image, 0, Uint8, self.h_compact_bytes) + value = Uint8(scale_token) + byte[self.h_image_bytes + lane * 16 + 2 * chunk + ew // 2] = value + group_barrier(eg) + + def pair(self, t, slot): + """Expanded row of unit row ``t``: its pair (grouped), or token t's ``slot``.""" + return Int64(slot) if self.grouped else Int64(t) * 16 + slot + + @cute.jit + def fc2(self, sm, workspace, output, values, unit, tile, ranges): + """Rows of ``unit``: slots of tokens 0..7, or grouped pair indices.""" + _, warp, lane, _ = thread() + if cutlass.const_expr(self.grouped): + slots = self.tables(sm)[1] + unit = self.local(ranges, unit, 1, Int32(unit >= ranges[1, 0])) + else: + slots = view(sm.iterator, self.off_misc + 2176, Int8, 1024) + for t in cutlass.range_constexpr(self.width): + slot = Int32(slots[unit * 8 + t]) + if slot >= 0: + other = cute.arch.shuffle_sync_down(values[t], 8) + if cutlass.const_expr(self.publish): + # Stage row [token, slot]: lane q < 4 stores this warp's 8q..8q+7. + word, q = pack_pair(values[t], other, policy="stage"), lane & 3 + words = [shuffle(word, s + 2 * q) for s in (0, 16, 1, 17)] + row = self.pair(t, slot) * 1792 + tile * 64 + (warp & 3) * 16 + stage = workspace.iterator + self.ws_stage + (row + 4 * q) * 4 + if lane < 4: + store_words(stage, words) + elif lane & 8 == 0: + column = tile * 128 + (warp & 3) * 32 + 4 * (lane % 8) + lane // 8 + dst = output.iterator + self.pair(t, slot) * 3584 + column + store_words(dst, (bf16x2(values[t], other),)) + + @cute.jit + def finalize(self, workspace, weights, batch, box): + """Each epilogue warp polls (and re-arms) the 16 staged rows of its 64-column + lines and multicasts their weighted sum into latent slot (parity, rank) of the + [3, TP, 64, 3584] mailbox (two-shot: pushes every 8 columns to their owner); + lines stride over every CTA's epilogue warps.""" + _, warp, lane, cta = thread() + for line in range(cta * 8 + warp - 4, batch * 56, self.grid * 8): + token, column = line // 56, line % 56 * 64 + source = workspace.iterator + self.ws_stage + source += (Int64(token) * 16 * 1792 + column // 2 + lane) * 4 + words = cute.make_rmem_tensor(16, Uint32) + ready = False + while not ready: + bad = False + for slot in cutlass.range_constexpr(16): + words[slot] = load_volatile(source + slot * 7168)[0] + bad |= words[slot] == Uint32(4294967295) + ready = cute.arch.vote_all_sync(not bad) + low, high = Float32(0), Float32(0) + for slot in cutlass.range_constexpr(16): + store_words(source + slot * 7168, (Uint32(4294967295),)) + weight = Float32(weights[token * 16 + slot]) + low = fma(bf16(words[slot]), weight, low) + high = fma(bf16(words[slot], True), weight, high) + word = pack_pair(low, high, policy="latent") + packed = [shuffle(word, 4 * (lane & 7) + i) for i in range(4)] + if lane < 8: + if cutlass.const_expr(self.publish == "twoshot"): + push(*box, token, column + 8 * lane, packed) + else: + mc, parity, rank = box + offset = ((parity * self.tp + rank) * 64 + token) * 3584 + column + offset += 8 * lane + multicast_words(mc + Int64(offset) * 2, packed) + + @cute.jit + def epilogue(self, sm, workspace, output, tmem, ranges, batch, parity, scalars): + _, warp, lane, _ = thread() + ew, eg = (warp & 3, (warp - 4) // 4) + full = view(sm.iterator, self.off_misc + self.bar_accfull, Uint64, 16) + empty = view(sm.iterator, self.off_misc + self.bar_accempty, Uint64, 16) + meta = view(sm.iterator, self.off_misc + 600, Int32, 64) + acc, phase, ended = (eg, Int32(0), False) + while not ended: + self.sync(full.iterator + acc, phase, True) + kind, a, b = (meta[4 * acc], meta[4 * acc + 1], meta[4 * acc + 2]) + if kind == 2: + ended = True + else: + raw = load_tmem(tmem + (Uint32(32 * ew) << 16) + acc * 8, 8) + values = cute.make_rmem_tensor(self.width, Float32) + for t in cutlass.range_constexpr(self.width): + values[t] = raw[t] + fence_tmem(before=True) + cute.arch.sync_warp() + if lane == 0: + cute.arch.mbarrier_arrive(empty.iterator + acc) + acc += 2 + if acc >= 16: + acc -= 16 + phase ^= 1 + if kind == 0: + if self.partials(workspace, values, a, ranges, batch): + self.store_hidden(sm, workspace, values, a, parity, scalars[:3]) + else: + self.fc2(sm, workspace, output, values, a, b, ranges) + if cutlass.const_expr(self.publish): + self.finalize(workspace, output, batch, scalars[3:]) + + @cute.jit + def __call__(self, *args): + launch(self.kernel, args, self.grid, 384, cluster=1) + + @cute.jit + def initialize(self, sm, maps, batch): + tid, _, _ = cute.arch.thread_idx() + barriers = view(sm.iterator, self.off_misc, Uint64, 65) + tmem_base = view(sm.iterator, self.off_misc + 856, Uint32, 1) + if tid == 0: + receive = range(528, 533) if self.capacity == 4 else () + for i in (*range(65), *receive): + acc = 2 * self.stages + count = 1 if not acc <= i < acc + 32 else 2 if i < acc + 16 else 4 + count = batch * 8 if acc + 36 <= i < acc + 50 else count + if self.grouped and i < self.stages: + count = 33 + cute.arch.mbarrier_init(barriers.iterator + i, count) + if cutlass.const_expr(self.steal): + queue = view(sm.iterator, self.off_misc + self.bar_queue, Uint64, 16) + for i in cutlass.range_constexpr(16): + cute.arch.mbarrier_init(queue.iterator + i, 1) + cute.arch.mbarrier_init_fence() + acquire_tensormap(maps.iterator) + acquire_tensormap(maps.iterator + 128) + if cutlass.const_expr(self.paired): + acquire_tensormap(maps.iterator + 640) + if tid // 32 == 3: + cute.arch.alloc_tmem(512, tmem_base.iterator) + cute.arch.relinquish_tmem_alloc_permit() + if cutlass.const_expr(not self.grouped): + self.clear_x(sm, tid, batch) + cute.arch.fence_proxy("async.shared", space="cta") + return tmem_base + + @cute.jit + def clear_x(self, sm, tid, batch): + """Zero the padded token rows and the scale image (capacity <= 8).""" + zero = (Uint32(0), Uint32(0), Uint32(0), Uint32(0)) + for i in range(tid, (8 - batch) * 224, 384): + token, chunk = (batch + i // 224, i % 224) + store_shared_words(sm.iterator + swizzled(chunk, token, self.off_x), zero) + xsf = sm.iterator + self.off_xsf + for i in range(tid, 224, 384): + store_shared_words(xsf + (i >> 3) * 128 + (i & 7) * 16, zero) + + @cute.jit + def rearm(self, workspace, call_count): + """Re-arm the other parity's images: all 128, or the last grouped call's.""" + _, _, lane, cta = thread() + items: cutlass.Constexpr = 128 * self.h_compact_bytes // 16 + lo, hi = (cta * items // self.grid, (cta + 1) * items // self.grid) + if cutlass.const_expr(self.grouped): + last = view(workspace.iterator, self.ws_calls + 640, Int32, 2) + used = max(last[call_count & 1 ^ 1], 0) * (self.h_compact_bytes // 16) + lo, hi = (cta * used // self.grid, (cta + 1) * used // self.grid) + other = self.image(workspace, call_count & 1 ^ 1, 0) + for i in range(lo + lane, hi, 32): + store_words(other + i * 16, (Uint32(4294967295),) * 4) + if cutlass.const_expr(self.steal): + # Every one of the 160 counters: grids change with the batch. + base = self.ws_steal + 640 * (call_count & 1 ^ 1) + counters = view(workspace.iterator, base, Int32, 160) + if lane == 0: + counters[cta] = Int32(-1) + if cta < 160 - self.grid: + counters[self.grid + cta] = Int32(-1) + + @cute.kernel + def kernel(self, args): + x, ids, maps, w13_scales, w2_scales, workspace, output, batch, *scalars = args + tid, warp, lane, cta = thread() + sm = alloc(cutlass.utils.SmemAllocator(), Uint8, self.smem_bytes, 1024) + tmem_base = self.initialize(sm, maps, batch) + cute.arch.griddepcontrol_wait() + counters = view(workspace.iterator, self.ws_counter, Int32, 160) + if cutlass.const_expr(self.grouped): + # Images and rearm alternate over grouped calls; ``x`` is the plan. + routed = view(sm.iterator, self.off_misc + 1140, Int32, 3) + if cutlass.const_expr(self.pregrouped): + self.fetch_windows(sm, workspace) + else: + self.group(sm, ids, workspace, batch) + self.windows(sm, self.partition(routed[0], cta)) + # Retire the routing scratch's generic accesses before weight-ring reuse. + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_threads() + distinct, call_count, generation = (routed[0], routed[1], routed[2]) + experts = self.tables(sm)[0] + else: + distinct, call_count, experts = self.route(sm, ids, counters, batch) + generation = call_count + layout = cute.make_layout((self.rounds, 4, 4), stride=(16, 4, 1)) + raw_x = cute.make_rmem_tensor(layout, Uint32) + if warp >= 2: + self.load_x(x, raw_x, batch) + fence_tmem() + tmem = tmem_base[0] + cute.arch.griddepcontrol_launch_dependents() + ranges = self.partition(distinct, cta) + if warp == 0: + if lane == 0: + weights = (sm, maps, w13_scales, w2_scales, ranges, experts) + if cutlass.const_expr(self.steal): + self.stream(*weights, (workspace, generation & 1)) + elif cutlass.const_expr(self.streamed): + self.stream(*weights) + else: + self.produce(*weights) + else: + if cutlass.const_expr(not self.grouped): + ready = view(sm.iterator, self.off_misc + self.bar_xready, Uint64, 14) + if warp >= 2: + self.quantize_x(sm, raw_x, ready.iterator, batch) + if warp == 1: + self.consume(sm, tmem) + elif warp == 2: + self.load_hidden(sm, workspace, ranges, generation & 1) + elif warp == 3: + if cutlass.const_expr(self.grouped): + self.gather_x(sm, x, ranges, batch) + self.rearm(workspace, generation) + else: + self.epilogue( + sm, workspace, output, tmem, ranges, batch, generation & 1, scalars + ) + release_tmem(tmem, 512, warp, 3) + if tid == 0: + counters[cta] = call_count + 1 + if cta < 160 - self.grid: + counters[self.grid + cta] = call_count + 1 + if cutlass.const_expr(self.grouped): + grouped = view(workspace.iterator, self.ws_calls, Int32, 162) + grouped[cta] = generation + 1 + if cta < 160 - self.grid: + grouped[self.grid + cta] = generation + 1 + if cta == 0: + grouped[160 + (generation & 1)] = distinct + + +class ExpertsGroupKernel(ExpertsKernel): + """One CTA groups the routes into the workspace unit table that pregrouped + ``ExpertsKernel`` read (launched after routing, beside the shared expert). Every + experts call advances all 160 call counters, so entry 0 stands for each CTA's.""" + + def __init__(self, *args, early=False): + super().__init__(*args) + self.threads = 1024 + # early: release the experts at entry, so their launch and prologue overlap + # this kernel's dependency wait and grouping (they still wait for its end). + self.early = early + + @cute.jit + def __call__(self, *args): + launch(self.kernel, args, 1, self.threads) + + @cute.kernel + def kernel(self, args): + ids, workspace, batch = args + tid = thread()[0] + sm = alloc(cutlass.utils.SmemAllocator(), Uint8, 32768 + 16, 1024) + units = view(sm.iterator, 32768, Int32, 3) + if cutlass.const_expr(self.early): + cute.arch.griddepcontrol_launch_dependents() + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(not self.early): + cute.arch.griddepcontrol_launch_dependents() + self.group(sm, ids, workspace, batch, units) + experts, _, words = self.tables(sm, False) + table = workspace.iterator + self.ws_group + out_experts = view(table, 16, Int32, 1024) + for i in range(tid, units[0], self.threads): + out_experts[i] = experts[i] + row = tuple(words[4 * i + k] for k in range(4)) + store_words(table + 16 + 4096 + 16 * i, row) + if tid == 0: + offsets = (self.ws_counter, self.ws_calls) + calls = [Uint32(view(workspace.iterator, o, Int32, 1)[0]) for o in offsets] + store_words(table, (Uint32(units[0]), *calls, Uint32(0))) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts_runtime.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts_runtime.py new file mode 100644 index 0000000000..02b46523b1 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/experts_runtime.py @@ -0,0 +1,207 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Persistent routed experts for B<=64 (B>8: grouped units of <=8 tokens per expert) +over one shared core and model-owned workspace. The entry wait completes the +predecessor (incl. workspace rearm/counters) before shared state is read; output rows +keep the caller's top16 order.""" + +import functools + +import cuda.bindings.driver as cuda +import cutlass as c +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import make_fake_stream + +from ..common.host import check_residency, fake, stream, tensor_maps, to_cute +from ..common.policy import batch_table +from .experts import ( + PLAN_BYTES, + ExpertsGeometry, + ExpertsGroupKernel, + ExpertsKernel, + ExpertsPlanKernel, + fc2_pairs, +) + + +def tensor(value, align=16): + return to_cute(value.reshape(-1), align) + + +def streamed(batch): + """B<=8 on the incremental producer: ~1-1.5% at B4-8, neutral at B1-2, which keep + the previous producer.""" + return bool(batch_table("4-8:1").get(batch, 0)) + + +def pregrouped(tp, batch): + """B>8 experts read the unit table a one-CTA kernel builds after routing, beside + the fused shared expert (hidden at B9-32; at B64 its ~2.3 us after routing outlast + the shared expert).""" + table = batch_table(PREGROUP[tp]) + return batch > 8 and bool(table.get(batch, 0)) + + +PREGROUP = {8: "9-32:1", 16: "9-32:1"} +# With the shared join skipped (moe.prepared.JOIN), TP8 B16 releases at entry. +EARLY = {8: "9-16:1", 16: ""} + + +def early_group(tp, batch): + """Pregrouped batches whose unit-table kernel releases the experts at its entry.""" + table = batch_table(EARLY[tp]) + return pregrouped(tp, batch) and bool(table.get(batch, 0)) + + +def stealing(batch): + """B>8 experts balance the FC2(B) tail by claims.""" + return batch > 8 and bool(batch_table("9-64:1").get(batch, 0)) + + +@functools.cache +def compile_experts(tp, experts, batch, grid, device, publish=False, *knobs): + """Only code and geometry are cached; no owner tensors or descriptors. ``knobs``: + (B<=8 stream, pregrouped, steal).""" + defaults = (streamed(batch), False, False) + stream, grouped, steal = knobs + defaults[len(knobs) :] + capacity = 1 << (batch - 1).bit_length() + moe = ExpertsKernel(tp, capacity, experts, grid, publish, stream, grouped, steal) + flat = lambda n, dtype=c.Uint8, align=16: fake((n,), (1,), dtype, align) + box = (c.Int64(0), c.Int32(0), c.Int32(0)) if publish else () + if publish == "twoshot": # receive slots, owners' bases, parity x TP + rank + box = (fake((3 * tp, 64, 3584 // tp)), fake((tp,), None, c.Int64), c.Int32(0)) + x, ids = fake((batch, 3584), (3584, 1)), flat(batch * 16, c.Int32) + x = flat(PLAN_BYTES) if moe.grouped else x # grouped: the MXFP8 plan + scales = (flat(experts * n) for n in (moe.w13_scale_bytes, moe.w2_scale_bytes)) + out = flat(batch * 16 * (1 if publish else 3584), c.BFloat16) + maps = flat(768 if moe.paired else 256, align=64) + tensors = (x, ids, maps, *scales, flat(moe.workspace_bytes), out) + scalars = (c.Int32(batch), *map(c.Float32, (100, 0.25, 0.04)), *box) + with torch.accelerator.device_index(device): + return cute.compile(moe, *tensors, *scalars, make_fake_stream()) + + +@functools.cache +def compile_group(tp, experts, batch, grid, device, early=False): + """The pregrouped experts' unit-table kernel: (ids, workspace, batch); ``early``: + it releases the experts at entry.""" + capacity = 1 << (batch - 1).bit_length() + kernel = ExpertsGroupKernel(tp, capacity, experts, grid, early=early) + flat = lambda n, dtype=c.Uint8: fake((n,), (1,), dtype, 16) + args = (flat(batch * 16, c.Int32), flat(kernel.workspace_bytes), c.Int32(batch)) + with torch.accelerator.device_index(device): + return cute.compile(kernel, *args, make_fake_stream()) + + +@functools.cache +def compile_plan(batch, device): + """Plan writer for callers without the routing kernel's copy role in front.""" + args = (fake((batch, 3584)), fake((PLAN_BYTES,), (1,), c.Uint8), make_fake_stream()) + with torch.accelerator.device_index(device): + return cute.compile(ExpertsPlanKernel(batch), *args) + + +def scratch(tp, device): + """Model-owned expert buffers: the workspace and the grouped plan (MXFP8 x, written + by the routing kernel or ``compile_plan``).""" + size, u8 = (ExpertsGeometry(tp).workspace_bytes, torch.uint8) + workspace = torch.full((size,), 255, device=device, dtype=u8) + plan = torch.empty(PLAN_BYTES, device=device, dtype=u8) + return {"workspace": workspace, "plan": plan} + + +class PersistentExperts: + def __init__(self, w13, s13, w2, s2, *, beta=4.0, linear_beta=25.0): + experts, rows, _ = w13.shape + self.tp, self.grid = (8, 136) if rows == 768 else (16, 152) + self.device, self.experts = w13.device, experts + self.weights = (w13, s13, w2, s2) + # TP8: B3-4 on 148 CTAs; B5-8 publish on 128-136 CTAs, leaving SMs for the + # one-wave tail; grouped B9-64 on 152. + grids = "3-4:148,5-6:128,7:132,9-64:152" if self.tp == 8 else "" + self.grids = batch_table(grids) + scalars = (beta * linear_beta, 1 / beta, 1 / linear_beta) + self.scalars = tuple(map(c.Float32, scalars)) + # The experts kernel uses maps 0-1 (L2 promotion on w13); the TP8 B1 cluster + # experts use maps 2-4 (w13, w2 two-chunk and one-chunk boxes, promoted); TP8 + # FC2 pairs (B4-64) use map 5 (w2 chunk 2 of 256 rows). + w2_tile = (192, 3, 64) if self.tp == 8 else (128, 2, 32) + specs = [(w13, rows, 1792, 28, 64, 2, True), (w2, 3584, *w2_tile, 2, False)] + if self.tp == 8: + specs += [specs[0]] + [(w2, 3584, 192, 3, 64, box, True) for box in (2, 1)] + dtype = cuda.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B + swizzle = cuda.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_128B + l2 = cuda.CUtensorMapL2promotion + l2 = (l2.CU_TENSOR_MAP_L2_PROMOTION_NONE, l2.CU_TENSOR_MAP_L2_PROMOTION_L2_256B) + maps = [ + (weight.data_ptr(), dtype, (128, experts * n, chunks), (row_bytes, stride)) + + ((128, 128, box), swizzle, l2[promote]) + for weight, n, row_bytes, chunks, stride, box, promote in specs + ] + if self.tp == 8: + shape = (128, experts * 3584, 3), (192, 64), (128, 256, 1) + maps.append((w2.data_ptr(), dtype, *shape, swizzle, l2[False])) + self.maps = tensor_maps(maps, self.device) + scales = (tensor(s.view(torch.uint8)) for s in (s13, s2)) + self.views = (to_cute(self.maps[:256], 64), *scales) + self.all_maps = to_cute(self.maps, 64) # TP8 FC2 pairs: maps 0-5 + self.prepared = {} + + def prepare( + self, batch, buffers, out, publish=(), planned=False, twoshot=False, early=False + ): + """``publish``: (multicast, parity, rank) of the WaveMoeTail latent slot to + publish (``twoshot``: the ``common.twoshot`` push box); ``planned``: B>8 calls + find the plan written (by the routing kernel), else run writes it from x first; + ``early``: the caller runs ``group`` (pregrouped batches) before ``run``.""" + mode = "twoshot" if twoshot else bool(publish) + grouped = pregrouped(self.tp, batch) + args: tuple = (batch, self.grid_for(batch), self.device.index, mode) + args += (streamed(batch), grouped, stealing(batch)) + code = compile_experts(self.tp, self.experts, *args) + # CTAs wait on other clusters' partials: the whole grid must be resident. + check_residency(code, self.device.index, args[1], 384, "Kimi K3 MoE experts") + ws, maps = tensor(buffers["workspace"]), self.views[0] + plan: tuple | None = None + capacity = 1 << (batch - 1).bit_length() + if fc2_pairs(self.tp, capacity, batch > 8 or streamed(batch)): + maps = self.all_maps + group = None + if batch > 8: + plan = (tensor(buffers["plan"]), compile_plan(batch, self.device.index)) + plan = plan[:1] if planned else plan + if grouped: + release = early_group(self.tp, batch) + grouping = compile_group(self.tp, self.experts, *args[:3], release) + group = grouping, ws, early + self.prepared[batch] = (out, tensor(out), ws, code, publish, plan, maps, group) + + def group(self, ids): + """Pregrouped batches prepared ``early``: the unit-table kernel (run it on the + stream after routing; ``run`` then skips it).""" + group = self.prepared[ids.shape[0]][7] + if group is not None and group[2]: + self.launch_group(ids, group) + + def launch_group(self, ids, group): + code, ws, _ = group + code(tensor(ids), ws, c.Int32(ids.shape[0]), stream(self.device)) + + def grid_for(self, batch): + return self.grids.get(batch, self.grid) + + def run(self, x, ids, weights=None): + batch = x.shape[0] + out, output, workspace, code, box, plan, maps, group = self.prepared[batch] + output = tensor(weights) if box else output + source = to_cute(x, 16) + if plan is not None: + if len(plan) > 1: + plan[1](source, plan[0], stream(self.device)) + source = plan[0] + if group is not None and not group[2]: + self.launch_group(ids, group) + args = (source, tensor(ids), maps, *self.views[1:], workspace, output) + code(*args, c.Int32(batch), *self.scalars, *box, stream(self.device)) + return out diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/front.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/front.py new file mode 100644 index 0000000000..0067ed5add --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/front.py @@ -0,0 +1,538 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MoE front: the weight-stationary [router shard; down shard] GEMM publishes logits +and the latent shard into a multicast mailbox row; the routing kernel polls all 896 +logits, ranks the top16 (sigmoid + bias, renormalized) and Lamport-copies the latent.""" + +import functools + +import cutlass as c +import cutlass.cute as cute +import torch +from cutlass import Int32, Uint32, Uint64 +from cutlass.cute.runtime import make_fake_stream + +from ..common.host import fake, stream, to_cute +from ..common.mailbox import poll, rearm +from ..common.policy import batch_table +from ..common.primitives import ( + alloc, + asm, + bits_float, + div_rn, + float_bits, + launch, + redux, + sigmoid, + store_words, + view, +) +from ..gemm.projection import PreparedProjection +from ..gemm.ws_gemm import WsGemm +from .experts import PLAN_BYTES, plan_chunk +from .shared import compile_down, compile_shared, situ, weight_maps +from .tail_runtime import CAP + + +@asm(Int32) +def shared_add(address, value): + code = "atom.shared.add.u32 $0,[$1],$2;" + return code, "=r,r,r,~{memory}", (address, Int32(value)) + + +@cute.jit +def ordered(score): + """Monotone uint32 key of a float; NaN maps to 0 (below every number).""" + bits = Uint32(float_bits(score)) + key = ~bits if (bits & Uint32(0x80000000)) != 0 else bits | Uint32(0x80000000) + return key if score == score else Uint32(0) + + +class RouteSitu: + def __init__( + self, batch, shared, renorm, early, roles, polled, plan=False + ): # fmt: skip + # B>8: few registers, so routing CTAs co-reside with the front and the fused + # shared expert, and polled logits skip the front's completion wait. + self.lean = batch > 8 + self.batch, self.shared, self.experts, self.renorm = batch, shared, 896, renorm + # early: release the persistent experts at entry (safe once down drains first); + # polled: logits arrive as FP32 words in mailbox columns [3584, 5376) (Lamport); + # plan: the copy role also writes the latent as MXFP8 into the experts' plan. + self.early, self.polled, self.plan = early, polled, plan + self.roles = tuple(r for r, on in zip(("route", "copy", "situ"), roles) if on) + + @cute.jit + def __call__(self, *args): + launch(self.kernel, args, len(self.roles) * self.batch, 256) + + @cute.kernel + def kernel(self, args): + logits, bias, ids, weights, mailbox, latent, gate_up, hidden = args[0:8] + scale, beta, linear_beta = args[8:11] + plan = args[11] if self.plan else None + tid, _, _ = cute.arch.thread_idx() + cta, _, _ = cute.arch.block_idx() + token, role = cta % self.batch, cta // self.batch + if c.const_expr(self.early): + cute.arch.griddepcontrol_launch_dependents() + for r in c.range_constexpr(len(self.roles)): + if role == r: + if c.const_expr(self.roles[r] == "route"): + self.route(logits, bias, ids, weights, scale, token, tid, mailbox) + elif c.const_expr(self.roles[r] == "copy"): + # Lamport: no dependency wait; the sentinel poll is the barrier. + base = mailbox.iterator + token * mailbox.shape[2] + tid * 8 + if tid < 192: + self.copy(base, latent, token, tid, 2, plan) + else: + self.copy(base, latent, token, tid, 1, plan) + else: + cute.arch.griddepcontrol_wait() + rounds = (self.shared + 255) // 256 # issue every load first + pair = cute.make_rmem_tensor((rounds, 2), c.Float32) + for i in c.range_constexpr(rounds): + if tid + 256 * i < self.shared: + pair[i, 0] = gate_up[token, tid + 256 * i] + pair[i, 1] = gate_up[token, self.shared + tid + 256 * i] + for i in c.range_constexpr(rounds): + if tid + 256 * i < self.shared: + gate = c.Float32(c.BFloat16(pair[i, 0])) + up = c.Float32(c.BFloat16(pair[i, 1])) + value = situ(gate, up, beta, linear_beta) + hidden[token, tid + 256 * i] = c.BFloat16(value) + # Late release: the persistent experts must not take SMs from the down GEMM. + if c.const_expr(not self.early): + cute.arch.griddepcontrol_launch_dependents() + + @cute.jit + def select(self, key, threshold, buffers): + """B>8 candidates (keys at or above ``threshold``) append to ``dense`` at a + per-warp offset from the shared count, then rank themselves: 8 threads per + candidate (4 past 32 candidates) each count their share of the candidates above + it (key, then lower expert). Ranks below 16 store the expert and its sigmoid.""" + dense, sig, chosen, picked, found = buffers + tid, _, _ = cute.arch.thread_idx() + lane = tid % 32 + packed = cute.make_rmem_tensor(4, Uint64) + ballots = cute.make_rmem_tensor(4, Uint32) + total = Int32(0) + for i in c.range_constexpr(4): # packed 0: not a candidate + expert = self.expert(tid, i) + hit = (expert < self.experts) & (key[i] >= threshold) + ballots[i] = cute.arch.vote_ballot_sync(hit) + packed[i] = Uint64(0) + if hit: + packed[i] = (Uint64(key[i]) << 32) | Uint64(~Uint32(expert)) + total += cute.arch.popc(ballots[i]) + offset = Int32(0) + if (lane == 0) & (total > 0): + offset = shared_add(found.iterator, total) + offset = cute.arch.shuffle_sync(offset, 0) + below = (Uint32(1) << lane) - Uint32(1) + for i in c.range_constexpr(4): + if packed[i] != Uint64(0): + dense[offset + cute.arch.popc(ballots[i] & below)] = packed[i] + offset += cute.arch.popc(ballots[i]) + cute.arch.sync_threads() + n = found[0] + if n <= 32: # usually: 8 threads per candidate, 4 slots each + part, index = tid % 8, tid // 8 + mine, rank = Uint64(0), Int32(0) + if index < n: + mine = Uint64(dense[index]) + for k in c.range_constexpr(4): # slots past n do not count + j = part + 8 * k + rank += Int32((j < n) & (Uint64(dense[j]) > mine)) + for shift in c.range_constexpr(3): + rank += cute.arch.shuffle_sync_bfly(rank, 1 << shift) + if (part == 0) & (index < n) & (rank < 16): + expert = Int32(~Uint32(mine)) + chosen[rank], picked[rank] = expert, sig[expert] + else: + part = tid % 4 + for first in c.range(0, n, 64, unroll=1): + index = first + tid // 4 + mine, rank = Uint64(0), Int32(0) + if index < n: + mine = Uint64(dense[index]) + if n <= 64: # every load issues at once; slots past n do not count + for k in c.range_constexpr(16): + j = part + 4 * k + rank += Int32((j < n) & (Uint64(dense[j]) > mine)) + else: + for j in c.range(part, n, 4, unroll=2): + rank += Int32(Uint64(dense[j]) > mine) + rank += cute.arch.shuffle_sync_bfly(rank, 1) + rank += cute.arch.shuffle_sync_bfly(rank, 2) + if (part == 0) & (index < n) & (rank < 16): + expert = Int32(~Uint32(mine)) + chosen[rank], picked[rank] = expert, sig[expert] + cute.arch.sync_threads() + + @cute.jit + def publish(self, chosen, picked, ids, weights, scale, token): + """Warp 0: the 16 ids by rank and their sigmoids renormalized (lane r holds + rank r; the butterfly sums in the same order for every lane).""" + tid, _, _ = cute.arch.thread_idx() + lane = tid % 32 + if tid < 32: + value = c.Float32(0) + if lane < 16: + ids[token, lane] = chosen[lane] + value = picked[lane] + total = value + for shift in c.range_constexpr(4): + total += cute.arch.shuffle_sync_bfly(total, 1 << shift) + if lane < 16: + value = value * scale + if c.const_expr(self.renorm): + value = div_rn(value, total + c.Float32(1e-20)) + weights[token, lane] = c.BFloat16(value) + + def expert(self, tid, i): + return tid * 4 + i if self.polled else tid + i * 256 + + @cute.jit + def copy(self, base, latent, token, tid, count: c.Constexpr[int], plan=None): + got = poll(base, 2048, count) + for i in c.range_constexpr(count): + dst = latent.iterator + token * latent.stride[0] + (tid + 256 * i) * 8 + words = [got[i, k] for k in range(4)] + store_words(dst, words) + rearm(base + i * 2048) + if c.const_expr(plan is not None): + plan_chunk(plan, token, tid + 256 * i, words) + + @cute.jit + def route(self, logits, bias, ids, weights, scale, token, tid, mailbox): + """Exact top16 by threads [0, 256): candidates at or above the 16th largest + lane-group maximum are compacted and ranked by an unrolled 64-way compare.""" + warp, lane = tid // 32, tid % 32 + smem = c.utils.SmemAllocator() + km = alloc(smem, Uint32, (8, 32)) + cand, dense = (alloc(smem, Uint64, shape) for shape in ((8, 128), 64)) + sig = alloc(smem, c.Float32, 1024) + counts = alloc(smem, Int32, 8) + cutoff = alloc(smem, Uint32, 1) + chosen = alloc(smem, Int32, 16) + if c.const_expr(self.lean): # candidate count; the chosen experts' sigmoids + found, picked = alloc(smem, Int32, 1), alloc(smem, c.Float32, 16) + # Thread t owns experts t + 256 i, or 4 t + i when polling 16-byte chunks. + offsets = cute.make_rmem_tensor(4, c.Float32) + for i in c.range_constexpr(4): # bias is a weight: load before the wait + offsets[i] = c.Float32(0) + if self.expert(tid, i) < self.experts: + offsets[i] = bias[self.expert(tid, i)] + raw = cute.make_rmem_tensor(4, c.Float32) + key, top = cute.make_rmem_tensor(4, Uint32), Uint32(0) + for i in c.range_constexpr(4): + key[i] = Uint32(0) + if c.const_expr(self.polled): + if tid < 224: + cell = mailbox.iterator + token * mailbox.shape[2] + 3584 + tid * 8 + got = poll(cell, 0, 1) + rearm(cell) + for i in c.range_constexpr(4): + raw[i] = bits_float(got[0, i]) + if c.const_expr(self.lean): # all four experts: one block, no branches + scores = [sigmoid(raw[i]) for i in range(4)] + for i in c.range_constexpr(4): + sig[4 * tid + i] = scores[i] + key[i] = ordered(scores[i] + offsets[i]) + if c.const_expr(not self.lean): + cute.arch.griddepcontrol_wait() + else: + cute.arch.griddepcontrol_wait() + for i in c.range_constexpr(4): + if self.expert(tid, i) < self.experts: + raw[i] = logits[token, self.expert(tid, i)] + for i in c.range_constexpr(4): + expert = self.expert(tid, i) + if c.const_expr(not (self.lean and self.polled)): # noqa: SIM102 + if expert < self.experts: + score = sigmoid(raw[i]) + sig[expert] = score + key[i] = ordered(score + offsets[i]) + top = max(top, key[i]) + km[warp, lane] = top + if c.const_expr(self.lean): # noqa: SIM102 + if tid == 0: + found[0] = Int32(0) + cute.arch.sync_threads() + if warp == 0: + group = km[0, lane] + for w in c.range_constexpr(1, 8): + group = max(group, km[w, lane]) + km[0, lane] = group + cute.arch.sync_warp() + rank = Int32(0) + for j in c.range_constexpr(32): + rank += Int32(km[0, j] > group) + low = redux(group if rank < 16 else Uint32(0xFFFFFFFF)) + if lane == 0: + cutoff[0] = low + cute.arch.sync_threads() + if c.const_expr(self.lean): + dense = view(cand.iterator, 0, Uint64, (1024,)) + self.select(key, cutoff[0], (dense, sig, chosen, picked, found)) + self.publish(chosen, picked, ids, weights, scale, token) + return + threshold, offset = cutoff[0], Int32(0) + for i in c.range_constexpr(4): + expert = self.expert(tid, i) + hit = (expert < self.experts) & (key[i] >= threshold) + ballot = cute.arch.vote_ballot_sync(hit) + if hit: + below = cute.arch.popc(ballot & ((Uint32(1) << lane) - Uint32(1))) + packed = (Uint64(key[i]) << 32) | Uint64(~Uint32(expert)) + cand[warp, offset + below] = packed + offset += cute.arch.popc(ballot) + if lane == 0: + counts[warp] = offset + cute.arch.sync_threads() + if warp == 0: + n = Int32(0) + for v in c.range_constexpr(8): + n += counts[v] + for half in c.range_constexpr(2): # compact the regions densely + index = lane + 32 * half + rest, region = Int32(index), Int32(0) + for v in c.range_constexpr(8): + if (region == v) & (rest >= counts[v]): + rest -= counts[v] + region += 1 + value = Uint64(0) + if index < n: + value = cand[region, rest] + dense[index] = value + cute.arch.sync_warp() + if n <= 64: + for half in c.range_constexpr(2): + mine, rank = dense[lane + 32 * half], Int32(0) + for j in c.range_constexpr(64): + rank += Int32(dense[j] > mine) + if (lane + 32 * half < n) & (rank < 16): + chosen[rank] = Int32(~Uint32(mine)) + else: # rare: rank every candidate against all regions + for q in range(lane, n, 32): + rest, region = Int32(q), Int32(0) + for v in c.range_constexpr(8): + if (region == v) & (rest >= counts[v]): + rest -= counts[v] + region += 1 + mine, rank = cand[region, rest], Int32(0) + for v in c.range_constexpr(8): + for j in range(counts[v]): + rank += Int32(cand[v, j] > mine) + if rank < 16: + chosen[rank] = Int32(~Uint32(mine)) + cute.arch.sync_warp() + expert, value = chosen[lane % 16], c.Float32(0) + if lane < 16: + ids[token, lane] = expert + value = sig[expert] + total = value + for shift in c.range_constexpr(4): + total += cute.arch.shuffle_sync_bfly(total, 1 << shift) + if lane < 16: + value = value * scale + if c.const_expr(self.renorm): + value = div_rn(value, total + c.Float32(1e-20)) + weights[token, lane] = c.BFloat16(value) + + +# Front GEMM geometry "ks.kb.tg.rt.ch" per batch, by packed rows (TP8 560, TP16 288), +# from GB300 sweeps. ks: cluster K-split; kb: 32-wide +# K blocks per warp, so a row tile takes 224 / (ks * kb) warps (at most 16); tg: tokens +# per group (0: the single-pass path); rt: 16-row tiles per CTA; ch: accumulator chains +# per 8-token tile (0: 4 / tiles per group). +FRONT = { + 560: "1-4:4.14.0.1.0,5-32:16.7.16.5.2,33-64:16.14.16.5.2", + 288: "1-4:4.14.0.1.0,5-64:16.7.16.3.2", +} +WIDTH = 3584 + 2 * 896 # front mailbox row: BF16 latent, then FP32 logit words +# Batches whose shared expert is the fused kernel (``SharedUmma``), not the split-K +# chain; of those, split batches run gate/up + SiTU beside the front and the down rows +# (``SharedDown``) on SPLIT SMs beside the routed experts (0: fused). +FUSED, SPLIT = "9-64:1", "9-12:6,13-24:4,25-64:2" +# Per-TP overrides of SPLIT entries, (q = 1, verify steps): on real prompts (GB300 +# SPEED-Bench 8k mixed) SharedDown on 2 SMs (~123 us at M32) outlasts the experts (TP8 +# q=1 B32: 3 SMs, -5.4% step). Verify steps (q > 1) share experts among a request's q +# tokens: TP8 DL7 M16/32/64 -5.9/-12.9/-1.1%, TP16 M32 -5.9%. +SPLITS = {8: ("25-32:3", "13-16:6,25-32:4,33-64:3"), 16: ("", "25-32:3")} +# Split batches whose front GEMM rows join the gate/up kernel as extra M-tiles (one +# kernel instead of two that slow each other ~2x side by side). +MERGED = {8: "41-64:1", 16: "33-64:1"} +# Shared-expert split-K tactics (M tile, N tile, split, stages) by width (TP8 768, TP16 +# 384): gate_up, then down. +SHARED_TACTICS = { + 768: ( + "1-8:128.8.4.6,9-16:64.16.4.10,17-32:64.32.4.8,33-64:64.64.4.5", + "1-8:128.8.1.6,9-16:64.16.1.6,17-32:64.32.1.6,33-64:64.64.1.6", + ), + 384: ( + "1-16:64.8.4.8,17-32:64.16.4.10,33-64:64.32.4.8", + "1-8:128.8.1.3,9-16:64.16.1.3,17-32:64.32.1.3,33-64:64.64.1.3", + ), +} + + +@functools.cache +def front_table(n): + return batch_table(FRONT.get(n, "")) + + +def front_ks(batch, n): + return front_table(n).get(batch, (4,))[0] + + +def front_geometry(batch, n, ks, kb=None, tg=None, rt=None, ch=None): + """(rt, kq, staged, tg, ch) of the front GEMM over ``n`` rows. Unset arguments come + from the table when ``ks`` is its K-split, else from the single pass; token groups + whose x slice does not fit in SMEM fall back to the single pass.""" + entry = front_table(n).get(batch, ()) + table = entry[1:] if entry[:1] == (ks,) else (14, 0, 1, 0) + kb, tg, rt, ch = (t if v is None else v for v, t in zip((kb, tg, rt, ch), table)) + kq = min(16, 224 // (ks * kb)) + kb, nt = 224 // (ks * kq), -(-batch // 8) + staged = max(kb - 8, 0) if kq == 16 else 0 + used = ( + rt * kq * (nt * 544 + max(staged, 1) * 1024) + -(-nt * rt // ks) * ks * kq * 512 + ) + if nt * 8 * (7168 // ks * 2 + 64) > 220 * 1024 - used: + tg = 0 + return rt, kq, staged, tg, ch + + +@functools.cache +def compile_route(batch, shared, renorm, roles, device, polled=False, plan=False): + kernel = RouteSitu(batch, shared, renorm, batch <= 2, roles, polled, plan) + f32 = c.Float32 + args: tuple = (fake((batch, 896), None, f32), fake((896,), None, f32)) + args += (fake((batch, 16), None, c.Int32), fake((batch, 16)), fake((1, 64, WIDTH))) + args += (fake((batch, 3584)), fake((batch, 2 * shared), None, f32)) + args += (fake((batch, shared)), *(f32(v) for v in (1, 4, 25))) + args += (fake((PLAN_BYTES,), None, c.Uint8),) * plan + (make_fake_stream(),) + with torch.accelerator.device_index(device): + return cute.compile(kernel, *args) + + +@functools.cache +def compile_publish(n, split, valid, batch, device, ks, geometry=None): + """Front GEMM that only publishes; routing then runs in the polled kernel.""" + rt, kq, staged, tg, ch = geometry or front_geometry(batch, n, ks) + kernel = WsGemm(n, 7168, batch, rt, kq, staged, split, valid, ks, tg, ch) + args: tuple = (fake((batch, 7168)), fake((n, 7168)), fake((1, 64, WIDTH))) + args += (c.Int64(0), c.Int32(0), c.Int32(0), fake((batch, 896), None, c.Float32)) + with torch.accelerator.device_index(device): + return cute.compile(kernel, *args, make_fake_stream()) + + +class SharedFront: + def __init__( + self, moe, packed, box, storage, batches, beta, linear_beta, plan=None, + counters=None, q=1, + ): # fmt: skip + """``plan``: the experts' plan buffer, written by the routing kernel at B>8; + ``counters``: the fused shared expert's model-owned barrier counters; ``q``: the + deployment's verify length (its SPLITS entries).""" + experts = moe.shared_experts + gate_up, down = experts.gate_up_proj.weight, experts.down_proj.weight + self.shared = shared = gate_up.shape[0] // 2 + tp = 6144 // shared + shard, valid = 3584 // tp, 896 // tp + self.device, self.packed, self.storage = packed.device, to_cute(packed), storage + gu, dn = map(batch_table, SHARED_TACTICS[shared]) + self.gate_up = PreparedProjection( + gate_up, {b: gu[b] for b in batches}, output_dtype=torch.float32 + ).prepare() + self.down = PreparedProjection(down, {b: dn[b] for b in batches}).prepare() + scale, renorm = moe.routed_scaling_factor, moe.moe_renormalize + self.scalars = tuple(map(c.Float32, (scale, beta, linear_beta))) + self.box: tuple = (to_cute(box.buffer), c.Int64(box.mc)) + self.box += (c.Int32(box.rank * valid), c.Int32(box.rank * shard)) + bias = to_cute(moe.gate.e_score_correction_bias) + dev, n = packed.device.index, packed.shape[0] + fused = batch_table(FUSED) + fused = {b for b in batches if b > 8 and counters is not None and fused.get(b)} + verify = q > 1 + split = batch_table(f"{SPLIT},{SPLITS[tp][verify]}") + self.split = {b: split[b] for b in fused if split.get(b)} # down SMs + if fused: # TMA maps and counters of the fused kernel + self.maps = weight_maps(gate_up, down) + self.fused = to_cute(self.maps, 64), to_cute(counters) + merged = batch_table(MERGED[tp]) + self.merged = {b for b in self.split if merged.get(b)} + if self.merged: # gate/up, down and front maps; unused RS arguments + self.merged_maps = weight_maps(gate_up, down, front=packed) + # Kept alive here: a CuTe view does not hold its tensor. + self.unused = torch.zeros(tp, dtype=torch.int64, device=packed.device) + self.merged_args = to_cute(self.merged_maps, 64), to_cute(self.unused) + geometry = (n, n - shard, valid, box.buffer.shape[-1]) + self.plans: dict[int, tuple] = {} + for b in batches: + views = tuple(t[:b] for t in storage) + logits, gate, ids, weights, h, _, latent = map(to_cute, views) + front = compile_publish(n, n - shard, valid, b, dev, front_ks(b, n)) + planned = plan is not None and b > 8 + roles = (True, True, False) + route = compile_route(b, shared, renorm, roles, dev, True, planned) + situ = compile_route(b, shared, renorm, (False, False, True), dev) + args = (logits, bias, ids, weights, self.box[0], latent, gate, h) + extra = (to_cute(plan),) if planned else () + code = None + if b in self.merged: + gate = compile_shared(tp, b, dev, True, geometry) + code = gate, compile_down(tp, b, dev, self.split[b]) + elif b in self.split: + gate = compile_shared(tp, b, dev, gate=True) + code = gate, compile_down(tp, b, dev, self.split[b]) + elif b in fused: + code = compile_shared(tp, b, dev) + self.plans[b] = views, front, route, situ, args + self.scalars, extra, code + + def route(self, x): + """Router logits + published down shard, then top16 and the latent copy (and + at B>8 the experts' MXFP8 plan). Merged batches publish them from the gate/up + kernel, which also leaves h for ``shared_expert``'s down rows.""" + views, front, route, _, args, extra, code = self.plans[x.shape[0]] + s = stream(self.device) + if x.shape[0] in self.merged: + maps, unused = self.merged_args + work = (to_cute(x), maps, args[7], self.fused[1], unused) + work += (Int32(0), Int32(0), Int32(CAP), *self.scalars[1:], *self.box[1:]) + code[0](*work, s) + else: + front(to_cute(x), self.packed, *self.box, args[0], s) + route(*args, *extra, s) + return views[2], views[3], views[6] + + def merged_batch(self, batch): + return batch in self.merged + + def planned(self, batch): + return bool(self.plans[batch][5]) + + def shared_expert(self, x, tail): + """The shared expert's rows into the tail's RS mailboxes (``tail.push``); the + fused kernel peer-stores them itself and leaves the returned rows unwritten. + Returns (rows, the split's down launch to run beside the experts, or None).""" + views, _, _, situ, args, _, code = self.plans[x.shape[0]] + s = stream(self.device) + if code is not None: + maps, counters = self.fused + mailbox = (tail.peers, *tail.ring, Int32(CAP)) + front = (to_cute(x), maps, args[7], counters, *mailbox, *self.scalars[1:]) + if x.shape[0] not in self.split: + code(*front, s) + return views[5], None + gate, down = code + if x.shape[0] not in self.merged: # merged: h comes from ``route`` + gate(*front, s) + return views[5], lambda: down(args[7], maps, *mailbox, stream(self.device)) + self.gate_up.run(x, views[1]) + situ(*args, s) + shared = self.down.run(views[4], views[5]) + tail.push(shared) + return shared, None diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/prepared.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/prepared.py new file mode 100644 index 0000000000..12d03fe2df --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/prepared.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Prepared MoE provider: one model owner serializes its scratch, no tensor pointers +are rebound at execution, and the side stream joins before the tail.""" + +import torch + +from ..common.policy import batch_table +from .experts_runtime import scratch as expert_scratch +from .shared import scratch as shared_scratch + + +class MoEScratch: + def __init__(self, tp, device, batches): + self.tp, self.device, self.batches = tp, torch.device(device), tuple(batches) + cap, h = max(batches), 6144 // tp + with torch.accelerator.device_index(self.device.index): + self.stream = torch.cuda.Stream(device=self.device) + self.begin, self.shared_done = torch.cuda.Event(), torch.cuda.Event() + self.gate_done = torch.cuda.Event() # the split shared expert's front half + f32, i32, bf16 = torch.float32, torch.int32, torch.bfloat16 + empty = lambda s, t=bf16: torch.empty(s, dtype=t, device=self.device) + self.latent = empty((cap, 3584)) + self.front: tuple = (empty((cap, 896), f32), empty((cap, 2 * h), f32)) + self.front += (empty((cap, 16), i32), empty((cap, 16)), empty((cap, h))) + self.front += (empty((cap, 7168)), self.latent) + self.expanded = empty((cap * 16, 3584)) + self.experts = expert_scratch(tp, self.device) + # The experts' plan (MXFP8 x), written by the routing kernel at B>8. + self.plan = self.experts["plan"] + self.shared = shared_scratch(self.device) # fused shared-expert counters + self.views = {} + for b in batches: + self.views[b] = (self.latent[:b], self.expanded[: b * 16], self.experts) + + +# TP8 B16: the experts start beside the split gate/up half (+1% E2E). +JOIN = {8: "9-16:0", 16: ""} + + +class PreparedMoE: + def __init__(self, front, experts, tail, state, *, prefetch=None, layer_idx=None): + self.front, self.experts, self.tail, self.state = front, experts, tail, state + self.prefetch, self.layer_idx = prefetch, layer_idx + self.small_batch = None # optional small-batch chain + # Split batches whose routed experts launch beside the shared expert's + # gate/up half instead of after it. + self.join = batch_table(JOIN[state.tp]) + self.sms = torch.cuda.get_device_properties(state.device).multi_processor_count + for b, sms in front.split.items(): # the split's down CTAs take the SMs left + experts.grids[b] = min(experts.grid_for(b), self.sms - sms) + for b in state.batches: + _, expanded, buffers = state.views[b] + publish, planned = tail.publish.get(b, ()), front.planned(b) + twoshot = b in tail.twoshots + experts.prepare(b, buffers, expanded, publish, planned, twoshot, True) + + def run(self, x, out, *, consume=None): + state, b = self.state, x.shape[0] + if self.small_batch is not None and b in self.small_batch.batches: + return self.small_batch.run(out, consume) + main = torch.cuda.current_stream(state.device) + # The shared expert runs on the side stream beside the routed experts, or, when + # the persistent expert grid spans every SM, beside the front only. Split: its + # gate/up half beside the front, its down rows beside the experts. + before = self.experts.grid_for(b) >= self.sms + if self.front.merged_batch(b): + # One kernel publishes the routing logits and leaves h; the down rows then + # run beside the experts. + ids, weights, latent = self.front.route(x) + state.begin.record(main) + with torch.cuda.stream(state.stream): + state.stream.wait_event(state.begin) + shared, beside = self.front.shared_expert(x, self.tail) + beside() + state.shared_done.record(state.stream) + self.experts.group(ids) + return self.finish(b, out, latent, ids, weights, shared, consume) + state.begin.record(main) + with torch.cuda.stream(state.stream): + state.stream.wait_event(state.begin) + shared, beside = self.front.shared_expert(x, self.tail) + if beside is not None: + state.gate_done.record(state.stream) + beside() + state.shared_done.record(state.stream) + ids, weights, latent = self.front.route(x) + self.experts.group(ids) # pregrouped unit table, beside the shared expert + if before or (beside is not None and self.join.get(b, 1)): + main.wait_event(state.shared_done if before else state.gate_done) + return self.finish(b, out, latent, ids, weights, shared, consume, before) + + def finish(self, b, out, latent, ids, weights, shared, consume, before=False): + """The routed experts, the next layer's prefetch fork, then the tail once the + shared expert's side stream is done.""" + main = torch.cuda.current_stream(self.state.device) + expanded = self.experts.run(latent, ids, weights) + if self.prefetch is not None: + self.prefetch.fork("next_attention", self.layer_idx, b) + if not before: + main.wait_event(self.state.shared_done) + return self.tail.run(expanded, weights, shared, out, consume=consume) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/routing.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/routing.py new file mode 100644 index 0000000000..23bf073d23 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/routing.py @@ -0,0 +1,254 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Top16 routing (threshold, regional candidates and exact composite ranking) and the +helpers the TP8 and TP16 small-batch fronts share.""" + +from typing import Any + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Uint32, Uint64 + +from ..common.primitives import ( + bits_float, + copy_evict, + float_bits, + gate_bf16_bits, + mailbox_load, + mailbox_store, + no_neg_zero, + published, + redux, + store_shared_words, + view, + zeros, +) + + +@cute.jit +def bf16_word(lo, hi): + return no_neg_zero(Uint32(gate_bf16_bits(lo)) | (Uint32(gate_bf16_bits(hi)) << 16)) + + +@cute.jit +def poll_words(words, pointer, valid): + """Copy four mailbox words into ``words``; clear ``valid`` if any is unpublished.""" + values = mailbox_load(pointer, count=4) + for j in cutlass.range_constexpr(4): + words[j] = values[j] + valid = valid & (values[j] != Uint32(0x80000000)) + return valid + + +@cute.jit +def gather_x(xmoe, par, cta, dst, first_chunk, active, tokens=1, h2=None): + """Poll this CTA's x_moe rows (and the TP16 H2 slices if ``h2 = (warp, lane, + pairs)``) until published, then stage the row chunks at ``dst``.""" + words = zeros((tokens, 2, 4), Uint32) + h2_words: Any = None + if cutlass.const_expr(h2 is not None): + h2_words = zeros((tokens, 4, 4), Uint32) + ready = False + while not ready: + valid = True + for token in cutlass.range_constexpr(tokens): + rows = (par * xmoe.shape[1] + cta % xmoe.shape[1]) * xmoe.shape[2] + base = xmoe.iterator + rows * 7168 + token * 7168 + if active: + for chunk in cutlass.range_constexpr(2): + pointer = base + (first_chunk + 28 * chunk) * 8 + valid = poll_words(words[token, chunk, None], pointer, valid) + if cutlass.const_expr(h2 is not None): + warp, lane, pairs = h2 + if (warp < 7) & (pairs > 0): + for chunk in cutlass.range_constexpr(4): + pointer = base + (32 * warp + lane + 224 * chunk) * 8 + valid = poll_words(h2_words[token, chunk, None], pointer, valid) + ready = cute.arch.vote_all_sync(valid) + if active: + for token in cutlass.range_constexpr(tokens): + for chunk in cutlass.range_constexpr(2): + pointer = dst + token * 14336 + (first_chunk + 28 * chunk) * 16 + values = tuple(words[token, chunk, j] for j in range(4)) + store_shared_words(pointer, values) + return words, h2_words + + +@cute.jit +def score_record(score, biased): + """Published (score, score+bias) bits; NaN or 0x80000000 bits become (0, -inf).""" + score_bits, biased_bits = Uint32(float_bits(score)), Uint32(float_bits(biased)) + sentinel = (score_bits == Uint32(0x80000000)) | (biased_bits == Uint32(0x80000000)) + if sentinel | (score != score) | (biased != biased): + score_bits, biased_bits = Uint32(0), Uint32(0xFF800000) + return score_bits, biased_bits + + +@cute.jit +def stage_rows(slots, pitch, slot, weights, first, rows, barrier, skip=0): + """Bulk-copy weight rows ``first + skip + i`` into slots ``slot + i``.""" + cute.arch.mbarrier_arrive_and_expect_tx(barrier, rows * 14336) + for row in range(rows): + dst = slots + (slot + row) * pitch + src = weights.iterator + (first + skip + row) * 7168 + copy_evict(dst, src, size=14336, barrier=barrier) + + +class Top16Fast: + """Eight-warp top16 of 896 experts, lower ids winning ties. ``scores`` holds FP32 + (score, score+bias) bits per [parity, token, expert]; ``run`` polls ``token``'s + parity-``par`` records until published and writes ``ids[16*token:]``. Runs on the + first 256 threads of the CTA with named barrier 1 and 13,200 bytes of SMEM.""" + + def __init__(self, batch): + self.batch = batch + + @cute.jit + def arrays(self, storage): + km = view(storage, 0, Uint32, (8, 32), (32, 1)) + gc = view(storage, 1024, Uint32, 32) + wc = view(storage, 1152, Uint32, 8) + threshold = view(storage, 1184, Uint32, 1) + cc = view(storage, 1200, Uint64, 904) + cs = view(storage, 8432, Uint32, 904) + cc2 = view(storage, 12048, Uint64, 64) + selected_scores = view(storage, 12560, Float32, 80) + selected_ids = view(storage, 12880, Int32, 80) + return km, gc, wc, threshold, cc, cs, cc2, selected_scores, selected_ids + + @cute.jit + def rank_small(self, arrays, counts, n, lane, ids, groups: cutlass.Constexpr[int]): + cc, cs, cc2, selected_scores, selected_ids = arrays[4:] + mine = cute.make_rmem_tensor(groups, Uint64) + score = cute.make_rmem_tensor(groups, Uint32) + for half in cutlass.range_constexpr(groups): + index = Uint32(lane + 32 * half) + total = Uint32(0) + slot = index + for warp in cutlass.range_constexpr(7): + total += counts[warp] + past = index >= total + gap = (128 if warp < 4 else 96) - counts[warp] + slot += gap if past else Uint32(0) + slot = min(slot, Uint32(895)) + mine[half] = cc[slot] + score[half] = cs[slot] + if index >= n: + mine[half] = Uint64(0) + cc2[index] = mine[half] + cute.arch.sync_warp() + for half in cutlass.range_constexpr(groups): + rank = zeros(4, Uint32) + for q in cutlass.range_constexpr(32 * groups): + rank[q % 4] += Uint32(cc2[q] > mine[half]) + r = (rank[0] + rank[1]) + (rank[2] + rank[3]) + index = Uint32(lane + 32 * half) + selected = (index < n) & (r < 16) + slot = r if selected else 16 + index + expert = Int32(Uint32(0xFFFFFFFF) - Uint32(mine[half])) + selected_ids[slot] = expert + selected_scores[slot] = bits_float(score[half]) + if selected: + mailbox_store(ids + r, (Uint32(expert),)) + + @cute.jit + def run(self, scores, ids, storage, par, token): + tid, _, _ = cute.arch.thread_idx() + warp, lane = tid // 32, tid % 32 + records = zeros((4, 2), Uint32) + pending = Uint32(0) + for i in cutlass.range_constexpr(4): + if warp + 8 * i < 28: + pending |= Uint32(1 << i) + while pending != 0: + for i in cutlass.range_constexpr(4): + if (pending & Uint32(1 << i)) != 0: + row = (par * self.batch + token) * 896 + lane + pointer = scores.iterator + (row + 32 * (warp + 8 * i)) * 2 + words = mailbox_load(pointer, count=2) + records[i, 0], records[i, 1] = words[0], words[1] + for i in cutlass.range_constexpr(4): + done = (pending & Uint32(1 << i)) != 0 + if published((records[i, j] for j in range(2)), done): + pending &= Uint32(0xFFFFFFFF ^ (1 << i)) + cute.arch.barrier(barrier_id=1, number_of_threads=256) + ids = ids.iterator + token * 16 + arrays = self.arrays(storage) + km, gc, wc, threshold, cc, cs, _, selected_scores, selected_ids = arrays + key = cute.make_rmem_tensor(4, Uint32) + maximum = Uint32(0) + for i in cutlass.range_constexpr(4): + biased = bits_float(records[i, 1]) + bits = records[i, 1] + negative = (bits & Uint32(0x80000000)) != 0 + ordered = ~bits if negative else bits | Uint32(0x80000000) + key[i] = ordered if (warp + 8 * i < 28) & (biased == biased) else Uint32(0) + maximum = max(maximum, key[i]) + km[warp, lane] = maximum + cute.arch.barrier(barrier_id=1, number_of_threads=256) + if warp == 0: + group_max = km[0, lane] + for w in cutlass.range_constexpr(1, 8): + group_max = max(group_max, km[w, lane]) + composite = (group_max & Uint32(0xFFFFFFE0)) | Uint32(31 - lane) + gc[lane] = composite + cute.arch.sync_warp() + ranks = zeros(4, Int32) + for q in cutlass.range_constexpr(32): + ranks[q % 4] += Int32(gc[q] > composite) + r = (ranks[0] + ranks[1]) + (ranks[2] + ranks[3]) + least = redux(composite if r < 16 else Uint32(0xFFFFFFFF)) + threshold[0] = least & Uint32(0xFFFFFFE0) + cute.arch.barrier(barrier_id=1, number_of_threads=256) + cutoff = threshold[0] + lane_mask = (Uint32(1) << lane) - Uint32(1) + region = Uint32(128 * warp if warp < 4 else 512 + 96 * (warp - 4)) + offset = Uint32(0) + for i in cutlass.range_constexpr(4): + live = warp + 8 * i < 28 + ballot = cute.arch.vote_ballot_sync(live & (key[i] >= cutoff)) + candidate = ((ballot >> lane) & 1) != 0 + compact = region + offset + cute.arch.popc(ballot & lane_mask) + position = compact if candidate else Uint32(896 + warp) + expert = Uint32(lane + 32 * (warp + 8 * i)) + cc[position] = (Uint64(key[i]) << 32) | Uint64(Uint32(0xFFFFFFFF) - expert) + cs[position] = records[i, 0] + offset += cute.arch.popc(ballot) + wc[warp] = offset + cute.arch.barrier(barrier_id=1, number_of_threads=256) + if warp == 0: + counts = cute.make_rmem_tensor(8, Uint32) + for w in cutlass.range_constexpr(8): + counts[w] = wc[w] + n = ((counts[0] + counts[1]) + (counts[2] + counts[3])) + ( + (counts[4] + counts[5]) + (counts[6] + counts[7]) + ) + if n <= 32: + self.rank_small(arrays, counts, n, lane, ids, 1) + elif n <= 64: + self.rank_small(arrays, counts, n, lane, ids, 2) + else: + prefix = cute.make_rmem_tensor(8, Uint32) + total = Uint32(0) + for w in cutlass.range_constexpr(8): + prefix[w] = total + total += counts[w] + for index in range(lane, n, 32): + slot = Uint32(index) + for w in cutlass.range_constexpr(7): + past = index >= prefix[w + 1] + gap = (128 if w < 4 else 96) - counts[w] + slot += gap if past else Uint32(0) + mine = cc[slot] + rank = Uint32(0) + for w in cutlass.range_constexpr(8): + base = 128 * w if w < 4 else 512 + 96 * (w - 4) + for q in range(counts[w]): + rank += Uint32(cc[base + q] > mine) + if rank < 16: + selected_expert = Int32(Uint32(0xFFFFFFFF) - Uint32(mine)) + selected_ids[rank] = selected_expert + selected_scores[rank] = bits_float(cs[slot]) + mailbox_store(ids + rank, (Uint32(selected_expert),)) + cute.arch.sync_warp() + return selected_scores diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/shared.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/shared.py new file mode 100644 index 0000000000..75e9762074 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/shared.py @@ -0,0 +1,684 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused shared expert for B>8: gate/up, SiTU, down and the routed tail's reduce-scatter +push in one persistent tcgen05 kernel, synchronized by DSMEM and one global-memory +barrier instead of kernel bounds. + +120 CTAs of 160 threads in clusters of ks = 120 / tiles (TP8 5, TP16 10). Loader warps +0-3 feed a 128B-swizzled ring of W [64, 64] (TMA) and x [N, 64] (cp.async) K blocks to +warp 4, which issues each block's four M=64 kind::f16 MMAs warp-uniformly +(``mma_bf16_block``). Gate/up M-tile m (32 gate rows and the matching 32 up rows) +belongs to cluster m, whose CTAs split its K. Each CTA pushes its FP32 partial rows +over DSMEM to the owner of their h columns (``width`` columns per owner); an owner +thread per (column, token) sums the parts in K (rank) order, rounds gate and up to BF16 +for SiTU, and the owner stores h rows as vectors. After the grid's h barrier, CTAs +below 112 run 64 down rows (the first ring's worth of weights already staged) and +peer-store them, -0 made +0, into the tail's RS mailboxes. The last CTA through the +barrier re-arms the counters.""" + +import functools + +import cuda.bindings.driver as cuda +import cutlass as c +import cutlass.cute as cute +import torch +from cutlass import Constexpr, Float32, Int32, Int64, Uint32 +from cutlass.cute.runtime import make_fake_stream + +from ..common.host import fake, tensor_maps +from ..common.primitives import ( + acquire_tensormap, + alloc, + asm, + at, + bits_float, + commit_asm, + copy_g2s, + copy_s2cluster, + div_rn, + fence_tmem, + float_bits, + launch, + load_tma2, + load_tmem, + mma_bf16_block, + multicast_words, + named_barrier, + pack_pair, + push_float, + release_tmem, + sigmoid, + smem_desc, + smem_words, + split, + store_relaxed, + store_words, + thread, + view, + wait, +) +from ..gemm.ws_gemm import positive_zero, positive_zero32 + +tanh = functools.partial(cute.math.tanh, fastmath=False) +GRID, ROWS = 120, 112 # CTAs; CTAs with 64 down rows +# With the MoE front rows: the K split by TP and a deeper ring that still fits beside +# three o-proj consumer CTAs (<= 30 KB each; the kernel launches at the consumer's +# entry and loads its first ring meanwhile). +FRONT_KS, FRONT_RING = {8: 4, 16: 6}, {8: 112 * 1024, 16: 96 * 1024} +COUNTERS = 4 # h ready, exit + + +@cute.jit +def situ(gate, up, beta, linear_beta, *, native: c.Constexpr = False): + if c.const_expr(native): + inv_beta, half = div_rn(c.Float32(1), beta), c.Float32(0.5) + g = half * beta * tanh(gate * inv_beta) * (c.Float32(1) + tanh(gate * half)) + if linear_beta > 0: + up = linear_beta * tanh(up * div_rn(c.Float32(1), linear_beta)) + return g * up + g = beta * tanh(div_rn(gate, beta)) * sigmoid(gate) + if linear_beta > 0: + up = linear_beta * tanh(div_rn(up, linear_beta)) + return g * up + + +@asm() +def prefetch_tma2(descriptor, column, row): + """L2 prefetch of the 2D TMA box at (``column``, ``row``).""" + code = "cp.async.bulk.prefetch.tensor.2d.L2.global.tile [$0,{$1,$2}];" + return code, "l,r,r", (descriptor, Int32(column), Int32(row)) + + +@asm() +def release_count(address): + return "red.release.gpu.global.add.u32 [$0],1;", "l,~{memory}", (address,) + + +@asm(Uint32) +def acquire_count(address): + return "ld.acquire.gpu.global.u32 $0,[$1];", "=r,l,~{memory}", (address,) + + +@asm(Uint32) +def fetch_count(address): + return "atom.relaxed.gpu.global.add.u32 $0,[$1],1;", "=r,l,~{memory}", (address,) + + +@cute.jit +def wait_count(address, target): + seen = Int32(acquire_count(address)) + while seen < target: + seen = Int32(acquire_count(address)) + + +def weight_maps(gate_up, down, front=None): + """TMA maps (W [64, 32] boxes, 128B swizzle) of gate/up, down, then (optional) the + MoE front GEMM's packed router / latent rows.""" + dtype = cuda.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_BFLOAT16 + swizzle = cuda.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_128B + l2 = cuda.CUtensorMapL2promotion.CU_TENSOR_MAP_L2_PROMOTION_L2_256B + weights = (gate_up, down) if front is None else (gate_up, down, front) + specs = [ + (w.data_ptr(), dtype, (w.shape[1], w.shape[0]), (2 * w.stride(0),)) + + ((64, 32), swizzle, l2) + for w in weights + ] + return tensor_maps(specs, gate_up.device) + + +class SharedBase: + """What the fused kernel and the split down kernel share: geometry, TMA weight and + cp.async row tiles, the M=64 accumulator readout and the RS push of 64 rows.""" + + def __init__(self, tp, batch): + self.tp, self.batch, self.hidden = tp, batch, 6144 // tp + self.down = self.hidden // 64 # down K blocks + self.n = (batch + 7) // 8 * 8 + self.columns = max(32, 1 << (self.n - 1).bit_length()) # one accumulator + + @cute.jit + def load_weight(self, tile, descriptor, barrier, row0, row1, column): + """W rows [row0, + 32) and [row1, + 32), 64 columns from ``column``, as a + K-major 128B-swizzled [64, 64] tile (two TMA boxes).""" + cute.arch.mbarrier_arrive_and_expect_tx(barrier, 8192) + load_tma2(tile, descriptor, barrier, column, row0) + load_tma2(tile + 4096, descriptor, barrier, column, row1) + + @cute.jit + def load_rows(self, tile, source, column): + """Rows [0, batch) of ``source``, 64 columns from ``column``, as the B tile + (threads [0, 128)).""" + tid, _, _, _ = thread() + count: Constexpr = self.batch * 8 + for i in c.range_constexpr((count + 127) // 128): + item = tid + 128 * i + row, chunk = item // 8, item % 8 + dst = tile + row // 8 * 1024 + row % 8 * 128 + (chunk ^ row % 8) * 16 + if c.const_expr((i + 1) * 128 <= count): + copy_g2s(dst, at(source, row, column + chunk * 8)) + else: + if item < count: # noqa: SIM102 + copy_g2s(dst, at(source, row, column + chunk * 8)) + + @cute.jit + def accumulated(self, tmem): + """This lane's accumulator row (M=64: rows 16 warp + lane, lane < 16); later + MMAs may then reuse the columns.""" + _, warp, _, _ = thread() + values = load_tmem(tmem + Uint32(32 * warp << 16), self.columns) + fence_tmem(before=True) + return values + + @cute.jit + def elected(self, barrier): + with cute.arch.elect_one(): + commit_asm(barrier) + + @cute.jit + def push(self, mailbox, staging, values, row0): + """Down rows [row0, + 64) (``values`` from ``accumulated``) -> FP32 [token, + row] tile at ``staging`` -> BF16 (-0 made +0) 16-byte peer stores into the RS + mailbox slot; ``mailbox`` is (peers, parity, rank, capacity). Warps 0-3.""" + peers, parity, rank, capacity = mailbox + tid, warp, lane, _ = thread() + tile = view(staging, 0, Float32, (self.n, 68), (68, 1)) + if lane < 16: + for t in c.range_constexpr(self.batch): + tile[t, 16 * warp + lane] = values[t] + named_barrier(1, 128) + shard: Constexpr = 7168 // self.tp + peer, column = row0 // shard, row0 % shard + count: Constexpr = self.batch * 8 + for i in c.range_constexpr((count + 127) // 128): + item = tid + 128 * i + if item < count: + t, part = item // 8, item % 8 + reduced = [] + for half in c.range_constexpr(2): + words = smem_words(at(tile, t, part * 8 + half * 4)) + reduced += [bits_float(w) for w in words] + packed = tuple( + pack_pair(reduced[2 * v], reduced[2 * v + 1], policy="latent") + for v in range(4) + ) + row = (parity * self.tp + rank) * capacity + t + offset = Int64(row * shard + column + part * 8) + store_relaxed(peers[peer] + offset * 2, packed) + named_barrier(1, 128) + + +class SharedUmma(SharedBase): + def __init__(self, tp, batch, budget=80 * 1024, gate=False, front=None): + super().__init__(tp, batch) + # gate: the split variant's front-window half (gate/up and SiTU into h; the + # down rows run in ``SharedDown`` beside the routed experts). front: (rows, + # split, valid, mailbox row) of the MoE front GEMM, whose packed router / + # latent rows join as extra M-tiles (gate mode, on the main stream after the + # o-proj consumer: the first ring's weights load before the dependency wait). + self.gate, self.front = (gate or front is not None, front) + self.tiles = 2 * self.hidden // 64 # gate/up M-tiles, one per cluster + self.grid = GRID + if front is not None: + self.gtiles = self.tiles + self.tiles += -(-front[0] // 64) + self.grid = FRONT_KS[tp] * self.tiles + self.ks = self.grid // self.tiles # cluster size: the K split + self.blocks = -(-112 // self.ks) # gate/up K blocks per CTA, at most + self.stage = 8192 + self.n * 128 + later = 0 if self.gate else self.down + self.stages = min(self.blocks + later, budget // self.stage) + # Owners 0 .. 32 / width - 1 take a tile's h columns ``width`` at a time. + self.width = 8 if self.ks < 8 else 4 + self.owners = 32 // self.width + self.received = self.ks * 2 * self.width * self.n * 4 # an owner's bytes + + @cute.jit + def __call__(self, *args): + options = dict(pdl=self.front is not None, min_blocks_per_mp=1) + launch(self.kernel, args, self.grid, 160, self.ks, **options) + + @cute.kernel + def kernel(self, args): + maps = args[1] + tid, warp, lane, cta = thread() + rank = cta % self.ks + stages, stage = self.stages, self.stage + smem = c.utils.SmemAllocator() + ring = alloc(smem, c.Uint8, stages * stage, 1024).iterator + received = alloc(smem, Float32, self.received // 4, 16).iterator + staged = alloc(smem, c.BFloat16, self.n * self.width, 16) # owned h rows + if c.const_expr(self.front is not None): # owned front rows [token, half] + staged = (staged, alloc(smem, Float32, self.n * 2 * (self.width + 1), 16)) + bars = alloc(smem, c.Uint64, 2 * stages + 3, 8).iterator + slot = alloc(smem, Uint32, 1) + if tid == 0: + for s in c.range_constexpr(stages): + cute.arch.mbarrier_init(bars + s, 129) # TMA + 128 cp.async arrivals + cute.arch.mbarrier_init(bars + stages + s, 1) + for i in c.range_constexpr(3): + cute.arch.mbarrier_init(bars + 2 * stages + i, 1) + cute.arch.mbarrier_init_fence() + # Every rank's gate and up partials of the owned columns. + nbytes = self.received if rank < self.owners else 0 + cute.arch.mbarrier_arrive_and_expect_tx(bars + 2 * stages + 2, nbytes) + acquire_tensormap(maps.iterator) + acquire_tensormap(maps.iterator + 128) + if c.const_expr(self.front is not None): + acquire_tensormap(maps.iterator + 256) + if warp == 4: + cute.arch.alloc_tmem(self.columns, slot.iterator) + cute.arch.relinquish_tmem_alloc_permit() + fence_tmem(before=True) + cute.arch.sync_threads() + fence_tmem() + cute.arch.cluster_arrive_relaxed() + cute.arch.cluster_wait() # every owner's receive barrier is armed + tmem = slot[0] + if c.const_expr(self.front is not None): + # Weights do not depend on the predecessor: the first ring's worth. + if tid == 0: + m, k0, q = self.span(cta) + for kb in c.range_constexpr(stages): + if kb < q: + self.load_tile(ring + kb * stage, maps, bars + kb, m, k0 + kb) + # The rest of the front tiles' boxes (on the routing path) into L2. + if m >= self.gtiles: + for kb in c.range(stages, q, unroll=1): + self.prefetch_tile(maps, m, k0 + kb) + cute.arch.griddepcontrol_wait() + cute.arch.griddepcontrol_launch_dependents() + if warp < 4: + self.stream(args, ring, (received, staged), bars, tmem, cta) + else: + self.multiply(ring, bars, tmem, cta) + release_tmem(tmem, self.columns, warp, 4) + cute.arch.cluster_arrive_relaxed() + cute.arch.cluster_wait() + + @cute.jit + def load_tile(self, tile, maps, barrier, m, k): + """K block ``k`` of M-tile ``m``: 32 gate and the matching 32 up rows, or (front + tiles past the gate/up ones) 64 consecutive front rows.""" + if m < self.gtiles: + rows = (32 * m, self.hidden + 32 * m) + self.load_weight(tile, maps.iterator, barrier, *rows, 64 * k) + else: + row = 64 * (m - self.gtiles) + self.load_weight(tile, maps.iterator + 256, barrier, row, row + 32, 64 * k) + + @cute.jit + def prefetch_tile(self, maps, m, k): + """L2 prefetch of ``load_tile``'s two boxes.""" + if m < self.gtiles: + prefetch_tma2(maps.iterator, 64 * k, 32 * m) + prefetch_tma2(maps.iterator, 64 * k, self.hidden + 32 * m) + else: + row = 64 * (m - self.gtiles) + prefetch_tma2(maps.iterator + 256, 64 * k, row) + prefetch_tma2(maps.iterator + 256, 64 * k, row + 32) + + @cute.jit + def span(self, cta): + """(tile, first K block, K blocks) of the CTA's gate/up work.""" + k0, count = split(112, cta % self.ks, self.ks) + return cta // self.ks, k0, count + + @cute.jit + def multiply(self, ring, bars, tmem, cta): + """Warp 4: every K block's four K16 MMAs and its stage release; gate/up + completes barrier 0 of three, down (the same columns) barrier 1.""" + stages, q = self.stages, self.span(cta)[2] + total = q + if c.const_expr(not self.gate): + total = q + (self.down if cta < ROWS else 0) + instruction: Constexpr = 1 << 4 | 1 << 7 | 1 << 10 | self.n >> 3 << 17 | 4 << 24 + done = bars + 2 * stages + for kb in c.range(total, unroll=1): + s = kb % stages + wait(bars + s, kb // stages % 2, scope="cta") + cute.arch.fence_proxy("async.shared", space="cta") + fence_tmem() + tile = ring + s * self.stage + a, b = smem_desc(tile, True), smem_desc(tile + 8192, True) + index = kb - q # block index within the product + if kb < q: + index = kb + release = Int32(kb + stages < total) + mma_bf16_block(tmem, a, b, instruction, index, bars + stages + s, release) + if kb == q - 1: + self.elected(done) + if c.const_expr(not self.gate): # noqa: SIM102 + if cta < ROWS: + self.elected(done + 1) + + @cute.jit + def stream(self, args, ring, buffers, bars, tmem, cta): + """Loader warps: gate/up blocks, the first down weights, the partial push, + owned h columns, the h barrier, then the down blocks' h (and remaining + weights) and the push.""" + x, maps, h, counters = args[0:4] + stages, stage, hidden = self.stages, self.stage, self.hidden + tid, _, _, _ = thread() + m, k0, q = self.span(cta) + gate_up, down = maps.iterator, maps.iterator + 128 + for kb in c.range(q, unroll=1): + s = kb % stages + tile = ring + s * stage + if kb >= stages: + wait(bars + stages + s, (kb // stages - 1) % 2, scope="cta") + k = k0 + kb + if c.const_expr(self.front is not None): + if (tid == 0) & (kb >= stages): # the first ring went before the wait + self.load_tile(tile, maps, bars + s, m, k) + elif tid == 0: + rows = (32 * m, hidden + 32 * m) + self.load_weight(tile, gate_up, bars + s, *rows, 64 * k) + self.load_rows(tile + 8192, x, 64 * k) + cute.arch.cp_async_mbarrier_arrive_noinc(bars + s) + received, staged = buffers + if c.const_expr(self.gate): + self.partial(received, bars, tmem, cta, ring) + self.reduce(args, received, staged, bars, cta) + return + early = min(self.down, stages) + row0 = 64 * cta + if cta < ROWS: + for j in c.range(early, unroll=1): + s = (q + j) % stages + if q + j >= stages: + wait(bars + stages + s, ((q + j) // stages - 1) % 2, scope="cta") + if tid == 0: + tile = ring + s * stage + self.load_weight(tile, down, bars + s, row0, row0 + 32, 64 * j) + self.partial(received, bars, tmem, cta) + self.reduce(args, received, staged, bars, cta) + if tid == 0: + wait_count(counters.iterator, GRID) + if fetch_count(counters.iterator + 1) == GRID - 1: + counters[0], counters[1] = Int32(0), Int32(0) + named_barrier(1, 128) + if cta < ROWS: + for j in c.range(self.down, unroll=1): + s = (q + j) % stages + tile = ring + s * stage + if j >= early: + wait(bars + stages + s, ((q + j) // stages - 1) % 2, scope="cta") + if tid == 0: + self.load_weight(tile, down, bars + s, row0, row0 + 32, 64 * j) + self.load_rows(tile + 8192, h, 64 * j) + cute.arch.cp_async_mbarrier_arrive_noinc(bars + s) + wait(bars + 2 * self.stages + 1, 0, scope="cta") + fence_tmem() + self.push(args[4:8], ring, self.accumulated(tmem), row0) + + @cute.jit + def partial(self, received, bars, tmem, cta, ring=None): + """Push the CTA's FP32 gate/up rows to their h columns' owners: row r (gate + r, or up r - 32) of column g = r % 32 lands in owner g // width's slot + [rank, r // 32, g % width]. With the front rows, the slots are staged in the + drained ring and each owner's leaves in one DSMEM bulk copy.""" + tid, warp, lane, _ = thread() + rank = cta % self.ks + wait(bars + 2 * self.stages, 0, scope="cta") + fence_tmem() + values = self.accumulated(tmem) + if c.const_expr(self.front is not None): + share: Constexpr = 2 * self.width * self.n # one owner's floats + if lane < 16: + row = 16 * warp + lane + g = row % 32 + owner = g // self.width + cell = ((owner * 2 + row // 32) * self.width + g % self.width) * self.n + for v in c.range_constexpr(self.n // 4): + words = tuple(float_bits(values[4 * v + i]) for i in range(4)) + at_v = cell + self.slot(4 * v, g % self.width) + store_words(ring + at_v * 4, words, shared=True) + cute.arch.fence_proxy("async.shared", space="cta") + named_barrier(1, 128) + if tid < self.owners: + arrived = bars + 2 * self.stages + 2 + source = ring + tid * share * 4 + dst = received + rank * share + copy_s2cluster(dst, source, arrived, tid, size=share * 4) + return + if lane < 16: + row = 16 * warp + lane + g = row % 32 + cell = ((rank * 2 + row // 32) * self.width + g % self.width) * self.n + arrived = bars + 2 * self.stages + 2 # the owner's receive barrier + for v in c.range_constexpr(self.n // 4): + chunk = tuple(values[4 * v + i] for i in range(4)) + push_float(received + cell + 4 * v, chunk, arrived, g // self.width) + + @cute.jit + def slot(self, t, j): + """Token ``t``'s place in a partial row of column ``j``: with the front rows + the row's 16-byte chunks rotate by ``j`` (conflict-free staging stores).""" + if c.const_expr(self.front is None): + return t + chunks: Constexpr = self.n // 4 + return (t // 4 + j) % chunks * 4 + t % 4 + + @cute.jit + def reduce(self, args, received, staged, bars, cta): + """Owner: a thread per (column, token) sums the parts in rank order and + applies SiTU; the rows leave in ``width``-column vectors. Then the CTA counts + into the h barrier (fused only).""" + h, counters = args[2:4] + beta, linear_beta = args[8:10] + tid, _, _, _ = thread() + m, rank = cta // self.ks, cta % self.ks + wait(bars + 2 * self.stages + 2, 0) # peers' DSMEM stores + parts = view(received, 0, Float32, self.received // 4) + if c.const_expr(self.front is None): + if rank < self.owners: + self.owned_h(h, parts, staged, m, rank, beta, linear_beta) + else: + staged, rows = staged + if rank < self.owners: + if m < self.gtiles: + self.owned_h(h, parts, staged, m, rank, beta, linear_beta) + else: + self.publish(args, parts, rows, m - self.gtiles, rank) + named_barrier(1, 128) + if c.const_expr(not self.gate): # noqa: SIM102 + if tid == 0: + cute.arch.fence_acq_rel_gpu() + release_count(counters.iterator) + + @cute.jit + def owned_h(self, h, parts, staged, m, rank, beta, linear_beta): + """Owner of gate/up tile ``m``: SiTU of the summed parts into ``width``-column + h vectors.""" + tid, _, _, _ = thread() + batch, width = self.batch, self.width + for i in c.range_constexpr((width * batch + 127) // 128): + item = tid + 128 * i + if item < width * batch: + j, t = item // batch, item % batch + sums = [Float32(0), Float32(0)] + for source in c.range_constexpr(self.ks): + for half in c.range_constexpr(2): + cell = (source * 2 + half) * width + j + sums[half] = sums[half] + parts[cell * self.n + self.slot(t, j)] + gate = Float32(c.BFloat16(sums[0])) + up = Float32(c.BFloat16(sums[1])) + value = c.BFloat16(situ(gate, up, beta, linear_beta)) + staged[t * width + j] = value + named_barrier(1, 128) + if tid < batch: + row = staged.iterator + tid * width + words = smem_words(row) if width == 8 else smem_words(row, count=2) + store_words(at(h, tid, 32 * m + width * rank), words) + + @cute.jit + def publish(self, args, parts, rows, f, rank): + """Owner of front tile ``f``: sums each (row, token) in rank order, then + publishes ``width``-row vectors like the front GEMM: FP32 router logits (rows + < valid) and BF16 latent columns (rows >= split), -0 made +0, as multicast + stores into the front mailbox (``mc``; this rank's ``slot`` / ``base``).""" + mc, slot, base = args[10:13] + _, split, valid, pitch = self.front + tid, _, _, _ = thread() + batch, width = self.batch, self.width + shard: Constexpr = 3584 // self.tp + for i in c.range_constexpr((width * batch + 127) // 128): + item = tid + 128 * i + if item < width * batch: + j, t = item // batch, item % batch + for half in c.range_constexpr(2): + total = Float32(0) + for source in c.range_constexpr(self.ks): + cell = (source * 2 + half) * width + j + total = total + parts[cell * self.n + self.slot(t, j)] + rows[(t * 2 + half) * (width + 1) + j] = total + named_barrier(1, 128) + if tid < 2 * batch: + t, half = tid // 2, tid % 2 + row = 64 * f + 32 * half + width * rank + values = [rows[tid * (width + 1) + v] for v in range(width)] + if row < valid: + for v in c.range_constexpr(width // 4): + words = tuple(positive_zero32(values[4 * v + i]) for i in range(4)) + offset = t * pitch + 3584 + 2 * (slot + row + 4 * v) + multicast_words(mc + Int64(offset) * 2, words, kind="bf16x2") + elif (row >= split) & (row - split < shard): + pairs = range(0, width, 2) + words = tuple(positive_zero(values[i], values[i + 1]) for i in pairs) + offset = t * pitch + base + row - split + multicast_words(mc + Int64(offset) * 2, words, kind="bf16x2") + + +class SharedDown(SharedBase): + """The split variant's down projection beside the routed experts: ``sms`` + persistent CTAs (one per SM the experts' grid leaves) stage h [N, hidden] once, + take every ``sms``-th 64-row tile, stream its W blocks through a TMA ring (warp 5) + into alternating TMEM accumulators (warp 4) and peer-store each tile like the + fused kernel (warps 0-3). Exactly the fused kernel's down numerics.""" + + def __init__(self, tp, batch, sms, budget=200 * 1024): + super().__init__(tp, batch) + self.sms, self.image = sms, self.down * self.n * 128 # h as B tiles + ring = budget - self.image - self.n * 68 * 4 + self.stages = min(24, ring // 8192) # W [64, 64] K blocks (8 KB) + + @cute.jit + def __call__(self, *args): + options = dict(pdl=False, min_blocks_per_mp=1) + launch(self.kernel, args, self.sms, 192, **options) + + @cute.kernel + def kernel(self, args): + h, maps = args[0:2] + tid, warp, lane, cta = thread() + stages, columns = self.stages, self.columns + smem = c.utils.SmemAllocator() + image = alloc(smem, c.Uint8, self.image, 1024).iterator + ring = alloc(smem, c.Uint8, stages * 8192, 1024).iterator + staging = alloc(smem, c.Uint8, self.n * 68 * 4, 16).iterator + # full [stages], empty [stages], accumulator full [2] / empty [2], h ready + bars = alloc(smem, c.Uint64, 2 * stages + 5, 8).iterator + slot = alloc(smem, Uint32, 1) + acc_full, acc_empty, ready = (bars + 2 * stages + i for i in (0, 2, 4)) + if tid == 0: + for s in c.range_constexpr(stages): + cute.arch.mbarrier_init(bars + s, 1) + cute.arch.mbarrier_init(bars + stages + s, 1) + for i in c.range_constexpr(2): + cute.arch.mbarrier_init(acc_full + i, 1) + cute.arch.mbarrier_init(acc_empty + i, 128) + cute.arch.mbarrier_init(ready, 128) + cute.arch.mbarrier_init_fence() + if warp == 4: + cute.arch.alloc_tmem(2 * columns, slot.iterator) + cute.arch.relinquish_tmem_alloc_permit() + fence_tmem(before=True) + cute.arch.sync_threads() + fence_tmem() + tmem = slot[0] + tiles = (ROWS - cta + self.sms - 1) // self.sms + if warp < 4: + for j in c.range_constexpr(self.down): + self.load_rows(image + j * self.n * 128, h, 64 * j) + cute.arch.cp_async_mbarrier_arrive_noinc(ready) + for i in c.range(tiles, unroll=1): + wait(acc_full + i % 2, i // 2 % 2, scope="cta") + fence_tmem() + values = self.accumulated(tmem + Uint32(i % 2 * columns)) + cute.arch.mbarrier_arrive(acc_empty + i % 2) + self.push(args[2:6], staging, values, 64 * (cta + i * self.sms)) + elif warp == 4: + self.multiply(image, ring, bars, tmem, tiles) + elif lane == 0: + down = maps.iterator + 128 + acquire_tensormap(down) # in the thread that issues its TMAs + for i in c.range(tiles, unroll=1): + row0 = 64 * (cta + i * self.sms) + for j in c.range(self.down, unroll=1): + g = i * self.down + j + s = g % stages + if g >= stages: + wait(bars + stages + s, (g // stages - 1) % 2, scope="cta") + self.load_weight( + ring + s * 8192, down, bars + s, row0, row0 + 32, 64 * j + ) + release_tmem(tmem, 2 * columns, warp, 4) + + @cute.jit + def multiply(self, image, ring, bars, tmem, tiles): + """Warp 4: each tile's K blocks into accumulator half i % 2.""" + stages, columns = self.stages, self.columns + instruction: Constexpr = 1 << 4 | 1 << 7 | 1 << 10 | self.n >> 3 << 17 | 4 << 24 + acc_full, acc_empty, ready = (bars + 2 * stages + i for i in (0, 2, 4)) + total = tiles * self.down + wait(ready, 0, scope="cta") + cute.arch.fence_proxy("async.shared", space="cta") + for i in c.range(tiles, unroll=1): + if i >= 2: + wait(acc_empty + i % 2, (i // 2 - 1) % 2, scope="cta") + fence_tmem() + d = tmem + Uint32(i % 2 * columns) + for j in c.range(self.down, unroll=1): + g = i * self.down + j + s = g % stages + wait(bars + s, g // stages % 2, scope="cta") + fence_tmem() + a = smem_desc(ring + s * 8192, True) + b = smem_desc(image + j * self.n * 128, True) + release = Int32(g + stages < total) + mma_bf16_block(d, a, b, instruction, j, bars + stages + s, release) + self.elected(acc_full + i % 2) + + +@functools.cache +def compile_shared(tp, batch, device, gate=False, front=None): + """Arguments: x, ``weight_maps``, h, counters, RS peers, parity, rank, capacity, + beta, linear beta. ``gate``: the split variant's gate/up + SiTU half; ``front`` + (rows, split, valid, mailbox row): with the MoE front GEMM's rows, whose results + go to the front mailbox (more arguments: multicast address, logit slot, latent + base; ``weight_maps`` with the front's).""" + maps = fake((256 if front is None else 384,), None, c.Uint8, 64) + args: tuple = (fake((batch, 7168)), maps, fake((batch, 6144 // tp))) + args += (fake((COUNTERS,), None, Int32), fake((tp,), None, Int64)) + args += (Int32(0), Int32(0), Int32(0), Float32(4), Float32(25)) + if front is not None: + args += (Int64(0), Int32(0), Int32(0)) + ring = 80 * 1024 if front is None else FRONT_RING[tp] + kernel = SharedUmma(tp, batch, ring, gate, front) + with torch.accelerator.device_index(device): + return cute.compile(kernel, *args, make_fake_stream()) + + +@functools.cache +def compile_down(tp, batch, device, sms): + """Arguments: h, ``weight_maps``, RS peers, parity, rank, capacity.""" + args: tuple = (fake((batch, 6144 // tp)), fake((256,), None, c.Uint8, 64)) + args += (fake((tp,), None, Int64), Int32(0), Int32(0), Int32(0)) + kernel = SharedDown(tp, batch, sms) + with torch.accelerator.device_index(device): + return cute.compile(kernel, *args, make_fake_stream()) + + +def scratch(device): + """Model-owned fused shared-expert counters (zeroed; re-armed by each launch).""" + return torch.zeros(COUNTERS, dtype=torch.int32, device=device) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/__init__.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/chain.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/chain.py new file mode 100644 index 0000000000..35383df906 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/chain.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Small-batch MoE chains: front -> routed experts -> (shared expert) -> polling tail. + +TP8 B1: ``MoeFrontTP8`` -> ``ClusterExpertsTP8`` -> ``MoeSharedTP8`` -> ``MoeTail`` +(front_poll, trigger 3). TP16 B1/B2: ``MoeFrontTP16`` (shared expert inside, on +designated SMs) -> ``ClusterExpertsTP16`` -> ``MoeTail``. The o-projection consumer +publishes the post-attention row into the x_moe mailbox (parity = layer index % 2) and +bumps an epoch hint; the co-resident front has staged its weights by then. MoE tail +polls the G/W/S slots (three parities, 0xffff-armed, re-armed after use) and publishes +the next AttnRes input.""" + +import functools + +import cutlass.cute as cute +import torch +from cutlass import Float32, Float64, Int32, Uint64 +from cutlass.cute.runtime import make_fake_stream + +from ...common.host import fake, layout, stream, to_cute +from ...oproj.prepared import OprojMailbox +from ..tail_runtime import POLL_CAP, compile_tail +from .experts import ClusterExpertsTP8, ClusterExpertsTP16 +from .front_tp8 import MoeFrontTP8, MoeSharedTP8, SharedSituTables +from .front_tp16 import MoeFrontTP16, initialize_designated + +# Kernel class and trailing scalar arguments of each small-batch launch. +_KERNELS = { + "front": (MoeFrontTP8, (Int32,) * 3 + (Uint64, Float32)), + "front16": (MoeFrontTP16, (Int32,) * 3 + (Uint64,) + (Float32,) * 4 + (Float64,)), + "experts_tp16": (ClusterExpertsTP16, (Float32,) * 3), + "experts_tp8": (ClusterExpertsTP8, (Float32,) * 3), + "shared": (MoeSharedTP8, (Int32,) * 2), + "tables": (SharedSituTables, (Float32,) * 2), +} + + +@functools.cache +def _compile(kind, specs, device, *args, **kwargs): + kernel, scalars = _KERNELS[kind] + operands = [fake(*spec) for spec in specs] + [t(0) for t in scalars] + with torch.accelerator.device_index(device): + return cute.compile(kernel(*args, **kwargs), *operands, make_fake_stream()) + + +def _launch(kind, tensors, *args, **kwargs): + specs = tuple(map(layout, tensors)) + code = _compile(kind, specs, tensors[0].device.index, *args, **kwargs) + return code, tuple(map(to_cute, tensors)) + + +def u32(t): + return t.view(torch.uint32) + + +class SmallBatchBox: + """Model-owned small-batch storage, shared serially by every MoE layer.""" + + def __init__(self, group, device, tp, rank): + def armed(shape, fill, dtype=torch.bfloat16): + value = torch.empty(shape, dtype=dtype, device=device) + value.view(torch.int32).fill_(fill) + return value + + zeros = functools.partial(torch.zeros, dtype=torch.int32, device=device) + self.tp, self.rank, self.device = tp, rank, device + self.sms = torch.cuda.get_device_properties(device).multi_processor_count + self.xmoe = armed((2, 1, 4, 7168), -0x80000000) + self.epoch = zeros(256) + self.scores = armed((7168,), -0x80000000, torch.int32) + self.ids = armed((16,), -1, torch.int32) + self.w, self.s = armed((3, POLL_CAP, 16), -1), armed((3, POLL_CAP, 7168), -1) + self.g = armed((3, POLL_CAP * 16, 3584), -1) + self.h2 = OprojMailbox(group, device, tp, rank, (2, POLL_CAP, 3584)) + self.latent = zeros(3584, dtype=torch.bfloat16) + self.hg, self.counter = zeros(6144), zeros(1) + self.tables = {} + self.batches = (1,) if tp == 8 else (1, 2) + if tp == 16: # resident-SM placement mask plus per-batch route/gate scratch + self.designated = zeros(self.sms + 16) + initialize_designated(self.designated) + self.scratch, i32 = {}, torch.int32 + for b in self.batches: + scores = armed((2 * b * 896 * 2,), -0x80000000, i32) + self.scratch[b] = scores, armed((b * 16,), -1, i32), zeros(2 * b * 768) + self.front_views = tuple(self.xmoe[p] for p in range(2)) + self.hint_views = tuple(self.epoch[p * 32 : p * 32 + 224] for p in range(2)) + + def situ_tables(self, beta, linear_beta): + key = float(beta), float(linear_beta) + if key not in self.tables: + table = torch.empty(131072, dtype=torch.float32, device=self.device) + code, args = _launch("tables", (table,)) + code(*args, Float32(beta), Float32(linear_beta), stream(self.device)) + self.tables[key] = table + return self.tables[key] + + +class SmallBatchChain: + """One MoE layer's small-batch chain; MoE tail weights come from the layer's + prepared tail.""" + + def __init__(self, box, moe, experts, tail, layer_idx, beta, linear_beta, prefetch): + self.box, self.tail_weights, self.tail_box = ( + box, + tail.tail_weights, + tail.tail_box, + ) + self.prefetch, self.layer_idx = prefetch, layer_idx + self.parity, self.slot = layer_idx % 2, int(tail.tail_weights.parity) + self.batches, shared = box.batches, moe.shared_experts + router, bias = moe.gate.weight, moe.gate.e_score_correction_bias + gate_up, shared_down = shared.gate_up_proj.weight, shared.down_proj.weight + rows = 3584 // box.tp + down = moe.routed_expert_down_proj.weight.narrow(0, box.rank * rows, rows) + scales = tuple(experts.weights[i].view(torch.uint8).reshape(-1) for i in (1, 3)) + scale, renorm = moe.routed_scaling_factor, moe.moe_renormalize + self.scalars: tuple = (Int32(box.rank), Uint64(box.h2.mc), Float32(scale)) + head: tuple = (router, bias, box.xmoe, u32(box.epoch)) + if box.tp == 8: + tensors = head + (u32(box.scores), box.ids, box.w, down, box.h2.buffer) + tensors += (box.latent, gate_up, u32(box.hg)) + self.front = _launch("front", tensors, grid=box.sms - 8, renorm=renorm) + tables = box.situ_tables(beta, linear_beta) + tensors = (box.xmoe, shared_down, tables, u32(box.hg), box.s) + self.shared = _launch("shared", tensors, grid=box.sms - 96) + tensors = (box.latent, box.ids, experts.maps[256:640].view(3, 128), *scales) + routed, args = _launch( + "experts_tp8", tensors + (box.g[self.slot, :16], box.counter) + ) + self.routed = routed, args + experts.scalars + self.tail = compile_tail(8, 1, box.device.index, "front_poll", 3) + return + maps, self.codes = experts.maps[:256].view(2, 128), {} + for b in self.batches: + scores, ids, hg = box.scratch[b] + tensors = head + (u32(scores), ids, box.w, down, box.h2.buffer, gate_up) + tensors += (u32(hg), shared_down, box.s, box.designated) + front = _launch("front16", tensors, b, grid=box.sms - 7 * b, renorm=renorm) + latent = box.h2.buffer[self.parity, :b].reshape(-1) + g = box.g[self.slot, : b * 16] + routed = _launch("experts_tp16", (latent, ids, maps, *scales, g), b) + tail = compile_tail(16, b, box.device.index, "front_poll", 2) + self.codes[b] = front, routed, tail + lb = linear_beta + self.scalars += (Float32(beta), Float32(1 / beta), Float32(lb), Float64(1 / lb)) + self.experts = experts.scalars + + def run(self, out, consume): + b, s = out.shape[0], stream(self.box.device) + par, slot = Int32(self.parity), Int32(self.slot) + if self.box.tp == 8: + self.front[0](*self.front[1], par, slot, *self.scalars, s) + self.routed[0](*self.routed[1], s) + self.shared[0](*self.shared[1], par, slot, s) + tail = self.tail + else: + (front, front_args), (routed, routed_args), tail = self.codes[b] + front(*front_args, par, slot, *self.scalars, s) + routed(*routed_args, *self.experts, s) + if self.prefetch is not None: + self.prefetch.fork("next_attention", self.layer_idx, b) + self.tail_weights.run( + self.tail_box, self.box.g, self.box.w, self.box.s, tail, b + ) + if consume(out) is not out: + raise RuntimeError("MoE tail output must be consumed from its mailbox ring") + return out diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/experts.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/experts.py new file mode 100644 index 0000000000..3376e77839 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/experts.py @@ -0,0 +1,541 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Small-batch routed experts: ``ClusterExpertsTP16`` (TP16 B1/B2, six-CTA split-K +clusters exchanging quantized hidden rows) and ``ClusterExpertsTP8`` (TP8 B1, one +six-CTA cluster per routing slot).""" + +from typing import Any + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Int64, Uint8, Uint32, Uint64 +from cutlass.cute.arch import store_async_dsmem + +from ...common.primitives import ( + acquire_tensormap_asm, + alloc, + chunk_mma, + commit_asm, + copy_evict, + copy_scale, + fence_tmem, + float_bits, + fp8_bits, + launch, + load_tma3, + load_tmem, + load_volatile, + load_words, + mma_steps, + mx_exponent, + mx_rescale, + pack_fp8, + prefetch_bytes, + prefetch_map, + published, + release_tmem, + require_zero_tmem, + shared_address, + situ, + smem_desc, + split, + store_bf16_pair, + store_pair, + store_shared_words, + thread, + unpack_bf16, + view, + wait, + wait_tmem, + warp_max_abs, +) + + +@cute.jit +def quantize_mx8(packed, row): + """(E8M0 scale, FP8 words) of eight BF16 values; four lanes share a 32-block.""" + values = cute.make_rmem_tensor(8, Float32) + for k in cutlass.range_constexpr(4): + values[2 * k], values[2 * k + 1] = unpack_bf16(packed[row, k]) + amax = Float32(0) + for k in cutlass.range_constexpr(8): + amax = max(amax, abs(values[k])) + for k in (1, 2): + amax = max(amax, cute.arch.shuffle_sync_bfly(amax, k)) + sb = mx_exponent(amax) + rs = mx_rescale(sb) + lo, hi = (pack_fp8([values[4 * h + k] * rs for k in range(4)]) for h in (0, 1)) + return sb, (lo, hi) + + +class _ClusterExperts: + """Shared launch of the six-CTA-cluster expert kernels.""" + + kernel: Any + threads: int + + @cute.jit + def __call__(self, *args): + launch(self.kernel, args, 96, self.threads, cluster=6, min_blocks_per_mp=1) + + +class ClusterExpertsTP16(_ClusterExperts): + """Paired CTAs split FC1 and exchange quantized hidden rows; the TP16 intermediate + has 192 live values padded to 256. IDs [B, 16], x BF16 [B, 3584], output BF16 + [B * 16, 3584]. Polled IDs are re-armed per cluster; the caller re-arms latent rows + after all consumers complete.""" + + threads = 288 + + def __init__(self, batch: int): + """Polled IDs and latent rows; successors are released before the final wait.""" + self.batch = batch + self.off_sfa = 6 * 32768 + self.off_x = self.off_sfa + 6 * 1024 + self.off_h = self.off_x + 7 * 2048 + self.off_xsf = self.off_h + batch * 2048 + self.off_hsf = self.off_xsf + 14 * 128 + 384 + self.off_partial = self.off_hsf + batch * 1152 + self.off_misc = self.off_partial + batch * 256 + self.off_acc1 = self.off_misc + 104 + self.off_acc2 = self.off_acc1 + batch * 8 + self.off_precv = self.off_acc2 + batch * 40 + self.off_hrecv = self.off_precv + batch * 8 + self.off_tmem = self.off_hrecv + batch * 8 + self.off_amax = self.off_tmem + 4 + self.off_stage = (self.off_amax + batch * 16 + 15) // 16 * 16 + # 512 B unused reserve keeps the tuned SMEM footprint. + self.smem_bytes = (self.off_stage + batch * 64 + 512 + 15) // 16 * 16 + cols = (8, 64, 128, 184) if batch == 1 else (16, 96, 144, 200) + self.col_acc2, self.col_sfa, self.col_xsf, self.col_hsf = cols + + @cute.jit + def get_id(self, ids, j, token): + expert = Int32(load_volatile(ids.iterator + token * 16 + j)[0]) + while expert < 0: + expert = Int32(load_volatile(ids.iterator + token * 16 + j)[0]) + return expert + + @cute.jit + def issue(self, sm, maps, s13, s2, full, n, e0, e1, rank, tile0, tiles): + slot = n % 6 + cute.arch.mbarrier_arrive_and_expect_tx(full.iterator + slot, 16384 + 1024) + if n < 7 * self.batch: + expert = e0 if n < 7 else e1 + ks = 7 * (rank % 2) + n % 7 + row = expert * 512 + rank // 2 * 128 + dst = sm.iterator + slot * 32768 + load_tma3(dst, maps.iterator, full.iterator + slot, row, ks * 2) + dst = sm.iterator + self.off_sfa + slot * 1024 + src = s13.iterator + Int64(expert) * 57344 + rank // 2 * 14336 + ks * 1024 + copy_evict(dst, src, size=1024, barrier=full.iterator + slot) + else: + u = n - 7 * self.batch + expert = e0 if u < tiles else e1 + tile = tile0 + u % tiles + dst, desc = sm.iterator + slot * 32768, maps.iterator + 128 + load_tma3(dst, desc, full.iterator + slot, expert * 3584 + tile * 128, 0) + dst = sm.iterator + self.off_sfa + slot * 1024 + src = s2.iterator + Int64(expert) * 28672 + tile * 1024 + copy_evict(dst, src, size=1024, barrier=full.iterator + slot) + + @cute.kernel + def kernel(self, args): + x, ids, maps, s13, s2, output, *situ_args = args + tid, warp, lane, cta = thread() + rank, j = cta % 6, cta // 6 + half, row_tile = rank % 2, rank // 2 + tile0, tiles = split(28, rank, 6) + sm = alloc(cutlass.utils.SmemAllocator(), Uint8, self.smem_bytes, 1024) + full = view(sm.iterator, self.off_misc, Uint64, 6) + empty = view(sm.iterator, self.off_misc + 48, Uint64, 6) + xall = view(sm.iterator, self.off_misc + 96, Uint64, 1) + acc1 = view(sm.iterator, self.off_acc1, Uint64, self.batch) + acc2 = view(sm.iterator, self.off_acc2, Uint64, self.batch * 5) + precv = view(sm.iterator, self.off_precv, Uint64, self.batch) + hrecv = view(sm.iterator, self.off_hrecv, Uint64, self.batch) + tmem_base = view(sm.iterator, self.off_tmem, Uint32, 1) + if tid == 0: + for s in cutlass.range_constexpr(6): + cute.arch.mbarrier_init(full.iterator + s, 1) + cute.arch.mbarrier_init(empty.iterator + s, 1) + cute.arch.mbarrier_init(xall.iterator, 7) + for t in cutlass.range_constexpr(self.batch): + cute.arch.mbarrier_init(acc1.iterator + t, 1) + for i in cutlass.range_constexpr(5): + cute.arch.mbarrier_init(acc2.iterator + 5 * t + i, 1) + cute.arch.mbarrier_init(precv.iterator + t, 1) + cute.arch.mbarrier_init(hrecv.iterator + t, 1) + cute.arch.mbarrier_arrive_and_expect_tx(precv.iterator + t, 256) + cute.arch.mbarrier_arrive_and_expect_tx(hrecv.iterator + t, 288) + cute.arch.mbarrier_init_fence() + acquire_tensormap_asm(maps.iterator) + acquire_tensormap_asm(maps.iterator + 128) + if warp == 2: + cute.arch.alloc_tmem(256, tmem_base.iterator) + cute.arch.relinquish_tmem_alloc_permit() + cute.arch.cluster_arrive_relaxed() + fence_tmem(before=True) + cute.arch.sync_threads() + fence_tmem() + tmem = tmem_base[0] + cute.arch.griddepcontrol_launch_dependents() + + if warp == 0: + if lane == 0: + e0 = e1 = self.get_id(ids, j, 0) + for n in cutlass.range_constexpr(6): + self.issue(sm, maps, s13, s2, full, n, e0, e1, rank, tile0, tiles) + ks = 7 * half + 6 + prefetch_map(maps.iterator, e0 * 512 + row_tile * 128, 2 * ks) + src = s13.iterator + Int64(e0) * 57344 + row_tile * 14336 + ks * 1024 + prefetch_bytes(src, 1024) + if cutlass.const_expr(self.batch == 2): + e1 = self.get_id(ids, j, 1) + for n in cutlass.range(6, self.batch * (7 + tiles), unroll=1): + slot = n % 6 + wait(empty.iterator + slot, ((n // 6) - 1) % 2) + self.issue(sm, maps, s13, s2, full, n, e0, e1, rank, tile0, tiles) + elif warp == 1: + require_zero_tmem(tmem) + adesc = smem_desc(sm.iterator, True) + sfa = smem_desc(sm.iterator + self.off_sfa, False) + xdesc = smem_desc(sm.iterator + self.off_x, True) + xsf = smem_desc(sm.iterator + self.off_xsf, False) + wait_tmem(xall.iterator) + with cute.arch.elect_one(): + for n in cutlass.range_constexpr(14): + copy_scale(self.col_xsf + n * 4, xsf + n * 8) + cute.arch.sync_warp() + for t in cutlass.range(self.batch, unroll=1): + for n in cutlass.range_constexpr(7): + g = 7 * t + n + slot = (t + n) % 6 + wait_tmem(full.iterator + slot, (g // 6) % 2) + with cute.arch.elect_one(): + col = self.col_sfa + 8 * slot + copy_scale(col, sfa + 64 * slot) + copy_scale(col + 4, sfa + 64 * slot + 32) + d, a, b = 8 * t, adesc + 2048 * slot, xdesc + 128 * n + mma_steps(d, a, b, col, self.col_xsf + 8 * n, n != 0, False) + commit_asm(empty.iterator + slot) + if cutlass.const_expr(n == 6): + commit_asm(acc1.iterator + t) + cute.arch.sync_warp() + for t in cutlass.range(self.batch, unroll=1): + wait(hrecv.iterator + t, test=True) + cute.arch.fence_proxy("async.shared", space="cta") + fence_tmem() + hdesc = smem_desc(sm.iterator + self.off_h + t * 2048, True) + hsf = smem_desc(sm.iterator + self.off_hsf + t * 1152, False) + with cute.arch.elect_one(): + for k in cutlass.range_constexpr(6): + copy_scale(self.col_hsf + 24 * t + 4 * k, hsf + k * 8) + cute.arch.sync_warp() + for i in cutlass.range_constexpr(5): + if i < tiles: + g = self.batch * 7 + t * tiles + i + slot = g % 6 + wait_tmem(full.iterator + slot, (g // 6) % 2) + with cute.arch.elect_one(): + col = self.col_sfa + 8 * slot + copy_scale(col, sfa + 64 * slot) + copy_scale(col + 4, sfa + 64 * slot + 32) + d, a = self.col_acc2 + 8 * (5 * t + i), adesc + 2048 * slot + mma_steps(d, a, hdesc, col, self.col_hsf + 24 * t, 0, True) + commit_asm(empty.iterator + slot) + commit_asm(acc2.iterator + 5 * t + i) + cute.arch.sync_warp() + else: + n = warp - 2 + packed = cute.make_rmem_tensor((self.batch, 4), Uint32) + for t in cutlass.range_constexpr(self.batch): + pointer = x.iterator + t * 3584 + 1792 * half + 256 * n + lane * 8 + ready = False + while not ready: + words = load_volatile(pointer, count=4) + ready = published(words) + for k in cutlass.range_constexpr(4): + packed[t, k] = words[k] + block, k = lane // 4, lane % 4 + kc, kb = 16 * n + 2 * block + k // 2, 8 * n + block + for t in cutlass.range_constexpr(self.batch): + sb, words = quantize_mx8(packed, t) + offset = self.off_x + kc // 8 * 1024 + t * 128 + ((kc % 8) ^ t) * 16 + offset += (k % 2) * 8 + store_shared_words(sm.iterator + offset, words) + if k == 0: + sm[self.off_xsf + kb // 4 * 128 + 16 * t + kb % 4] = Uint8(sb) + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_warp() + if lane == 0: + cute.arch.mbarrier_arrive(xall.iterator) + if (warp >= 4) & (warp < 8): + ew = warp % 4 + own = ew // 2 == half + tl = tmem + (Uint32(32 * ew) << 16) + cute.arch.cluster_wait() + partial = view(sm.iterator, self.off_partial, Float32, self.batch * 64) + maxima = view(sm.iterator, self.off_amax, Float32, self.batch * 4) + for t in cutlass.range(self.batch, unroll=1): + wait_tmem(acc1.iterator + t) + value = load_tmem(tl + 8 * t + t) + index = 64 * t + (ew % 2) * 32 + lane + if not own: + dst = cute.recast_ptr(partial.iterator + index, dtype=Uint32) + bits = Uint32(float_bits(value)) + store_async_dsmem(dst, bits, precv.iterator + t, rank ^ 1) + else: + wait(precv.iterator + t, test=True) + value += partial[index] + other = cute.arch.shuffle_sync_bfly(value, 8) + up = (lane & 8) == 0 + hidden = situ(other, value, *situ_args) if up else Float32(0) + amax = warp_max_abs(hidden) + if lane == 0: + maxima[4 * t + ew] = amax + cute.arch.barrier(barrier_id=1, number_of_threads=64) + sb = mx_exponent(max(amax, maxima[4 * t + (ew ^ 1)])) + if up: + byte = fp8_bits(hidden * mx_rescale(sb)) + stage = self.off_stage + 64 * t + 16 * ew + sm[stage + 2 * (lane % 8) + lane // 16] = byte + cute.arch.sync_warp() + if lane < 6: + src = sm.iterator + self.off_stage + 64 * t + 16 * ew + h_words = load_words(src, count=4, shared=True) + kc = 4 * row_tile + ew + h_offset = self.off_h + 2048 * t + kc // 8 * 1024 + h_offset += kc % 8 * 16 + dst = sm.iterator + cute.assume(h_offset, divby=16) + store_async_dsmem(dst, h_words, hrecv.iterator + t, lane) + if ew % 2 == 1: + sw = (Uint32(sb) * Uint32(0x01010101),) * 4 + scale_offset = self.off_hsf + 1152 * t + 128 * rank + dst = sm.iterator + cute.assume(scale_offset, divby=16) + store_async_dsmem(dst, sw, hrecv.iterator + t, lane) + for t in cutlass.range(self.batch, unroll=1): + for i in range(tiles): + wait_tmem(acc2.iterator + 5 * t + i) + value = load_tmem(tl + self.col_acc2 + 8 * (5 * t + i)) + other = cute.arch.shuffle_sync_down(value, 8) + tile = tile0 + i + column = tile * 128 + 32 * ew + 4 * (lane % 8) + lane // 8 + if (lane & 8) == 0: + dst = output.iterator + (t * 16 + j) * 3584 + column + store_pair(dst, value, other) + release_tmem(tmem, 256, warp, 2) + if (tid == 0) & (rank == 0): + for t in cutlass.range(self.batch, unroll=1): + ids[t * 16 + j] = Int32(-1) + if (cta == 0) & (tid == 0): # one-thread exit wait, as in the front + cute.arch.griddepcontrol_wait() + + +def _align_1024(value): + return (value + 1023) // 1024 * 1024 + + +class ClusterExpertsTP8(_ClusterExperts): + """FP4 x MXFP8 FC1 -> SiTU -> MXFP8 -> FC2 of x BF16 [3584] and IDs int32 [16] into + BF16 [16, 3584]. Weights stay in the loaded MoE weights' shuffled layout (three TMA + maps, two scale arrays); mailbox mode consumes/re-arms IDs and the latent row via a + grid ticket.""" + + slots = 6 + slot_bytes = 32768 + off_sfa = slots * slot_bytes + off_x = _align_1024(off_sfa + slots * 1024) + off_xsf = _align_1024(off_x + 27 * 128 + 1024) + off_h = _align_1024(off_xsf + 28 * 128 + 384) + off_hsf = off_h + 3 * 1024 + off_hsc = off_hsf + 3 * 128 + 384 + off_partial = off_hsc + 32 + off_stage = off_partial + 1024 + off_amax = off_stage + 64 + off_bar = off_amax + 16 + off_tmem = off_bar + (2 * slots + 8 + 3) * 8 + smem_bytes = off_tmem + 16 + threads = 256 + + @cute.jit + def issue(self, sm, maps, scales, index, row, chunk, offset, n, full, empty, count): + slot = n % 6 + phase = (n // 6) % 2 + if n >= 6: + wait(empty.iterator + slot, phase ^ 1) + cute.arch.mbarrier_arrive_and_expect_tx(full.iterator + slot, count * 8704) + dst, desc = sm.iterator + slot * 32768, maps.iterator + index * 128 + load_tma3(dst, desc, full.iterator + slot, row, chunk) + dst, src = sm.iterator + self.off_sfa + slot * 1024, scales.iterator + offset + copy_evict(dst, src, size=count * 512, barrier=full.iterator + slot) + + @cute.kernel + def kernel(self, args): + x, ids, maps, s13, s2, output, counter, *situ_args = args + tid, warp, lane, cta = thread() + rank, j = cta % 6, cta // 6 + sm = alloc(cutlass.utils.SmemAllocator(), Uint8, self.smem_bytes, 1024) + full = view(sm.iterator, self.off_bar, Uint64, 6) + empty = view(sm.iterator, self.off_bar + 48, Uint64, 6) + accb = view(sm.iterator, self.off_bar + 96, Uint64, 8) + xbar, hbar = accb.iterator + 8, accb.iterator + 9 + hsc = view(sm.iterator, self.off_hsc, Uint32, 6) + tmem_base = view(sm.iterator, self.off_tmem, Uint32, 1) + t2lo, t2hi = 28 * rank // 6, 28 * (rank + 1) // 6 + if tid == 0: + for s in cutlass.range_constexpr(6): + cute.arch.mbarrier_init(full.iterator + s, 1) + cute.arch.mbarrier_init(empty.iterator + s, 1) + for a in cutlass.range_constexpr(8): + cute.arch.mbarrier_init(accb.iterator + a, 1) + cute.arch.mbarrier_init(xbar, 6) + cute.arch.mbarrier_init(hbar, 1) + cute.arch.mbarrier_arrive_and_expect_tx(hbar, 6 * 68) + cute.arch.mbarrier_init_fence() + for i in cutlass.range_constexpr(3): + acquire_tensormap_asm(maps.iterator + i * 128) + if warp == 1: + cute.arch.alloc_tmem(256, tmem_base.iterator) + cute.arch.relinquish_tmem_alloc_permit() + fence_tmem(before=True) + cute.arch.sync_threads() + cute.arch.cluster_arrive_relaxed() + fence_tmem() + tmem = tmem_base[0] + cute.arch.griddepcontrol_launch_dependents() + cute.arch.griddepcontrol_wait() + if warp == 0: + if lane == 0: + expert = ids[j] + for n in range(14): + row, chunk = expert * 768 + rank * 128, 2 * n + offset = Int64(expert) * 86016 + rank * 14336 + n * 1024 + self.issue(sm, maps, s13, 0, row, chunk, offset, n, full, empty, 2) + n = Int32(14) + for tile in range(t2lo, t2hi): + offset = Int64(expert) * 43008 + tile * 1536 + row = expert * 3584 + tile * 128 + self.issue(sm, maps, s2, 1, row, 0, offset, n, full, empty, 2) + n += 1 + offset += 1024 + self.issue(sm, maps, s2, 2, row, 2, offset, n, full, empty, 1) + n += 1 + elif warp == 1: + adesc = smem_desc(sm.iterator, True) + sfa = smem_desc(sm.iterator + self.off_sfa, False) + xdesc = smem_desc(sm.iterator + self.off_x, True) + xsf = smem_desc(sm.iterator + self.off_xsf, False) + hdesc = smem_desc(sm.iterator + self.off_h, True) + hsf = smem_desc(sm.iterator + self.off_hsf, False) + wait_tmem(xbar) + for n in range(14): + slot = n % 6 + wait_tmem(full.iterator + slot, (n // 6) % 2) + for half in cutlass.range_constexpr(2): + kc = 2 * n + half + empty_address = acc_address = Uint32(0) + if cutlass.const_expr(half == 1): + empty_address = shared_address(empty.iterator + slot) + if n == 13: + acc_address = shared_address(accb.iterator) + a = (adesc + slot * 2048 + half * 1024, xdesc + kc * 8) + sf = (tmem + 64 + 8 * slot + 4 * half, sfa + slot * 64 + half * 32) + sfx = (tmem + 128 + 4 * kc, xsf + kc * 8, 1, Int32(kc != 0)) + chunk_mma(tmem, *a, *sf, *sfx, empty_address, acc_address) + cute.arch.sync_warp() + wait(hbar) + if lane < 3: + word = (hsc[2 * lane] & 0xFFFF) | (hsc[2 * lane + 1] << 16) + dst = sm.iterator + self.off_hsf + lane * 128 + store_shared_words(dst, (word,) * 4) + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_warp() + n = Int32(14) + for tile in range(t2lo, t2hi): + a = 2 + tile - t2lo + for stage in cutlass.range_constexpr(2): + slot = n % 6 + wait_tmem(full.iterator + slot, (n // 6) % 2) + for half in cutlass.range_constexpr(2 if stage == 0 else 1): + kh: cutlass.Constexpr = 2 * stage + half + empty_address = acc_address = Uint32(0) + if cutlass.const_expr(half == (1 if stage == 0 else 0)): + empty_address = shared_address(empty.iterator + slot) + if cutlass.const_expr(kh == 2): + acc_address = shared_address(accb.iterator + a) + ops: tuple = (tmem + 8 * a, adesc + slot * 2048 + half * 1024) + ops += (hdesc + kh * 64, tmem + 64 + 8 * slot + 4 * half) + ops += (sfa + slot * 64 + half * 32, tmem + 240 + 4 * kh) + ops += (hsf + kh * 8, Int32(tile == t2lo), int(kh != 0)) + chunk_mma(*ops, empty_address, acc_address) + cute.arch.sync_warp() + n += 1 + else: + xtid = tid - 64 + packed = cute.make_rmem_tensor((3, 4), Uint32) + for i in cutlass.range_constexpr(3): + if 32 * (warp - 2) + 192 * i < 448: + words = load_words( + x.iterator + 8 * (xtid + 192 * i), count=4, mutable=True + ) + for k in cutlass.range_constexpr(4): + packed[i, k] = words[k] + for i in cutlass.range_constexpr(3): + f = xtid + 192 * i + if 32 * (warp - 2) + 192 * i < 448: + sb, words = quantize_mx8(packed, i) + kc, chunk = f // 16, (f % 16) // 2 + offset = self.off_x + kc * 128 + (chunk ^ kc % 8) * 16 + f % 2 * 8 + store_shared_words(sm.iterator + offset, words) + kb = f // 4 + if f % 4 == 0: + sm[self.off_xsf + (kb // 4) * 128 + kb % 4] = Uint8(sb) + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_warp() + if lane == 0: + cute.arch.mbarrier_arrive(xbar) + if warp >= 4: + ew = warp % 4 + address = tmem + (Uint32(32 * ew) << 16) + up = (lane & 8) == 0 + local = 2 * (lane % 8) + lane // 16 + cute.arch.cluster_wait() + wait_tmem(accb.iterator) + value = load_tmem(address) + other = cute.arch.shuffle_sync_bfly(value, 8) + hidden = situ(other, value, *situ_args) if up else Float32(0) + amax = abs(hidden) + for off in cutlass.range_constexpr(5): + amax = max(amax, cute.arch.shuffle_sync_bfly(amax, 1 << off)) + maxima = view(sm.iterator, self.off_amax, Float32, 4) + if lane == 0: + maxima[ew] = amax + cute.arch.barrier(barrier_id=1, number_of_threads=128) + sb = mx_exponent(max(amax, maxima[ew ^ 1]), True) + if up: + byte = fp8_bits(hidden * mx_rescale(sb)) + sm[self.off_stage + ew * 16 + local] = byte + cute.arch.sync_warp() + if lane == 0: + src = sm.iterator + self.off_stage + ew * 16 + words = load_words(src, count=4, shared=True) + kc = 4 * rank + ew + offset = self.off_h + (kc // 8) * 1024 + (kc % 8) * 16 + sw = Uint32(mx_exponent(max(maxima[0], maxima[1]), True)) + sw |= Uint32(mx_exponent(max(maxima[2], maxima[3]), True)) << 8 + for peer in range(6): + dst = sm.iterator + cute.assume(offset, divby=16) + store_async_dsmem(dst, words, hbar, peer) + if ew == 0: + store_async_dsmem(hsc.iterator + rank, sw, hbar, peer) + for tile in range(t2lo, t2hi): + a = 2 + tile - t2lo + wait_tmem(accb.iterator + a) + value = load_tmem(address + 8 * a) + other = cute.arch.shuffle_sync_down(value, 8) + if up: + column = tile * 128 + 32 * ew + 4 * (lane % 8) + lane // 8 + dst = output.iterator + j * 3584 + column + store_bf16_pair(dst, value, other) + release_tmem(tmem, 256, warp, 1) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/front_tp16.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/front_tp16.py new file mode 100644 index 0000000000..7e80d43048 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/front_tp16.py @@ -0,0 +1,580 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""TP16 B1/B2 MoE front: SM placement, early router/latent-down staging, MMA router +and shared gate/up, FP32 routing, latent multicast/gather and shared SiTU/down rows.""" + +import functools + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass import BFloat16, Float32, Int32, Uint8, Uint32, Uint64 +from cutlass.cute.runtime import make_fake_stream, make_fake_tensor + +from .....ops.cute_dsl.latent_moe_tail.primitives import to_cute +from ...common.primitives import ( + alloc, + asm, + bf16, + copy_bulk, + div_rn, + early_weights, + fma, + gate_bf16_bits, + init_barriers, + launch, + load_f32x4, + mailbox_load, + mailbox_store, + mma_into, + multicast_words, + named_barrier, + prefetch_bytes, + shared_situ, + sigmoid, + sleep_ns, + smem_words, + split, + thread, + unpack_bf16, + view, + wait, + warp_sum, + zeros, +) +from ..routing import Top16Fast, bf16_word, gather_x, score_record, stage_rows + + +@asm() +def sreg(name, dtype=Uint32, *, pure=False): + bits, kind = (64, "l") if dtype is Uint64 else (32, "r") + return f"mov.u{bits} $0,%{name};", f"={kind}", (), dtype, pure + + +@asm(Uint32) +def claim_add(pointer): + return "atom.relaxed.gpu.global.add.u32 $0,[$1],1;", "=r,l,~{memory}", (pointer,) + + +@asm() +def counter_reduce(pointer, value, *, op="add", release=False): + ordering = "release" if release else "relaxed" + code = f"red.{ordering}.gpu.global.{op}.{'b32' if op == 'or' else 'u32'} [$0],$1;" + return code, "l,r,~{memory}", (pointer, Uint32(value)) + + +@asm(Uint32) +def load_arrivals(pointer): + return "ld.acquire.gpu.global.u32 $0,[$1];", "=r,l,~{memory}", (pointer,) + + +class FrontPlacement: + """One claiming thread per CTA; every CTA publishes its arrival.""" + + def __init__(self, batch): + self.batch = batch + self.offset = 8 if batch == 2 else 0 + + @cute.jit + def claim(self, epoch, designated, par, cta): + pointer = epoch.iterator + (2 + par) * 32 + self.offset + sm = sreg("smid", pure=True) + claimed = Int32(-1) + if sm < designated.shape[0] - 16: # noqa: SIM102 - preserve guarded load + if designated[sm] != 0: + claimed = Int32(claim_add(pointer)) + counter_reduce(pointer + 2, Uint32(0xFFFFFFFF) - Uint32(cta), op="max") + if cutlass.const_expr(self.batch == 2): + mask = pointer + 3 + (cta // 32) + counter_reduce(mask, Uint32(1) << (cta % 32), op="or") + counter_reduce(pointer + 1, Uint32(1), release=True) + return claimed + + @cute.jit + def rearm(self, epoch, par): + pointer = epoch.iterator + (2 + (par ^ 1)) * 32 + self.offset + for i in cutlass.range_constexpr(8 if self.batch == 2 else 3): + mailbox_store(pointer + i, (Uint32(0),)) + + @cute.jit + def role(self, epoch, par, cta, claimed, grid): + pointer = epoch.iterator + (2 + par) * 32 + self.offset + count = Int32(0) + role = Int32(-1) + if claimed >= 0: + arrived = load_arrivals(pointer + 1) + while arrived < grid: + sleep_ns(64) + arrived = load_arrivals(pointer + 1) + count = Int32(mailbox_load(pointer)[0]) + first = Uint32(0xFFFFFFFF) - mailbox_load(pointer + 2)[0] + if cutlass.const_expr(self.batch == 1): + role = claimed | (Int32(first == cta) << 16) + else: + order = Int32(0) + for word in cutlass.range_constexpr(5): + bits = mailbox_load(pointer + 3 + word)[0] + if word < cta // 32: + order += Int32(cute.arch.popc(bits)) + elif word == cta // 32: + below = (Uint32(1) << (cta % 32)) - Uint32(1) + order += Int32(cute.arch.popc(bits & below)) + role = order | (Int32(order < 2) << 16) + return role, count + + +class FrontPlacementProbe: + def __init__(self, cluster): + self.cluster = cluster + + @cute.jit + def __call__(self, *args): + grid, smem = args[0].shape[0] // 2, 220 * 1024 + launch(self.kernel, args, grid, 128, self.cluster, pdl=False, smem=smem) + + @cute.kernel + def kernel(self, args): + output, spin = args + tid, _, _, cta = thread() + sm, cluster = sreg("smid", pure=True), sreg("clusterid.x") + start = sreg("globaltimer", Uint64) + if tid == 0: + output[2 * cta] = Int32(sm) + output[2 * cta + 1] = Int32(cluster) + while sreg("globaltimer", Uint64) - start < spin: + pass + + +@functools.cache +def _compile_probe(device, cluster, length): + output = make_fake_tensor(Int32, (length,), (1,), assumed_align=16) + with torch.accelerator.device_index(device): + probe = FrontPlacementProbe(cluster) + return cute.compile(probe, output, Uint64(20000), make_fake_stream()) + + +def plan_designated(groups, *, engine_clusters=16, cluster_size=6): + """Largest-capacity-first plan of the routed-expert clusters over the SM groups.""" + capacities = [len(group) // cluster_size for group in groups] + if sum(capacities) < engine_clusters: + raise ValueError("SM probe found insufficient routed-engine cluster capacity") + order = sorted(range(len(groups)), key=lambda i: (-capacities[i], -len(groups[i]))) + counts = [0] * len(groups) + remaining = engine_clusters + for i in order: + counts[i] = min(capacities[i], remaining) + remaining -= counts[i] + keep = set() + for group, count in zip(groups, counts): + quota = len(group) - cluster_size * count + chosen: list[int] = [] + for sm in sorted(group): + if sm % 2 == 0 and sm + 1 in group and len(chosen) + 2 <= quota: + chosen.extend((sm, sm + 1)) + for sm in sorted(group): + if len(chosen) < quota and sm not in chosen: + chosen.append(sm) + keep.update(chosen) + return keep, counts + + +def initialize_designated(table: torch.Tensor): + """Probe SM groups and fill the model-owned int32[nsm+16] table; returns CPU audit + data. Probe storage is freed before capture. Probes: 1- and 3-wave grids of + 2..16-CTA clusters, 220 KiB SMEM (one CTA per SM), a 20 us global-timer spin.""" + device = table.device + nsm = torch.cuda.get_device_properties(device).multi_processor_count + parent = list(range(4096)) + seen = set() + launches = [] + + def find(sm): + while parent[sm] != sm: + parent[sm] = parent[parent[sm]] + sm = parent[sm] + return sm + + with torch.accelerator.device_index(device.index): + stream = cuda.CUstream(torch.cuda.current_stream(device).cuda_stream) + for cluster in range(2, 17): + for waves in (1, 3): + clusters = max(1, (nsm // cluster) * waves) + output = table.new_full((2 * cluster * clusters,), -1) + try: + compiled = _compile_probe(device.index, cluster, output.numel()) + compiled(to_cute(output), Uint64(20000), stream) + except RuntimeError as e: + launches.append(dict(cluster=cluster, waves=waves, error=str(e))) + break + torch.accelerator.synchronize(device) + members: dict[int, list[int]] = {} + for sm, group in output.view(-1, 2).cpu().tolist(): + if sm >= 0: + members.setdefault(group, []).append(sm) + seen.add(sm) + for group in members.values(): + for sm in group[1:]: + parent[find(sm)] = find(group[0]) + launches.append(dict(cluster=cluster, waves=waves, clusters=clusters)) + if len(seen) != nsm: + raise RuntimeError(f"Placement probe saw {len(seen)} of {nsm} SMs") + group_map: dict[int, list[int]] = {} + for sm in sorted(seen): + group_map.setdefault(find(sm), []).append(sm) + groups = sorted(group_map.values(), key=lambda group: group[0]) + keep, counts = plan_designated(groups) + values = [int(sm in keep) for sm in range(nsm + 16)] + table.copy_(torch.tensor(values, dtype=torch.int32, device=device)) + return dict( + groups=groups, + engine_clusters=counts, + designated=sorted(keep), + launches=launches, + ) + + +class MoeFrontTP16: + def __init__(self, batch: int, *, grid=None, renorm=True): + self.batch = batch + self.grid = grid if grid is not None else 141 if batch == 1 else 134 + self.renorm = renorm + self.shared_bytes = 15 * 14336 + (512 * batch + 128) * 4 + 16 + 8 + self.placement = FrontPlacement(batch) + + @cute.jit + def __call__(self, *args): + launch(self.kernel, args, self.grid, 512, min_blocks_per_mp=1) + + @cute.jit + def gemv(self, rows, nrows, activation, partials, warp, lane): + group, quarter = lane // 4, lane % 4 + acc = zeros((4, 4), Float32) + for iteration in cutlass.range_constexpr(14): + offset = (warp * 448 + iteration * 32 + 8 * quarter) * 2 + weight = x = (Uint32(0),) * 4 + if group < nrows: + weight = smem_words(rows + group * 14336 + offset) + if group < 2 * self.batch: + x = smem_words(activation + (group // 2) * 14336 + offset) + b0 = x[2] if (group & 1) != 0 else x[0] + b1 = x[3] if (group & 1) != 0 else x[1] + chain: cutlass.Constexpr = iteration % 4 + a = (weight[0], weight[2], weight[1], weight[3]) + mma_into(acc[chain, None], a, (b0, b1), bf16=True) + if (quarter < self.batch) & (group < nrows): + chain_sum = lambda c: acc[c, 0] + acc[c, 3] + total = (chain_sum(0) + chain_sum(1)) + (chain_sum(2) + chain_sum(3)) + partials[(warp * 16 + group) * self.batch + quarter] = total + + @cute.jit + def gate_pass(self, slots, rows, start, bars, partials, hg, par, tid): + wait(bars + 1) + warp, lane = tid // 32, tid % 32 + self.gemv(slots + 9 * 14336, rows, slots + 7 * 14336, partials, warp, lane) + cute.arch.sync_threads() + if tid < self.batch * rows: + token, row = tid // rows, tid % rows + total = Float32(0) + for w in cutlass.range_constexpr(16): + total += partials[(w * 16 + row) * self.batch + token] + bits = Uint32(gate_bf16_bits(total)) + pointer = hg.iterator + (par * self.batch + token) * 768 + start + row + mailbox_store(pointer, ((bits << 16) | Uint32(1),)) + + @cute.jit + def poll_flag(self, pointer): + word = mailbox_load(pointer)[0] + while (word & 1) == 0: + if cutlass.const_expr(self.batch == 2): + sleep_ns(128) + word = mailbox_load(pointer)[0] + return word + + @cute.jit + def shared_dot(self, packed, hq, token): + value = Float32(0) + for chunk in cutlass.range_constexpr(3): + w = (*unpack_bf16(packed[chunk, 0]), *unpack_bf16(packed[chunk, 1])) + for j in cutlass.range_constexpr(3, -1, -1): + value = fma(w[j], hq[token][chunk][j], value) + return warp_sum(value, 32) + + @cute.jit + def load_down(self, down, packed, pair, end_pair, lane, fits=True): + for row in cutlass.range_constexpr(4): + for chunk in cutlass.range_constexpr(3): + words = (Uint32(0), Uint32(0)) + if fits & (pair + row // 2 < end_pair): + words = early_weights( + down.iterator + (2 * pair + row) * 384 + 4 * (lane + 32 * chunk) + ) + packed[row, chunk, 0], packed[row, chunk, 1] = words[0], words[1] + + @cute.jit + def down_dots(self, packed, values, hq, token): + for row in cutlass.range_constexpr(4): + values[token, row] = self.shared_dot(packed[row, None, None], hq, token) + + @cute.jit + def lane_word(self, values, token, lane): + lo = values[token, 0] if lane == 0 else values[token, 2] + hi = values[token, 1] if lane == 0 else values[token, 3] + return bf16_word(lo, hi) + + @cute.jit + def store_word(self, output, par3, token, column, word): + row = par3 * output.shape[1] + token + mailbox_store(output.iterator + row * 7168 + 2 * column, (word,)) + + @cute.jit + def stream_pairs(self, down, hq, output, par3, begin_pair, end_pair, warp, lane): + packed = cute.make_rmem_tensor((4, 3, 2), Uint32) + values = cute.make_rmem_tensor((self.batch, 4), Float32) + for pair in range(begin_pair + 2 * warp, end_pair, 32): + self.load_down(down, packed, pair, end_pair, lane) + for token in cutlass.range_constexpr(self.batch): + self.down_dots(packed, values, hq, token) + if (lane < 2) & (pair + lane < end_pair): + word = self.lane_word(values, token, lane) + self.store_word(output, par3, token, pair + lane, word) + + @cute.jit + def shared_outputs(self, down, hsm, output, par3, first, end, warp, lane): + begin_pair, end_pair = (first + 1) // 2, end // 2 + hq = [ + [ + load_f32x4(hsm.iterator + t * 384 + 4 * (lane + 32 * c), shared=True) + for c in range(3) + ] + for t in range(self.batch) + ] + packed = cute.make_rmem_tensor((4, 3, 2), Uint32) + values = cute.make_rmem_tensor((self.batch, 4), Float32) + saved = cute.make_rmem_tensor((3, self.batch), Uint32) + fits = (end_pair - begin_pair) <= 96 + if cutlass.const_expr(self.batch == 2): + for iteration in cutlass.range_constexpr(3): + pair = begin_pair + 2 * warp + 32 * iteration + self.load_down(down, packed, pair, end_pair, lane, fits) + for token in cutlass.range_constexpr(2): + self.down_dots(packed, values, hq, token) + saved[iteration, token] = self.lane_word(values, token, lane) + if cutlass.const_expr(self.batch == 1) or not fits: + self.stream_pairs(down, hq, output, par3, begin_pair, end_pair, warp, lane) + boundary = cute.make_rmem_tensor(self.batch, Uint32) + if ((end & 1) != 0) & (warp == 0): + edge = end // 2 + for row in cutlass.range_constexpr(2): + for chunk in cutlass.range_constexpr(3): + words = early_weights( + down.iterator + (2 * edge + row) * 384 + 4 * (lane + 32 * chunk) + ) + packed[row, chunk, 0], packed[row, chunk, 1] = words[0], words[1] + for token in cutlass.range_constexpr(self.batch): + lo = self.shared_dot(packed[0, None, None], hq, token) + hi = self.shared_dot(packed[1, None, None], hq, token) + boundary[token] = bf16_word(lo, hi) + if cutlass.const_expr(self.batch == 2): + for iteration in cutlass.range_constexpr(3): + pair = begin_pair + 2 * warp + 32 * iteration + if fits & (lane < 2) & (pair + lane < end_pair): + for token in cutlass.range_constexpr(2): + word = saved[iteration, token] + self.store_word(output, par3, token, pair + lane, word) + if ((end & 1) != 0) & (warp == 0) & (lane == 0): + for token in cutlass.range_constexpr(self.batch): + self.store_word(output, par3, token, end // 2, boundary[token]) + + @cute.kernel + def kernel(self, args): + router, bias, xmoe, epoch, scores, ids, weights, down, latent = args[0:9] + gate_up, hg, shared_down, output, designated, par, par3, rank = args[9:17] + multicast_base, scale, situ = args[17], args[18], args[19:23] + tid, warp, lane, cta = thread() + router_first, router_rows = split(896, cta, self.grid) + pair_first, pairs = split(112, cta, self.grid) + gate_first, gate_rows = Int32(0), Int32(0) + if cta >= self.batch: + part, parts = cta - self.batch, self.grid - self.batch + gate_first, gate_rows = split(768, part, parts) + slots = alloc(cutlass.utils.SmemAllocator(), Uint8, self.shared_bytes).iterator + top, size = 15 * 14336, 256 * self.batch + partials = view(slots, top, Float32, size) + latent_partial = view(slots, top + 4 * size, Float32, (16, 8), (8, 1)) + gate_partial = view(slots, top + 4 * size + 512, Float32, size) + bars = cute.recast_ptr(slots + top + 8 * size + 512, dtype=Uint64) + roles = view(slots, top + 8 * size + 528, Int32, 2) + hsm = view(slots, 0, Float32, self.batch * 384) + gate_args = (slots, gate_rows, gate_first, bars, gate_partial, hg, par, tid) + cute.arch.griddepcontrol_launch_dependents() + active_lane = lane < 28 + first_chunk = 56 * warp + (lane if active_lane else 0) + h2_weights = zeros((2, 4 if self.batch == 2 else 2, 4), Uint32) + for row in cutlass.range_constexpr(2): + for chunk in cutlass.range_constexpr(4 if self.batch == 2 else 2): + valid = active_lane & (pairs > 0) + index = first_chunk + 28 * chunk + if cutlass.const_expr(self.batch == 2): + valid = (warp < 7) & (pairs > 0) + index = 32 * warp + lane + 224 * chunk + if valid: + pointer = down.iterator + (2 * pair_first + row) * 7168 + index * 8 + words = early_weights(pointer, count=4) + for j in cutlass.range_constexpr(4): + h2_weights[row, chunk, j] = words[j] + if tid == 0: + init_barriers(bars, 2, async_proxy=True) + claimed = self.placement.claim(epoch, designated, par, cta) + stage_rows(slots, 14336, 0, router, router_first, router_rows, bars) + if cta == 0: + mailbox_store(epoch.iterator + (par ^ 1) * 32, (Uint32(0),)) + self.placement.rearm(epoch, par) + observed = mailbox_load(epoch.iterator + par * 32)[0] + while observed < 28 * self.batch: + sleep_ns(32) + observed = mailbox_load(epoch.iterator + par * 32)[0] + role, count = self.placement.role(epoch, par, cta, claimed, self.grid) + roles[0], roles[1] = role, count + else: + if (tid >= 32) & (tid < 32 + self.batch * router_rows): + offset = tid - 32 + token, row = offset // router_rows, offset % router_rows + index = ((par ^ 1) * self.batch + token) * 896 + router_first + row + mailbox_store(scores.iterator + index * 2, (Uint32(0x80000000),) * 2) + if (tid >= 64) & (tid < 64 + self.batch * gate_rows): + offset = tid - 64 + token, row = offset // max(gate_rows, 1), offset % max(gate_rows, 1) + index = ((par ^ 1) * self.batch + token) * 768 + gate_first + row + hg[index] = Uint32(0) + vectors: cutlass.Constexpr = xmoe.shape[1] * xmoe.shape[2] * 896 + rearm = ((xmoe, vectors, vectors), (latent, 3584, self.batch * 448)) + for tensor, stride, count in rearm: + for vector in range(cta * 511 + tid - 1, count, self.grid * 511): + pointer = tensor.iterator + (par ^ 1) * stride * 8 + vector * 8 + mailbox_store(pointer, (Uint32(0x80000000),) * 4) + my_bias = Float32(0) + if tid < self.batch * router_rows: + my_bias = bias[router_first + tid % router_rows] + cute.arch.sync_threads() + role, count = roles[0], roles[1] + shared_resident = role >= 0 + consumer = shared_resident & ((role & 0x10000) != 0) + dense = role & 0xFFFF + consumer_token = dense if self.batch == 2 else Int32(0) + shared_first, shared_end = Int32(0), Int32(0) + if shared_resident: + shared_first = 7168 * dense // count + shared_end = 7168 * (dense + 1) // count + h2 = (warp, lane, pairs) if self.batch == 2 else None + words_x, words_h2 = gather_x( + xmoe, par, cta, slots + 7 * 14336, first_chunk, active_lane, self.batch, h2 + ) + if cutlass.const_expr(self.batch == 1): + for row in cutlass.range_constexpr(2): + value = Float32(0) + if pairs > 0: + for chunk in cutlass.range_constexpr(2): + for j in cutlass.range_constexpr(4): + w0, w1 = unpack_bf16(h2_weights[row, chunk, j]) + x0, x1 = unpack_bf16(words_x[0, chunk, j]) + value = fma(w0, x0, fma(w1, x1, value)) + value = warp_sum(value, 32) + if lane == 0: + latent_partial[warp, row] = value + elif (warp < 7) & (pairs > 0): + for token in cutlass.range_constexpr(2): + for row in cutlass.range_constexpr(2): + value = Float32(0) + for chunk in cutlass.range_constexpr(4): + for j in cutlass.range_constexpr(4): + w0, w1 = unpack_bf16(h2_weights[row, chunk, j]) + x0, x1 = unpack_bf16(words_h2[token, chunk, j]) + value = fma(w1, x1, fma(w0, x0, value)) + value = warp_sum(value, 32) + if lane == 0: + latent_partial[warp, 2 * token + row] = value + cute.arch.sync_warp() + if cutlass.const_expr(self.batch == 1): + if warp == 1: + cute.arch.barrier(barrier_id=3, number_of_threads=512) + else: + named_barrier(3, 512, arrive=True) + elif (warp == 1) & (pairs > 0): + cute.arch.barrier(barrier_id=3, number_of_threads=224) + elif (warp < 7) & (pairs > 0): + named_barrier(3, 224, arrive=True) + if (warp == 1) & (pairs > 0) & (lane < self.batch): + lo, hi = Float32(0), Float32(0) + for w in cutlass.range_constexpr(16 if self.batch == 1 else 7): + lo += latent_partial[w, 2 * lane] + hi += latent_partial[w, 2 * lane + 1] + word = bf16_word(lo, hi) + offset = par * 8 * 3584 + lane * 3584 + rank * 224 + 2 * pair_first + multicast_words(multicast_base + Uint64(offset * 2), (word,), kind="b32") + if warp == 1: + cute.arch.sync_warp() + if tid == 0: + if gate_rows > 0: + cute.arch.mbarrier_arrive_and_expect_tx(bars + 1, gate_rows * 14336) + dst, src = slots + 9 * 14336, gate_up.iterator + gate_first * 7168 + copy_bulk(dst, src, bars + 1, gate_rows * 14336) + if shared_resident & (shared_end > shared_first): + src = shared_down.iterator + shared_first * 384 + prefetch_bytes(src, (shared_end - shared_first) * 768) + wait(bars) + self.gemv(slots, router_rows, slots + 7 * 14336, partials, warp, lane) + cute.arch.sync_threads() + if tid < self.batch * router_rows: + token, row = tid // router_rows, tid % router_rows + logit = Float32(0) + for w in cutlass.range_constexpr(16): + logit += partials[(w * 16 + row) * self.batch + token] + score = sigmoid(logit) + record = score_record(score, score + my_bias) + index = (par * self.batch + token) * 896 + router_first + row + mailbox_store(scores.iterator + index * 2, record) + if (gate_rows > 0) & (~consumer): + self.gate_pass(*gate_args) + cute.arch.sync_threads() + if consumer & (warp < 8): + top16 = Top16Fast(self.batch) + selected = top16.run(scores, ids, slots, par, consumer_token) + if warp == 0: + score = selected[lane] if lane < 16 else Float32(0) + total = warp_sum(score, 16) + if lane < 16: + normalized = score + if cutlass.const_expr(self.renorm): + # Like the TP8 front and vLLM's grouped_topk kernels. + normalized = div_rn(score, total + Float32(1e-20)) + bits = Uint32(gate_bf16_bits(normalized * scale)) + if bits == 0x8000: + bits = Uint32(0) + high = cute.arch.shuffle_sync_down(bits, 1, mask=0xFFFF) + if (lane & 1) == 0: + slot = par3 * weights.shape[1] + consumer_token + pointer = weights.iterator + slot * 16 + lane + mailbox_store(pointer, (bits | (high << 16),)) + cute.arch.sync_warp() + if consumer & (gate_rows > 0): + cute.arch.sync_threads() + self.gate_pass(*gate_args) + if shared_resident: + if consumer: + cute.arch.sync_threads() + for element in range(tid, self.batch * 384, 512): + token, column = element // 384, element % 384 + gp = hg.iterator + (par * self.batch + token) * 768 + column + gate, up = self.poll_flag(gp), self.poll_flag(gp + 384) + result = shared_situ(bf16(gate, True), bf16(up, True), *situ) + hsm[element] = Float32(BFloat16(result)) + cute.arch.sync_threads() + self.shared_outputs( + shared_down, hsm, output, par3, shared_first, shared_end, warp, lane + ) + cute.arch.sync_threads() + # One-thread exit wait: completion then covers the o-proj consumer (and its + # mailbox re-arm) for later waits. + if (cta == 0) & (tid == 0): + cute.arch.griddepcontrol_wait() diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/front_tp8.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/front_tp8.py new file mode 100644 index 0000000000..47a8a0f7c9 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/front_tp8.py @@ -0,0 +1,361 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""TP8 B1 small-batch front (full router staged in SMEM, x_moe polling, threshold top16, +latent multicast and gather, shared gate/up) and the shared-expert finisher.""" + +import cutlass +import cutlass.cute as cute +from cutlass import BFloat16, Float32, Int32, Uint8, Uint32, Uint64 + +from ...common.primitives import ( + alloc, + bf16, + copy_bulk, + copy_g2s, + div_rn, + early_weights, + fma, + gate_bf16_bits, + init_barriers, + launch, + load_f32x4, + mailbox_load, + mailbox_store, + mma_into, + multicast_words, + named_barrier, + published, + routing_add, + routing_sigmoid, + routing_weight, + sigmoid, + smem_words, + split, + store_gate_mailbox, + store_relaxed, + table_load, + thread, + unpack_bf16, + view, + wait, + warp_sum, + zeros, +) +from ..routing import Top16Fast, bf16_word, gather_x, score_record, stage_rows + + +class MoeSharedTP8: + ring_bytes = 57344 + row_bytes = 14400 + lsu_rows = 92 + shared_bytes = ring_bytes + row_bytes + lsu_rows * 1536 + 4096 + 24 + + def __init__(self, grid=52): + self.grid = grid + + @cute.jit + def __call__(self, *args): + launch(self.kernel, args, self.grid, 512, min_blocks_per_mp=1) + + @cute.jit + def stage(self, dst, down, low, high, barrier): + cute.arch.mbarrier_arrive_and_expect_tx(barrier, (high - low) * 1536) + if high > low: + copy_bulk(dst, down.iterator + low * 768, barrier, (high - low) * 1536) + + @cute.jit + def rows_out(self, base, low, start, high, hsm, output, par3, lane): + packed = cute.make_rmem_tensor((6, 2), Uint32) + for row in range(start, high, 16): + for chunk in cutlass.range_constexpr(6): + pointer = base + (row - low) * 1536 + (lane + 32 * chunk) * 8 + values = smem_words(pointer, count=2) + packed[chunk, 0], packed[chunk, 1] = values[0], values[1] + self.row_out(row, packed, hsm, output, par3, lane) + + @cute.jit + def row_out(self, row, packed, hsm, output, par3, lane): + value = Float32(0) + for chunk in cutlass.range_constexpr(6): + h = load_f32x4(hsm.iterator + 4 * (lane + 32 * chunk), shared=True) + w0, w1 = unpack_bf16(packed[chunk, 0]) + w2, w3 = unpack_bf16(packed[chunk, 1]) + value = fma(w0, h[0], fma(w1, h[1], fma(w2, h[2], fma(w3, h[3], value)))) + value = warp_sum(value, 32) + if lane == 0: + pointer = output.iterator + par3 * output.shape[1] * 7168 + row + store_gate_mailbox(pointer, gate_bf16_bits(value)) + + @cute.kernel + def kernel(self, args): + xmoe, down, tables, hg, output, par, par3 = args[0:7] + tid, warp, lane, cta = thread() + first = 7168 * cta // self.grid + end = 7168 * (cta + 1) // self.grid + first_end = min(end, first + 37) + second_end = min(end, first_end + 9) + lsu_count = min(92, end - second_end) + storage = alloc(cutlass.utils.SmemAllocator(), Uint8, self.shared_bytes, 1024) + ring = storage.iterator + xslot = ring + self.ring_bytes + lsu = xslot + self.row_bytes + partial = lsu + self.lsu_rows * 1536 + hsm = view(partial, 1024, Float32, 768) + bars = cute.recast_ptr(partial + 4096, dtype=Uint64) + cute.arch.griddepcontrol_launch_dependents() + if tid == 0: + init_barriers(bars, 3, async_proxy=True) + self.stage(ring, down, first, first_end, bars + 1) + active = lane < 28 + gather_x(xmoe, par, cta, xslot, 56 * warp + (lane if active else 0), active) + cute.arch.sync_warp() + if lsu_count > 0: + for vector in range(tid, lsu_count * 96, 512): + dst = lsu + vector * 16 + copy_g2s(dst, down.iterator + second_end * 768 + vector * 8) + cute.arch.cp_async_commit_group() + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_threads() + if tid == 0: + cute.arch.fence_proxy("async.shared", space="cta") + self.stage(xslot, down, first_end, second_end, bars + 2) + if tid < 384: + column = 2 * tid + gate = up = (Uint32(0), Uint32(0)) + valid = False + while not valid: + gate = mailbox_load(hg.iterator + par * 2 * 1536 + column, count=2) + up = mailbox_load(hg.iterator + par * 2 * 1536 + 768 + column, count=2) + valid = ((gate[0] & gate[1] & up[0] & up[1]) & Uint32(1)) != 0 + ga0 = table_load(tables.iterator + (gate[0] >> 16)) + ua0 = table_load(tables.iterator + 65536 + (up[0] >> 16)) + ga1 = table_load(tables.iterator + (gate[1] >> 16)) + ua1 = table_load(tables.iterator + 65536 + (up[1] >> 16)) + hsm[column] = Float32(BFloat16(ga0 * ua0)) + hsm[column + 1] = Float32(BFloat16(ga1 * ua1)) + cute.arch.sync_threads() + out = (hsm, output, par3, lane) + if lsu_count > 0: + cute.arch.cp_async_wait_group(0) + cute.arch.sync_threads() + self.rows_out( + lsu, second_end, second_end + warp, second_end + lsu_count, *out + ) + parts = ((ring, first, first_end), (xslot, first_end, second_end)) + for part, (base, low, high) in enumerate(parts): + if high > low: + if warp == 0: + wait(bars + 1 + part) + cute.arch.sync_threads() + self.rows_out(base, low, low + warp, high, *out) + if cta == 0: + cute.arch.griddepcontrol_wait() + + +class SharedSituTables: + """Separate FP32 gate/up lookup tables over all BF16 patterns.""" + + @cute.jit + def __call__(self, *args): + launch(self.kernel, args, 256, 256, pdl=False) + + @cute.kernel + def kernel(self, args): + tables, beta, linear_beta = args[0:3] + tid, _, _, cta = thread() + index = 256 * cta + tid + value = bf16(Uint32(index)) + gate = div_rn(value, beta) + probability = sigmoid(value) + tables[index] = (beta * cute.math.tanh(gate, fastmath=False)) * probability + up = value + if linear_beta > 0: + scaled = div_rn(value, linear_beta) + up = linear_beta * cute.math.tanh(scaled, fastmath=False) + tables[65536 + index] = up + + +class MoeFrontTP8: + row_bytes = 14400 + shared_bytes = 15 * row_bytes + (16 * 16 + 16 * 4) * 4 + 24 + + def __init__(self, grid=140, *, renorm=True): + self.grid = grid + self.renorm = renorm + + @cute.jit + def __call__(self, *args): + launch(self.kernel, args, self.grid, 512, min_blocks_per_mp=1) + + @cute.jit + def gemv(self, rows, nrows, activation, partials, warp, lane): + group, quarter = lane // 4, lane % 4 + acc = zeros((8, 4), Float32) + for iteration in cutlass.range_constexpr(14): + offset = (warp * 448 + iteration * 32 + 8 * quarter) * 2 + lo = hi = x = (Uint32(0),) * 4 + if group < nrows: + lo = smem_words(rows + group * self.row_bytes + offset) + if group + 8 < nrows: + hi = smem_words(rows + (group + 8) * self.row_bytes + offset) + if group < 1: + x = smem_words(activation + offset) + for half in cutlass.range_constexpr(2): + chain: cutlass.Constexpr = (2 * iteration + half) % 8 + a = (lo[2 * half], hi[2 * half], lo[2 * half + 1], hi[2 * half + 1]) + mma_into(acc[chain, None], a, (x[2 * half], x[2 * half + 1]), bf16=True) + if quarter == 0: + for half in cutlass.range_constexpr(2): + comp: cutlass.Constexpr = 2 * half + partials[warp * 16 + group + 8 * half] = ( + (acc[0, comp] + acc[1, comp]) + (acc[2, comp] + acc[3, comp]) + ) + ((acc[4, comp] + acc[5, comp]) + (acc[6, comp] + acc[7, comp])) + + @cute.jit + def gate_pass(self, rows, count, start, barrier, ctx): + activation, partials, hg, par, tid, warp, lane = ctx + wait(barrier) + self.gemv(rows, count, activation, partials, warp, lane) + cute.arch.sync_threads() + if tid < count: + value = Float32(0) + for w in cutlass.range_constexpr(16): + value += partials[w * 16 + tid] + bits = Uint32(gate_bf16_bits(value)) + pointer = hg.iterator + (par * 2) * 1536 + start + tid + mailbox_store(pointer, ((bits << 16) | Uint32(1),)) + cute.arch.sync_threads() + + @cute.kernel + def kernel(self, args): + router, bias, xmoe, epoch, scores, ids, weights, down = args[:8] + latent, latent_out, gate_up, hg, par, par3, rank, multicast_base = args[8:16] + scale = args[16] + tid, warp, lane, cta = thread() + first_router, router_rows = split(896, cta, self.grid) + first_pair, pairs = split(224, cta, self.grid) + gu_start, gu_rows = Int32(0), Int32(0) + if cta > 0: + gu_start, gu_rows = split(1536, cta - 1, self.grid - 1) + gu_first = min(gu_rows, 7) + gu_second = gu_rows - gu_first + storage = alloc(cutlass.utils.SmemAllocator(), Uint8, self.shared_bytes, 128) + slots = storage.iterator + activation = slots + 7 * self.row_bytes + dst = (slots, self.row_bytes) + top = 15 * self.row_bytes + partials = view(slots, top, Float32, 256) + latent_partial = view(slots, top + 1024, Float32, (16, 4), (4, 1)) + bars = cute.recast_ptr(slots + top + 1280, dtype=Uint64) + cute.arch.griddepcontrol_launch_dependents() + active_lane = lane < 28 + first_chunk = 56 * warp + (lane if active_lane else 0) + second_chunk = first_chunk + 28 + latent_weights = zeros((4, 2, 4), Uint32) + for row in cutlass.range_constexpr(4): + if active_lane & (row < 2 * pairs): + for chunk in cutlass.range_constexpr(2): + col = first_chunk if chunk == 0 else second_chunk + pointer = down.iterator + (2 * first_pair + row) * 7168 + 8 * col + words = early_weights(pointer, count=4) + for j in cutlass.range_constexpr(4): + latent_weights[row, chunk, j] = words[j] + cute.arch.cp_async_commit_group() + if tid == 0: + init_barriers(bars, 3, async_proxy=True) + stage_rows(*dst, 0, router, first_router, router_rows, bars) + if cta == 0: + mailbox_store(epoch.iterator + (par ^ 1) * 32, (Uint32(0),)) + observed = mailbox_load(epoch.iterator + par * 32)[0] + while observed < 28: + observed = mailbox_load(epoch.iterator + par * 32)[0] + else: + for i in range(tid - 1, 2 * router_rows, 511): + token, row = i // router_rows, i % router_rows + index = ((par ^ 1) * 2 + token) * 896 + first_router + row + mailbox_store(scores.iterator + index * 2, (Uint32(0x80000000),) * 2) + if cta > 0: + for i in range(tid - 1, 2 * gu_rows, 511): + token, row = i // gu_rows, i % gu_rows + hg[(par ^ 1) * 2 * 1536 + token * 1536 + gu_start + row] = Uint32(0) + slot_vectors: cutlass.Constexpr = xmoe.shape[1] * xmoe.shape[2] * 896 + for i in range(cta * 511 + tid - 1, slot_vectors, self.grid * 511): + pointer = xmoe.iterator + (par ^ 1) * slot_vectors * 8 + i * 8 + mailbox_store(pointer, (Uint32(0x80000000),) * 4) + my_bias = Float32(0) + if tid < router_rows: + my_bias = bias[first_router + tid] + cute.arch.sync_threads() + input_words, _ = gather_x(xmoe, par, cta, activation, first_chunk, active_lane) + for row in cutlass.range_constexpr(4): + if row < 2 * pairs: + total = Float32(0) + for chunk in cutlass.range_constexpr(2): + for j in cutlass.range_constexpr(4): + w0, w1 = unpack_bf16(latent_weights[row, chunk, j]) + x0, x1 = unpack_bf16(input_words[0, chunk, j]) + total = fma(w0, x0, fma(w1, x1, total)) + total = warp_sum(total, 32) + if lane == 0: + latent_partial[warp, row] = total + cute.arch.sync_warp() + if warp == 1: + cute.arch.barrier(barrier_id=3, number_of_threads=512) + if lane < pairs: + lo, hi = Float32(0), Float32(0) + for w in cutlass.range_constexpr(16): + lo += latent_partial[w, 2 * lane] + hi += latent_partial[w, 2 * lane + 1] + packed = bf16_word(lo, hi) + offset = par * 8 * 3584 * 2 + rank * 896 + (first_pair + lane) * 4 + multicast_words(multicast_base + Uint64(offset), (packed,), kind="b32") + else: + named_barrier(3, 512, arrive=True) + if tid == 0: # noqa: SIM102 + if gu_first > 0: + stage_rows(*dst, 8, gate_up, gu_start, gu_first, bars + 1) + wait(bars) + self.gemv(slots, router_rows, activation, partials, warp, lane) + cute.arch.sync_threads() + if tid == 0: # noqa: SIM102 + if gu_second > 0: + cute.arch.fence_proxy("async.shared", space="cta") + stage_rows(*dst, 0, gate_up, gu_start, gu_second, bars + 2, gu_first) + if tid < router_rows: + logit = Float32(0) + for w in cutlass.range_constexpr(16): + logit += partials[w * 16 + tid] + score = routing_sigmoid(logit) + record = score_record(score, routing_add(score, my_bias)) + index = par * 2 * 896 + first_router + tid + mailbox_store(scores.iterator + index * 2, record) + if cta == 0: + if warp < 8: + selected = Top16Fast(2).run(scores, ids, slots, par, 0) + if warp == 0: + score = selected[lane] if lane < 16 else Float32(0) + total = warp_sum(score, 16) + if lane < 16: + value = routing_weight(score, total, scale, renorm=self.renorm) + pointer = weights.iterator + par3 * weights.shape[1] * 16 + lane + store_gate_mailbox(pointer, gate_bf16_bits(value)) + cute.arch.sync_warp() + for vector in range(tid, 448, 256): + pointer = latent.iterator + par * 8 * 3584 + vector * 8 + words = mailbox_load(pointer, count=4, scope="sys") + ready = published(words) + while not ready: + words = mailbox_load(pointer, count=4, scope="sys") + ready = published(words) + mailbox_store(latent_out.iterator + vector * 8, words) + store_relaxed(pointer, (Uint32(0x80000000),) * 4) + cute.arch.griddepcontrol_wait() + else: + cute.arch.sync_threads() + ctx = (activation, partials, hg, par, tid, warp, lane) + if gu_first > 0: + self.gate_pass( + slots + 8 * self.row_bytes, gu_first, gu_start, bars + 1, ctx + ) + if gu_second > 0: + self.gate_pass(slots, gu_second, gu_start + gu_first, bars + 2, ctx) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/tail.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/tail.py new file mode 100644 index 0000000000..0ae448631d --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/tail.py @@ -0,0 +1,1031 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Clustered MoE tail (TP8 B1..4, TP16 B1..8): routed finalization/publication, +all-reduce, RMS tree, K-split up projection, shared reduce-scatter and output multicast. +TMA weight loads overlap the mailbox wait. The caller owns the three rotating mailboxes +and must use different parities on successive calls. TP8 B5+ and TP16 B9+ run the +one-wave WaveMoeTail over latent rows published by the routed experts and shards +pushed by WaveMoeTail.push.""" + +from functools import reduce +from operator import or_ + +import cutlass +import cutlass.cute as cute +from cutlass import BFloat16, Float32, Int32, Int64, Uint8, Uint32, Uint64 + +from ..common.primitives import ( + acquire_tensormap, + alloc, + bf16, + copy_g2s, + fence_tmem, + fma, + fma_bf16, + front_dirty, + init_barriers, + inv_rms, + launch, + load_f32x4, + load_tma2, + load_tma3, + load_tmem, + load_volatile, + load_words, + mailbox_store, + mma, + mma_bf16_block, + mul2, + multicast_words, + no_neg_zero, + pack_pair, + push_float, + release_tmem, + round_pair_square, + scatter_reduce, + sleep_ns, + smem_desc, + smem_words, + store_relaxed, + store_shared_words, + store_words, + thread, + view, + wait, + wait_tmem, + zeros, +) + +shuffle = cute.arch.shuffle_sync + + +def scatter_sums(values, lane, width): + """``warp_sum(values[g], width)`` of every g (4 or 8 values) as a reduce-scatter: + each butterfly level keeps half of the values and receives the partner's other + half, so every partial sum (and the total) is the butterfly's. Returns the lane's + total; value g's sits in lanes g << (log2(width) - log2(len(values))).""" + values, offset = list(values), width // 2 + while len(values) > 1: + half, upper = len(values) // 2, (lane & offset) != 0 + kept = [] + for j in range(half): + low, high = values[j], values[half + j] + keep = Float32(cutlass.select_(upper, high, low)) + send = Float32(cutlass.select_(upper, low, high)) + kept.append(keep + cute.arch.shuffle_sync_bfly(send, offset)) + values, offset = kept, offset // 2 + total = values[0] + while offset >= 1: + total += cute.arch.shuffle_sync_bfly(total, offset) + offset //= 2 + return total + + +def gather_sums(total, count, width, padded): + """Every lane's copy of ``scatter_sums``' first ``count`` totals.""" + shift = width.bit_length() - padded.bit_length() + return [shuffle(total, g << shift) for g in range(count)] + + +def spread(pointer, lane, swap): + first = smem_words(pointer + 8 * (2 * lane + swap)) + return first, smem_words(pointer + 8 * (2 * lane + 1 - swap)) + + +class MoeTail: + """One MoE tail launch over caller-owned persistent buffers. ``front_poll`` inputs + have three parity slots armed with 0xffff BF16; latent/shared mailboxes [3, TP, + capacity, 3584 or 7168/TP] are armed with 0x80000000. Only complete chains may use + the early B1 trigger, which waits for the predecessor before releasing + successors.""" + + threads = 128 + options: dict = {} + + def __init__(self, tp: int, batch: int, *, mode="front", trigger=2): + self.tp, self.batch, self.mode, self.trigger = tp, batch, mode, trigger + self.cluster, self.tokens = (7 if tp == 8 else 8), min(batch, 2) + self.poll_ns, self.front_poll_ns = (0, 0) if tp == 8 else (64, 128) + self.cols, self.shard = 3584 // self.cluster, 7168 // tp + self.frags, self.nc = self.cols // 8, self.shard // 32 + self.grid = self.cluster * self.nc * ((batch + self.tokens - 1) // self.tokens) + + @cute.jit + def __call__(self, *args): + common = dict(cluster=self.cluster, preferred_smem_carveout=100) + launch(self.kernel, args, self.grid, self.threads, **common, **self.options) + + @cute.jit + def poll(self, base, stride, second, second_stride, dual): + values = cute.make_rmem_tensor((16, 4), Uint32) + sources = (base, stride, second, second_stride) + pending = Uint32(0) + for source in cutlass.range_constexpr(16): + if self.tp == 16 or source < 8 or dual: + pending |= self.poll_source(values, source, sources) + while pending != 0: + if cutlass.const_expr(self.poll_ns): + sleep_ns(Uint32(self.poll_ns)) + next_pending = Uint32(0) + for source in cutlass.range_constexpr(16): + if (pending & (Uint32(1) << source)) != 0: + next_pending |= self.poll_source(values, source, sources) + pending = next_pending + return values + + @cute.jit + def poll_source(self, values, source: cutlass.Constexpr[int], sources): + base, stride, second, second_stride = sources + if cutlass.const_expr(self.tp == 16 or source < 8): + word = load_volatile(base + source * stride, count=4) + else: + word = load_volatile(second + (source - 8) * second_stride, count=4) + for k in cutlass.range_constexpr(4): + values[source, k] = word[k] + return Uint32(reduce(or_, (v == 0x80000000 for v in word))) << source + + @cute.jit + def sum_sources(self, values, start: cutlass.Constexpr[int]): + acc = zeros(8, Float32) + for source in cutlass.range_constexpr(self.tp): + for k in cutlass.range_constexpr(8): + acc[k] = acc[k] + bf16(values[source + start, k // 2], k % 2) + return acc + + def poll_strides(self, io): + if self.mode != "front_poll": + return 0, 0, 0 + return tuple(t.shape[0] // 3 for t in io[2:5]) + + @cute.jit + def front(self, io, latent_mc, par, rank, capacity, block, warp, lane): + _, _, g, weights, s, peers = io + if cutlass.const_expr(self.mode == "front"): + cute.arch.griddepcontrol_wait() + gs, ws, ss = self.poll_strides(io) + start = block + self.grid * warp + for task in cutlass.range(start, self.batch * 168, self.grid * 4, unroll=1): + if task < self.batch * 56: + token, line = task // 56, task % 56 + src = g.iterator + par * gs + token * 16 * 3584 + 64 * line + 2 * lane + wp = weights.iterator + par * ws + token * 16 + w = cute.make_rmem_tensor(8, Uint32) + rows = cute.make_rmem_tensor(16, Uint32) + dirty = Int32(1) + while dirty != 0: + for h in cutlass.range_constexpr(2): + for k, v in enumerate(load_volatile(wp + 8 * h, count=4)): + w[4 * h + k] = v + for j in cutlass.range_constexpr(16): + rows[j] = load_volatile(src + j * 3584)[0] + if cutlass.const_expr(self.mode == "front_poll"): + words = [w[k] for k in range(8)] + [rows[j] for j in range(16)] + bad = reduce(or_, map(front_dirty, words)) + dirty = Int32(cute.arch.vote_any_sync(bad)) + if self.front_poll_ns > 0 and dirty != 0: + sleep_ns(Uint32(self.front_poll_ns)) + else: + dirty = 0 + a0, a1 = Float32(0), Float32(0) + for j in cutlass.range_constexpr(16): + weight = bf16(w[j // 2], j % 2) + a0 = fma(bf16(rows[j]), weight, a0) + a1 = fma(bf16(rows[j], True), weight, a1) + word = pack_pair(a0, a1, policy="latent") + packed = tuple(shuffle(word, (4 * lane + k) & 31) for k in range(4)) + if lane < 8: + row = (par * self.tp + rank) * capacity + token + offset = row * 3584 + 64 * line + 8 * lane + multicast_words(latent_mc + Int64(offset) * 2, packed) + elif lane < 8: + u = task - self.batch * 56 + token, line = u // 112, u % 112 + destination = line // (self.shard // 64) + own_line = line % (self.shard // 64) + src = s.iterator + par * ss + token * 7168 + 64 * line + 8 * lane + word = load_volatile(src, count=4) + if cutlass.const_expr(self.mode == "front_poll"): + while reduce(or_, map(front_dirty, word)): + if cutlass.const_expr(self.front_poll_ns): + sleep_ns(Uint32(self.front_poll_ns)) + word = load_volatile(src, count=4) + packed = tuple(no_neg_zero(v) for v in word) + row = (par * self.tp + rank) * capacity + token + offset = row * self.shard + 64 * own_line + 8 * lane + store_relaxed(peers[destination] + Int64(offset) * 2, packed) + + @cute.jit + def gemv(self, sw, sx, partial, bars, q, warp, lane, mt: cutlass.Constexpr[int]): + weight, xs, out = sw.iterator, sx.iterator, partial.iterator + barrier = bars.iterator + 2 + acc = zeros(8 * mt, Float32) + if cutlass.const_expr(self.tp == 8): + # Two bank-spread loads; columns 16*lane .. 16*lane+15 accumulate in order. + swap = (lane >> 2) & 1 + xl = cute.make_rmem_tensor((mt, 8), Uint32) + for token in cutlass.range_constexpr(mt): + aa, bb = spread(xs + token * self.cols, lane, swap) + for k in cutlass.range_constexpr(4): + xl[token, k] = bb[k] if swap != 0 else aa[k] + xl[token, k + 4] = aa[k] if swap != 0 else bb[k] + for row in cutlass.range_constexpr(8): + wa, wb = spread(weight + (warp * 8 + row) * self.cols, lane, swap) + for token in cutlass.range_constexpr(mt): + value = acc[row * mt + token] + for k in cutlass.range_constexpr(8): + ww = (wb + wa)[k] if swap != 0 else (wa + wb)[k] + value = fma_bf16(ww, xl[token, k], value) + value = fma_bf16(ww, xl[token, k], value, high=True) + acc[row * mt + token] = value + elif lane < 28: + for half in cutlass.range_constexpr(2): + chunk = lane + 28 * half + x = [smem_words(xs + t * self.cols + 8 * chunk) for t in range(mt)] + for row in cutlass.range_constexpr(8): + ww = smem_words(weight + (warp * 8 + row) * self.cols + 8 * chunk) + for token in cutlass.range_constexpr(mt): + value = acc[row * mt + token] + for k in cutlass.range_constexpr(4): + value = fma_bf16(ww[k], x[token][k], value) + value = fma_bf16(ww[k], x[token][k], value, high=True) + acc[row * mt + token] = value + # One final row/token per lane group, immediately pushed to the cluster leader. + value, selected = scatter_reduce(acc, lane, 8 * mt) + if (lane & ((4 if mt == 1 else 2) - 1)) == 0: + row, token = selected // mt, selected % mt + push_float(out + (q * 32 + warp * 8 + row) * mt + token, value, barrier, 0) + + @cute.kernel + def kernel(self, args): + latent, shared, weight_map, gamma, gemm2, weights, shared_output = args[0:7] + peers, output, latent_mc, output_mc, par, rank, capacity, eps = args[7:15] + tid, _, _ = cute.arch.thread_idx() + block, _, _ = cute.arch.block_idx() + q = cute.arch.block_idx_in_cluster() + warp, lane = tid // 32, tid % 32 + cluster = block // self.cluster + group, row_cluster = cluster // self.nc, cluster % self.nc + token0, row0 = group * self.tokens, row_cluster * 32 + local_tokens = min(self.tokens, self.batch - token0) + allocator = cutlass.utils.SmemAllocator() + sw = alloc(allocator, BFloat16, 32 * self.cols, 128) + sg = alloc(allocator, BFloat16, self.cols) + sx = alloc(allocator, BFloat16, self.tokens * self.cols) + partial = alloc(allocator, Float32, self.cluster * 32 * self.tokens) + squares = alloc(allocator, Float32, self.cluster * 2 * self.tokens) + shared_sum = alloc(allocator, Float32, self.tokens * 32) + bars = alloc(allocator, Uint64, 4, 8) + bar, expect = bars.iterator, cute.arch.mbarrier_arrive_and_expect_tx + if tid == 127: + init_barriers(bar, 4) + expect(bar + 3, self.cols * 2) + gamma_q = gamma.iterator + q * self.cols + copy_g2s(sg.iterator, gamma_q, size=self.cols * 2, barrier=bar + 3) + expect(bar + 1, self.cluster * 2 * local_tokens * 4) + expect(bar + 2, self.cluster * 32 * local_tokens * 4 if q == 0 else 0) + expect(bar, 32 * self.cols * 2) + acquire_tensormap(weight_map.iterator) + load_tma3(sw.iterator, weight_map.iterator, bar, 2 * q, row0) + cute.arch.cluster_arrive_relaxed() + io = (latent, shared, gemm2, weights, shared_output, peers) + self.front(io, latent_mc, par, rank, capacity, block, warp, lane) + pt, pk = tid // 64, tid % 64 + lat_thread = (pt < local_tokens) & (pk < self.frags) + rs_j = tid - 64 if self.tokens == 1 else pk if self.tp == 8 else pk - self.frags + rs_thread = (q == 0) & (rs_j >= 0) & (rs_j < 4) + if cutlass.const_expr(self.tokens != 1): + rs_thread &= pt < local_tokens + token = Int32(0) if cutlass.const_expr(self.tokens == 1) else pt + s2 = Float32(0) + if lat_thread | rs_thread: + slot = par * self.tp * capacity + token0 + token + lat_ptr = latent.iterator + (slot * 3584 + q * self.cols + 8 * pk) + rs_ptr = shared.iterator + (slot * self.shard + row0 + 8 * max(rs_j, 0)) + base = lat_ptr if lat_thread else rs_ptr + stride = capacity * (3584 if lat_thread else self.shard) + dual = rs_thread & lat_thread + values = self.poll(base, stride, rs_ptr, capacity * self.shard, dual) + sums = self.sum_sources(values, 0) + if lat_thread: + packed = [] + for k in cutlass.range_constexpr(4): + packed.append(pack_pair(sums[2 * k], sums[2 * k + 1])) + square = round_pair_square(packed[k]) + s2 += bf16(square) + s2 += bf16(square, True) + store_shared_words(sx.iterator + token * self.cols + pk * 8, packed) + else: + for k in cutlass.range_constexpr(8): + shared_sum[token * 32 + 8 * rs_j + k] = Float32(BFloat16(sums[k])) + if cutlass.const_expr(self.tp == 8 and self.tokens == 2): # noqa: SIM102 + if dual: + rs = self.sum_sources(values, 8) + for k in cutlass.range_constexpr(8): + shared_sum[token * 32 + 8 * rs_j + k] = Float32(BFloat16(rs[k])) + cute.arch.sync_warp() + for stage in cutlass.range_constexpr(5): + s2 += cute.arch.shuffle_sync_bfly(s2, 16 >> stage) + cute.arch.cluster_wait() + if pt < local_tokens and lane < self.cluster: + sq = squares.iterator + (2 * q + (warp & 1)) * local_tokens + pt + push_float(sq, s2, bar + 1, lane) + cute.arch.sync_warp() + wait(bar + 1) + cute.arch.sync_warp() + if cutlass.const_expr(self.trigger == 3): + cute.arch.griddepcontrol_wait() + cute.arch.griddepcontrol_launch_dependents() + wait(bar + 3) + if lat_thread: + total = Float32(0) + for peer in cutlass.range_constexpr(2 * self.cluster): + total += squares[peer * local_tokens + pt] + inverse = inv_rms(total, eps) + gx = smem_words(sg.iterator + pk * 8) + xx = smem_words(sx.iterator + pt * self.cols + pk * 8) + scaled = lambda k, h: bf16(xx[k], h) * inverse * bf16(gx[k], h) + packed = [pack_pair(scaled(k, 0), scaled(k, 1)) for k in range(4)] + store_shared_words(sx.iterator + pt * self.cols + pk * 8, packed) + wait(bar) + cute.arch.sync_warp() + cute.arch.sync_threads() + if self.tokens == 1 or self.batch % 2 == 0 or local_tokens == 2: + self.gemv(sw, sx, partial, bars, q, warp, lane, self.tokens) + else: + self.gemv(sw, sx, partial, bars, q, warp, lane, 1) + if q == 0: # noqa: SIM102 - Keep the owner CTA branch separate from its warp selection. + if warp < local_tokens: + wait(bar + 2) + cute.arch.sync_warp() + value = Float32(0) + for peer in cutlass.range_constexpr(self.cluster): + value += partial[(peer * 32 + lane) * local_tokens + warp] + rounded = Float32(BFloat16(value)) + out = Float32(BFloat16(rounded + shared_sum[warp * 32 + lane])) + # Pack paired rows, then four pairs into a 16-byte writer. + pair = pack_pair(out, shuffle(out, (lane + 1) & 31), policy="latent") + words = tuple(shuffle(pair, (8 * lane + 2 * k) & 31) for k in range(4)) + if lane < 4: + offset = (token0 + warp) * 7168 + rank * self.shard + row0 + multicast_words(output_mc + Int64(offset + 8 * lane) * 2, words) + if cutlass.const_expr(self.trigger == 2): + # front_poll: the next AttnRes reads the prefix before any wait, and no + # kernel since the o-proj consumer that wrote it (the MoE front and + # experts) waits before releasing: this grid must, at every batch. + if cutlass.const_expr(self.mode == "front_poll"): + cute.arch.griddepcontrol_wait() + cute.arch.griddepcontrol_launch_dependents() + self.rearm(io, capacity, *args[15:17], block, tid) + # Orders deferred rearm, workspace reuse and producer lifetimes (all modes). + cute.arch.griddepcontrol_wait() + + @cute.jit + def rearm(self, io, capacity, parity, batch, block, tid, reduced=None): + """Re-arms ``parity``'s first ``batch`` tokens (and WaveMoeTail's reduced + rows).""" + latent, shared, g, w, s, _ = io + if parity >= 0: + nl, nr = self.tp * batch * 448, self.tp * batch * (self.shard // 8) + ng, nw, ns = Int32(0), Int32(0), Int32(0) + if cutlass.const_expr(self.mode == "front_poll"): + gs, ws, ss = self.poll_strides(io) + ng = min(batch, gs // (16 * 3584)) * 16 * 448 + nw = min(batch, ws // 16) * 2 + ns = min(batch, ss // 7168) * 896 + elif cutlass.const_expr(reduced is not None): + ng = batch * 448 + ring = (parity, capacity, batch) + start, step = block * self.threads + tid, self.grid * self.threads + for i in cutlass.range(start, nl + nr + ng + nw + ns, step, unroll=1): + if i < nl: + self.rearm_row(latent, 3584, i, ring) + elif i < nl + nr: + self.rearm_row(shared, self.shard, i - nl, ring) + elif cutlass.const_expr(self.mode == "front_poll"): + j = i - nl - nr + fill = (Uint32(0xFFFFFFFF),) * 4 + if j < ng: + store_words(g.iterator + parity * gs + 8 * j, fill) + elif j < ng + nw: + store_words(w.iterator + parity * ws + 8 * (j - ng), fill) + else: + store_words(s.iterator + parity * ss + 8 * (j - ng - nw), fill) + elif cutlass.const_expr(reduced is not None): + rows = reduced.iterator + parity * capacity * 3584 + store_words(rows + 8 * (i - nl - nr), (Uint32(0x80000000),) * 4) + + @cute.jit + def rearm_row(self, box, width: cutlass.Constexpr[int], i, ring): + parity, capacity, batch = ring + row, frag = i // (width // 8), i % (width // 8) + source, token = row // batch, row % batch + dst = box.iterator + ((parity * self.tp + source) * capacity + token) * width + store_words(dst + 8 * frag, (Uint32(0x80000000),) * 4) + + +def pow2_spans(n): + """(start, count) spans of n columns, power-of-two counts, largest first.""" + bits = [1 << i for i in reversed(range(n.bit_length())) if n >> i & 1] + return tuple((sum(bits[:i]), b) for i, b in enumerate(bits)) + + +class WaveMoeTail(MoeTail): + """One-wave tail over published inputs (TP8 B5+, TP16 B9+): seven 16-CTA clusters + split K into 224-column slices of a 128-row (TP16: 64-row) tile. Weight fragments + stay in registers through the latent poll, RMS all-gather and normalization; row + warps (TP16: two sets taking alternate groups) run two seven-step BF16 MMA chains + per 8-token group, and each 64-row output line has one owner CTA and one multicast + writer. Past one group (and at TP16) each cluster instead all-reduces every seventh + token of its slices once and shares them as BF16 rows through ``reduced``, so all + 112 CTAs must be co-resident; RS lines spread over warp pairs, every group's RMS + sums travel as one vector per peer, all groups normalize before any MMA, the + groups' MMA chains interleave per K step, and the previous layer's rows are + re-armed after the output multicast. ``twoshot`` (B9+): the latent + mailbox is this rank's ``common.twoshot`` receive slots [3, TP, capacity, 3584/TP] + (pushed by the experts), the CTAs of its owned slices reduce them like TP16 but + multicast the rows into every rank's ``reduced`` (the gather; its multicast + address replaces the latent one), so each GPU receives 2 x B x 7 KB of latent + instead of TP x B x 7 KB.""" + + threads, options = 256, dict(min_blocks_per_mp=1) + + def __init__(self, batch, tp=8, twoshot=False, early=False, umma=False): + super().__init__(tp, batch, mode="publish") + assert not twoshot or batch > 8 + assert not umma or batch > 8 + # ``early``: release the next AttnRes at entry, so its block statistics and + # exchange overlap this kernel (it reads only earlier layers' data before its + # Lamport poll of the output rows), instead of after the output multicast. + self.early = early + self.lines, self.groups = 2 if tp == 8 else 1, (batch + 7) // 8 + self.slots = (batch * self.lines + 15) // 16 # owned output lines per CTA + self.twoshot, self.width = twoshot, 3584 // tp if twoshot else 3584 + # Past one group (and at TP16) clusters share reduced latent rows; TP8 B<=8 + # polls its sources directly. + self.reduce = self.groups > 1 or tp == 16 + self.grid, self.poll_ns, self.cluster = 112, 64, 16 + # Shared memory: per group x [8, 224] BF16 and MMA partials [8 warps, 8, 16] + # (TP16 pads to whole warp sets); gathered squares (one group keeps each + # slice's 28 lane sums, more one sum); per owned line its partials [16 slices, + # 64] and RS rows [64]; two barriers. + n, s = self.groups + self.groups % (2 // self.lines), self.slots + # Past one group each (warp, slice) row holds the groups' sums, padded to v4. + self.padded = -(-self.groups // 4) * 4 + self.squares = 3584 if n == 1 else 128 * max(n, self.padded) + partials = 7680 * n + 4 * self.squares + self.offsets: tuple = (3584 * n, 7680 * n, partials, partials + 4096 * s) + self.offsets += (partials + 4352 * s + 32,) + self.smem = 116 * 1024 # one CTA per SM + # ``umma`` (past one group): the CTA's [M, 224] x [N, 224] up projection on + # tcgen05 (M = 64 x lines rows, N = the padded tokens, K padded to 256) from + # 128B-swizzled W and x tiles (W: four TMA boxes [M, 64] at entry, the last + # running into the next slice, whose columns meet the zeroed x padding; the + # weight argument is then the shard's tensor map), D in TMEM, staged as + # [token, row] in the drained W tiles for the owners' pushes. FP32 + # accumulation in the tensor core's order. + self.umma = umma + if umma: + self.m = 64 * self.lines + self.n = -(-8 * self.groups // (16 if self.lines == 2 else 8)) + self.n *= 16 if self.lines == 2 else 8 + self.columns = max(32, 1 << (self.n - 1).bit_length()) + w_bytes, x_bytes = self.m * 512, self.n * 512 + squares = w_bytes + x_bytes + partials = squares + 4 * self.squares + self.offsets = (0, squares, partials, partials + 4096 * s) + self.offsets += (partials + 4352 * s + 32,) + self.smem = max(self.smem, self.offsets[4] + 40) + + def group_smem(self, sm, g): + """Group g's x (BF16 pointer) and MMA partials.""" + x = view(sm, 3584 * g, BFloat16, 1792).iterator + return x, view(sm, self.offsets[0] + 4096 * g, Float32, 1024) + + def line_smem(self, sm, s): + """Owned line s's partials [16 slices, 64] and RS rows.""" + partial, rs = self.offsets[2:4] + cells = view(sm, partial + 4096 * s, Float32, 1024) + return cells, view(sm, rs + 256 * s, Float32, 64) + + @cute.kernel + def kernel(self, args): + latent, shared, weight, gamma = args[0:4] + output_mc, par, rank, capacity, eps = args[10:15] + tid, warp, lane, block = thread() + if cutlass.const_expr(self.early): + cute.arch.griddepcontrol_launch_dependents() + q = cute.arch.block_idx_in_cluster() + n, batch, lines, rows = self.groups, self.batch, self.lines, 64 * self.lines + row0, gq, tq, owner = block // 16 * rows, lane // 4, lane % 4, q < lines * batch + own_token, own_line = q // lines, q % lines + rs_j, lat = 4 * warp + lane - 28, lane < 28 + owns = lambda s: owner if s == 0 else q < lines * batch - 16 * s + rs_lane = lambda g: (lane >= 28) & owns(g) & (rs_j < 8) + align = 1024 if self.umma else 128 + sm = alloc(cutlass.utils.SmemAllocator(), Uint8, self.smem, align).iterator + groups = [self.group_smem(sm, g) for g in range(n)] + owned = [self.line_smem(sm, s) for s in range(self.slots)] + squares = view(sm, self.offsets[1], Float32, self.squares) + bar = view(sm, self.offsets[4], Uint64, 4).iterator + tmem_slot = view(sm, self.offsets[4] + 32, Uint32, 1) + if cutlass.const_expr(self.umma): # noqa: SIM102 + if warp == 0: + cute.arch.alloc_tmem(self.columns, tmem_slot.iterator) + cute.arch.relinquish_tmem_alloc_permit() + if tid == 0: + init_barriers(bar, 4 if self.umma else 2) + if cutlass.const_expr(self.umma): # W [M, 64] x 4 by TMA on barrier 3 + acquire_tensormap(weight.iterator) + cute.arch.mbarrier_arrive_and_expect_tx(bar + 3, self.m * 512) + for kb in cutlass.range_constexpr(4): + tile = sm + kb * (self.m * 128) + load_tma2(tile, weight.iterator, bar + 3, 224 * q + 64 * kb, row0) + # One group gathers 28 lane sums per (token, slice), more one sum for + # every warp's token. + gathered = 16 * batch * 28 * 4 if n == 1 else 16 * 8 * self.padded * 4 + cute.arch.mbarrier_arrive_and_expect_tx(bar, gathered) + size = 4096 if owns(self.slots - 1) else 0 # earlier lines are all owned + size = size + 4096 * (self.slots - 1) if self.slots > 1 else size + cute.arch.mbarrier_arrive_and_expect_tx(bar + 1, size) + cute.arch.cluster_arrive_relaxed() + xs = groups[0][0] + if cutlass.const_expr(not self.umma): + for i in cutlass.range(tid, (8 * n - batch) * 28, 256, unroll=1): + store_shared_words(xs + batch * 224 + 8 * i, (Uint32(0),) * 4) + gx, gamma_q = cute.make_rmem_tensor(4, Uint32), gamma.iterator + q * 224 + if lat & (warp < batch): + for k, v in enumerate(load_words(gamma_q + 8 * lane, count=4)): + gx[k] = v + # Rows 16 warp + gq (+ 8) of the cluster's tile (TP16: warp mod 4), columns + # 224 q + 32 k + 8 tq. + rw = warp if lines == 2 else warp % 4 + base = weight.iterator + (row0 + 16 * rw + gq) * 3584 + 224 * q + 8 * tq + fetch = lambda p: load_words(p, count=4, evict_first=True, mutable=True) + a = None + if cutlass.const_expr(not self.umma): + a = [ + [fetch(base + h * 8 * 3584 + 32 * k) for h in range(2)] + for k in range(7) + ] + xwords = [cute.make_rmem_tensor(4, Uint32) for _ in range(n)] + reduced = args[17] if self.reduce else None + ring = (*args[15:17], block, tid, reduced) + if cutlass.const_expr(self.reduce): + s2 = zeros(n, Float32) + ids: tuple = (q, warp, lane, block // 16) + places = (row0, owned[0][1].iterator) + gather = (args[9], rank) if self.twoshot else None + boxes = (latent, reduced, shared) + self.poll_inputs(boxes, (par, capacity), ids, places, gather) + if lat: + self.read_reduced(reduced, par * capacity, xwords, s2, ids) + else: + # Lanes 0..27 of warp t poll token t's latent; owners' lanes 28..31 the RS. + sq = Float32(0) + if lat & (warp < batch) | rs_lane(0): + slot = par * self.tp * capacity + rs = (slot + own_token) * self.shard + row0 + 64 * own_line + rs += 8 * max(rs_j, 0) + ptr = latent.iterator + (slot + warp) * 3584 + 224 * q + 8 * lane + ptr = ptr if lat else shared.iterator + rs + stride = capacity * (3584 if lat else self.shard) + sums = self.sum_sources(self.poll(ptr, stride, ptr, 0, False), 0) + if lat: + sq = self.latent_words(sums, xwords[0]) + else: + self.rs_words(sums, owned[0][1], rs_j) + s2 = [sq] + cute.arch.sync_warp() + cute.arch.cluster_wait() + if cutlass.const_expr(n == 1): + self.push_squares(squares, s2[0], bar, (q, warp, lane)) + wait(bar) + cute.arch.sync_warp() + row = (xwords[0], gx, eps) + self.normalize(squares, row, groups[0][0], (warp, lane, lat)) + else: + self.exchange_squares(squares, s2, bar, (q, warp, lane)) + wait(bar) + cute.arch.sync_warp() + outs = [groups[g][0] for g in range(n)] + if cutlass.const_expr(self.umma): # x tiles [N, 64] x 4 after W + outs = [sm + self.m * 512 + 1024 * g for g in range(n)] + self.normalize_groups(squares, (xwords, gx, eps), outs, (warp, lane, lat)) + if cutlass.const_expr(self.umma): # x stores -> tcgen05 + cute.arch.fence_proxy("async.shared", space="cta") + fence_tmem(before=True) + cute.arch.sync_warp() + cute.arch.sync_threads() + ids = (q, warp, lane, rw, gq, tq) + if cutlass.const_expr(n == 1): + self.mma_group(a, groups[0], owned, bar, 0, ids) + elif cutlass.const_expr(self.umma): + tmem = tmem_slot[0] + self.umma_partials(sm, tmem, bar, (q, warp, lane, tid)) + else: + size = (self.offsets[1] - self.offsets[0]) // 4 + products = view(sm, self.offsets[0], Float32, size) + self.mma_groups(a, xs, products, owned, bar, ids) + line = (lane, rank, row0, own_token, own_line) + for s in cutlass.range_constexpr(self.slots): + if owns(s) & (warp == s): + self.write_line(owned[s], bar, output_mc, line, s) + if cutlass.const_expr(not self.early): + cute.arch.griddepcontrol_launch_dependents() + # The previous layer's rows, off the output path. + if cutlass.const_expr(n == 1): + self.rearm(args[0:2] + args[4:8], capacity, *ring) + else: + self.rearm_rows(args[0:2], capacity, *ring) + if cutlass.const_expr(self.umma): + release_tmem(tmem, self.columns, warp, 0) + cute.arch.griddepcontrol_wait() + + @cute.jit + def load_sources(self, values, ptr, stride): + """Issue the TP source loads of a lane's 16-byte column without waiting.""" + for source in cutlass.range_constexpr(self.tp): + word = load_volatile(ptr + source * stride, count=4) + for k in cutlass.range_constexpr(4): + values[source, k] = word[k] + + @cute.jit + def settle(self, values, ptr, stride): + """Re-poll loaded sources until no word is the 0x80000000 sentinel.""" + pending = Uint32(0) + for source in cutlass.range_constexpr(self.tp): + words = (values[source, k] for k in range(4)) + pending |= Uint32(reduce(or_, (v == 0x80000000 for v in words))) << source + while pending != 0: + sleep_ns(Uint32(self.poll_ns)) + again = Uint32(0) + for source in cutlass.range_constexpr(self.tp): + if (pending & (Uint32(1) << source)) != 0: + again |= self.poll_row(values, source, ptr + source * stride) + pending = again + + @cute.jit + def column_sums(self, values): + """``sum_sources`` (same order and rounding) as one BF16 FMA by 1.0 per term.""" + acc, one = zeros(8, Float32), Uint32(0x3F803F80) + for source in cutlass.range_constexpr(self.tp): + for k in cutlass.range_constexpr(8): + acc[k] = fma_bf16(values[source, k // 2], one, acc[k], high=k % 2 == 1) + return acc + + @cute.jit + def rearm_rows(self, boxes, capacity, parity, batch, block, tid, reduced): + """Division-free ``rearm``: each source's latent and RS rows, then the reduced + rows, of the previous layer (complete before this launch).""" + if parity >= 0: + start, step = block * self.threads + tid, self.grid * self.threads + fill, widths = (Uint32(0x80000000),) * 4, (self.width, self.shard) + for s in cutlass.range_constexpr(self.tp): + first = (parity * self.tp + s) * capacity + for b in cutlass.range_constexpr(2): + row, w = boxes[b].iterator + first * widths[b], widths[b] + for i in cutlass.range(start, batch * (w // 8), step, unroll=1): + store_words(row + 8 * i, fill) + if cutlass.const_expr(reduced is not None): + row = reduced.iterator + parity * capacity * 3584 + for i in cutlass.range(start, batch * 448, step, unroll=1): + store_words(row + 8 * i, fill) + + @cute.jit + def poll_inputs(self, boxes, ring, ids, places, gather=None): + """Phase 1 and the RS lines in one convergent pass (every round's loads issued + first). Lanes 0..27 all-reduce the cluster's tokens tile + 7 j of the CTA's + slice once into BF16 ``reduced`` rows (``gather`` (multicast address, rank): + only the CTAs of owned slices, from the receive slots into every rank's); + lanes 28..31 of warps 2 p and 2 p + 1 sum owned lines p, p + 4's RS shard + (line 16 i + q: token (16 i + q) / lines).""" + (latent, reduced, shared), (par, capacity) = boxes, ring + (q, warp, lane, tile), (row0, rs_rows) = ids, places + lat, width = lane < 28, self.width + col = 224 * q + 8 * lane + local = col if gather is None else col - gather[1] * width + src = latent.iterator + (par * self.tp * capacity + tile) * width + local + dst = reduced.iterator + (par * capacity + tile) * 3584 + col + reduces = q >= 0 if gather is None else q * self.tp // 16 == gather[1] + rs_j = 4 * (warp % 2) + lane - 28 + rs = shared.iterator + par * self.tp * capacity * self.shard + row0 + 8 * rs_j + stride = capacity * width if lat else capacity * self.shard + latent_rounds = ((self.batch - 1) // 7 + 8) // 8 + rounds = max(latent_rounds, (self.slots + 3) // 4) + tokens = [7 * (warp + 8 * r) for r in range(rounds)] + lines = [4 * r + warp // 2 for r in range(rounds)] + live, ptrs = [], [] + for r in cutlass.range_constexpr(rounds): + token, line = tokens[r], lines[r] + on = lane < 0 # no latent round r + if cutlass.const_expr(r < latent_rounds): + on = reduces & (tile + token < self.batch) + rs_on = (line < self.slots) & (16 * line + q < self.lines * self.batch) + live.append(on if lat else rs_on) + rs_row = rs + (16 * line + q) // self.lines * self.shard + rs_row += 64 * ((16 * line + q) % self.lines) + ptrs.append(src + token * width if lat else rs_row) + values = [cute.make_rmem_tensor((self.tp, 4), Uint32) for _ in range(rounds)] + for r in cutlass.range_constexpr(rounds): + if live[r]: + self.load_sources(values[r], ptrs[r], stride) + for r in cutlass.range_constexpr(rounds): + if live[r]: + self.settle(values[r], ptrs[r], stride) + sums = self.column_sums(values[r]) + if lat: + if cutlass.const_expr(gather is None): + self.store_reduced(sums, dst + tokens[r] * 3584) + else: + offset = (par * capacity + tile + tokens[r]) * 3584 + col + self.store_reduced(sums, gather[0] + Int64(offset) * 2, True) + else: + rows = view(rs_rows, 64 * lines[r], Float32, 64) + self.rs_words(sums, rows, rs_j) + + @cute.jit + def store_reduced(self, sums, dst, multicast: cutlass.Constexpr[bool] = False): + pack = lambda j: pack_pair(sums[2 * j], sums[2 * j + 1], policy="latent") + if cutlass.const_expr(multicast): + multicast_words(dst, [pack(j) for j in range(4)]) + else: + mailbox_store(dst, [pack(j) for j in range(4)]) + + @cute.jit + def read_reduced(self, reduced, slot, xwords, s2, ids): + """Phase 2: the lane's words of tokens 8 g + warp and their squares.""" + q, warp, lane, _ = ids + base = reduced.iterator + (slot + warp) * 3584 + 224 * q + 8 * lane + values = self.poll_rows(base, 8 * 3584, (self.batch - warp + 7) // 8) + for g in cutlass.range_constexpr(self.groups): + for k in cutlass.range_constexpr(4): + xwords[g][k] = values[g, k] + s2[g] = self.square_sum(xwords[g]) + + @cute.jit + def poll_rows(self, base, stride, count): + """Lamport-poll the first ``count`` 16-byte rows base + r stride.""" + values = cute.make_rmem_tensor((self.groups, 4), Uint32) + pending = Uint32(0) + for r in cutlass.range_constexpr(self.groups): + if r < count: + pending |= self.poll_row(values, r, base + r * stride) + while pending != 0: + sleep_ns(Uint32(self.poll_ns)) + again = Uint32(0) + for r in cutlass.range_constexpr(self.groups): + if (pending & (Uint32(1) << r)) != 0: + again |= self.poll_row(values, r, base + r * stride) + pending = again + return values + + @cute.jit + def poll_row(self, values, r: cutlass.Constexpr[int], ptr): + word = load_volatile(ptr, count=4) + for k in cutlass.range_constexpr(4): + values[r, k] = word[k] + return Uint32(reduce(or_, (v == 0x80000000 for v in word))) << r + + @cute.jit + def mma_group(self, a, group, owned, bar, g: cutlass.Constexpr[int], ids): + """Group g's MMA chains; each 4-row partial goes to its 64-row line's owner.""" + (q, warp, lane, rw, gq, tq), (xg, pg) = ids, group + b = [smem_words(xg + gq * 224 + 8 * tq + 32 * k) for k in range(7)] + d = e = (Float32(0),) * 4 + for k in cutlass.range_constexpr(7): + (w0, w1), x = a[k], b[k] + d = mma((w0[0], w1[0], w0[1], w1[1]), x[0:2], d, bf16=True) + e = mma((w0[2], w1[2], w0[3], w1[3]), x[2:4], e, bf16=True) + for i in cutlass.range_constexpr(4): # [warp, token, row] partial products + pg[warp * 128 + (2 * tq + i % 2) * 16 + gq + 8 * (i // 2)] = d[i] + e[i] + cute.arch.sync_warp() + token, row = lane // 4, 16 * rw + 4 * (lane % 4) + if token < self.batch - 8 * g: + values = load_f32x4(pg.iterator + warp * 128 + 4 * lane, shared=True) + # Line lines (8 g + token) + row // 64 is owned by CTA (line % 16). + peer = self.lines * token + row // 64 + 8 * g * self.lines % 16 + cells = owned[8 * g * self.lines // 16][0] + push_float(cells.iterator + q * 64 + row % 64, values, bar + 1, peer) + + @cute.jit + def mma_groups(self, a, xs, products, owned, bar, ids): + """``mma_group`` over this warp's groups g = sets j + set (TP16: two warp + sets): every MMA chain (the groups' chains interleaved per K step), one warp + sync, then the owners' pushes (slot j).""" + q, warp, lane, rw, gq, tq = ids + sets = 2 // self.lines + h = 0 if sets == 1 else warp // 4 + mine = (self.groups + sets - 1) // sets + x = [xs + (sets * j + h) * 1792 + gq * 224 + 8 * tq for j in range(mine)] + d = [(Float32(0),) * 4 for _ in range(mine)] + e = [(Float32(0),) * 4 for _ in range(mine)] + for k in cutlass.range_constexpr(7): + w0, w1 = a[k] + for j in cutlass.range_constexpr(mine): + xk = smem_words(x[j] + 32 * k) + d[j] = mma((w0[0], w1[0], w0[1], w1[1]), xk[0:2], d[j], bf16=True) + e[j] = mma((w0[2], w1[2], w0[3], w1[3]), xk[2:4], e[j], bf16=True) + acc = [[d[j][i] + e[j][i] for i in range(4)] for j in range(mine)] + for j in cutlass.range_constexpr(mine): + cell = (sets * j + h) * 1024 + warp * 128 + 2 * tq * 16 + gq + for i in cutlass.range_constexpr(4): + products[cell + i % 2 * 16 + 8 * (i // 2)] = acc[j][i] + cute.arch.sync_warp() + token, row = lane // 4, 16 * rw + 4 * (lane % 4) + peer = self.lines * token + row // 64 + (8 * h if sets > 1 else 0) + for j in cutlass.range_constexpr(mine): + if token < self.batch - 8 * (sets * j + h): + src = products.iterator + (sets * j + h) * 1024 + warp * 128 + 4 * lane + values = load_f32x4(src, shared=True) + cells = owned[j][0] + push_float(cells.iterator + q * 64 + row % 64, values, bar + 1, peer) + + @cute.jit + def write_line(self, line, bar, output_mc, ids, s: cutlass.Constexpr[int]): + """Owned line s: sum the 16 slices' partials, add the RS rows and multicast the + 64 output rows.""" + (cells, rs_sum), (lane, rank, row0, own_token, own_line) = line, ids + wait(bar + 1) + cute.arch.sync_warp() + out = [] + for h in cutlass.range_constexpr(2): + parts = (cells[64 * p + 2 * lane + h] for p in range(16)) + value = Float32(BFloat16(sum(parts, Float32(0)))) + out.append(value + rs_sum[2 * lane + h]) + pair = pack_pair(*out, policy="latent") + words = tuple(shuffle(pair, (4 * lane + k) & 31) for k in range(4)) + if lane < 8: + t = 16 // self.lines * s + own_token + offset = t * 7168 + rank * self.shard + row0 + 64 * own_line + multicast_words(output_mc + Int64(offset + 8 * lane) * 2, words) + + @cute.jit + def latent_words(self, sums, words): + """BF16 latent words of a lane's eight columns; returns their square sum.""" + s2 = Float32(0) + for j in cutlass.range_constexpr(4): + words[j] = pack_pair(sums[2 * j], sums[2 * j + 1]) + square = round_pair_square(words[j]) + s2 += bf16(square) + s2 += bf16(square, True) + return s2 + + @cute.jit + def square_sum(self, words): + s2 = Float32(0) + for j in cutlass.range_constexpr(4): + square = round_pair_square(words[j]) + s2 += bf16(square) + s2 += bf16(square, True) + return s2 + + @cute.jit + def rs_words(self, sums, rs_sum, rs_j): + for j in cutlass.range_constexpr(8): + rs_sum[8 * rs_j + j] = Float32(BFloat16(sums[j])) + + @cute.jit + def push_squares(self, squares, s2, bar, ids): + """One group: the warp's token's 28 lane sums of every slice to every peer.""" + q, warp, lane = ids + if warp < self.batch: + chunk, first = lane % 7, lane // 7 + vector = tuple(shuffle(s2, 4 * chunk + i) for i in range(4)) + if first < 4: + dst = squares.iterator + warp * 448 + 28 * q + 4 * chunk + for peer in cutlass.range(first, 16, 4, unroll=1): + push_float(dst, vector, bar, peer) + + @cute.jit + def exchange_squares(self, squares, s2, bar, ids): + """Past one group: every group's warp sum for the warp's token; lane p < 16 + pushes them to peer p's (warp, slice) row as v4 words.""" + q, warp, lane = ids + values = [s2[g] for g in range(self.groups)] + values += [Float32(0)] * (self.padded - self.groups) + total = scatter_sums(values, lane, 32) + totals = gather_sums(total, self.padded, 32, self.padded) + if lane < 16: + dst = squares.iterator + (warp * 16 + q) * self.padded + for c in cutlass.range_constexpr(self.padded // 4): + push_float(dst + 4 * c, tuple(totals[4 * c : 4 * c + 4]), bar, lane) + + @cute.jit + def normalize_groups(self, squares, row, xs, ids): + """Past one group: every group's RMS (the 16 slices' sums, ``square_total``'s + order) first, then the scaled latent words (padding tokens too).""" + (words, gx, eps), (warp, lane, lat) = row, ids + base = squares.iterator + (warp * 16 + lane % 16) * self.padded + sums = [] + for c in cutlass.range_constexpr(self.padded // 4): + sums += list(load_f32x4(base + 4 * c, shared=True)) + inverse = inv_rms(scatter_sums(sums, lane, 16), eps, m8=True) + inverse = gather_sums(inverse, self.groups, 16, self.padded) + if lat: + for g in cutlass.range_constexpr(self.groups): + self.scale_words(words[g], gx, inverse[g], xs[g], warp, lane) + else: + if cutlass.const_expr(self.umma): # K columns 224..255 of the x tiles + for g in cutlass.range_constexpr(self.groups): + dst = self.x_chunk(xs[g], warp, lane) + store_shared_words(dst, (Uint32(0),) * 4) + + @cute.jit + def normalize(self, squares, row, xs, ids): + """One group: RMS-normalize the warp's token into ``xs``.""" + (words, gx, eps), (warp, lane, _) = row, ids + if warp < self.batch: + total = self.square_total(squares, warp, lane) + self.scale(words, gx, inv_rms(total, eps, m8=True), xs, ids) + + @cute.jit + def square_total(self, squares, warp, lane): + groups = cute.make_rmem_tensor(16, Float32) + for c in cutlass.range_constexpr(16): + groups[c] = Float32(0) + if lane < (16 if c % 4 == 3 else 32): + groups[c] = squares[warp * 448 + c // 4 * 112 + c % 4 * 32 + lane] + group_sum, _ = scatter_reduce(groups, lane, 16) + return sum((shuffle(group_sum, 2 * c) for c in range(16)), Float32(0)) + + @cute.jit + def scale(self, words, gx, inverse, xs, ids): + warp, lane, lat = ids + if lat: + self.scale_words(words, gx, inverse, xs, warp, lane) + + def scale_words(self, words, gx, inverse, xs, warp, lane): + """(x * inverse) * gamma per column, as packed FP32x2 multiplies.""" + pair = lambda w: (bf16(w), bf16(w, True)) + scaled = lambda k: mul2(mul2(pair(words[k]), (inverse, inverse)), pair(gx[k])) + packed = [pack_pair(*scaled(k)) for k in range(4)] + dst = xs + warp * 224 + 8 * lane + if cutlass.const_expr(self.umma): + dst = self.x_chunk(xs, warp, lane) + store_shared_words(dst, packed) + + def x_chunk(self, tile, row, chunk): + """16-byte chunk ``chunk`` (K columns 8 chunk ..) of row ``row`` in 128B- + swizzled [rows, 64] K tiles of ``self.n`` rows (``tile``: the 8-row block).""" + atom, part = chunk // 8, chunk % 8 + return tile + atom * (self.n * 128) + row * 128 + (part ^ row) * 16 + + @cute.jit + def umma_partials(self, sm, tmem, bar, ids): + """D [M, N] = W x^T on tcgen05 (warp 0 issues 4 K tiles), staged [token, row] + in the W tiles, then the 4-row partials pushed to their lines' owners.""" + q, warp, lane, tid = ids + m, n = self.m, self.n + done = bar + 2 + instruction: cutlass.Constexpr = 1 << 4 | 1 << 7 | 1 << 10 + instruction |= n >> 3 << 17 | m >> 4 << 24 + if warp == 0: + wait(bar + 3) # W tiles + fence_tmem() + for kb in cutlass.range_constexpr(4): + a = smem_desc(sm + kb * (m * 128), True) + b = smem_desc(sm + m * 512 + kb * (n * 128), True) + mma_bf16_block(tmem, a, b, instruction, int(kb > 0), done, int(kb == 3)) + wait_tmem(done) + # The MMAs (async proxy) read the W tiles the staging now overwrites. + cute.arch.fence_proxy("async.shared", space="cta") + # Warp w reads TMEM lanes 32 (w mod 4) .. (M = 64: the first 16 of each), + # half w / 4 of the columns. + quad, half, cols = warp % 4, warp // 4, n // 2 + values = () + for start, count in pow2_spans(cols): # TMEM loads: power-of-two columns + got = load_tmem(tmem + (32 * quad << 16) + half * cols + start, count) + values += got if count > 1 else (got,) + row = 32 * quad + lane if m == 128 else 16 * quad + lane + stage = view(sm, 0, Float32, n * m) + if lane < (32 if m == 128 else 16): + for t in cutlass.range_constexpr(cols): + stage[(half * cols + t) * m + row] = values[t] + fence_tmem(before=True) + cute.arch.sync_threads() + partials = view(sm, self.offsets[2], Float32, 1024 * self.slots).iterator + for p in cutlass.range(tid, self.batch * (m // 4), 256, unroll=1): + token, r = p // (m // 4), p % (m // 4) * 4 + line = self.lines * token + r // 64 + values4 = load_f32x4(stage.iterator + token * m + r, shared=True) + dst = partials + line // 16 * 1024 + q * 64 + r % 64 + push_float(dst, values4, bar + 1, line % 16) + + @cute.jit + def push(self, *args): + """WaveMoeTail's shared inputs: peer-store this rank's shared-expert rows [B, + 7168] into the RS mailboxes (-0 made +0); an ordinary launch after the down + GEMM.""" + launch(self.push_kernel, args, self.batch, 128, pdl=False) + + @cute.kernel + def push_kernel(self, args): + shared, peers, parity, rank, capacity = args + tid, _, _, token = thread() + for i in cutlass.range_constexpr(7): + column = 8 * (tid + 128 * i) + words = load_words(shared.iterator + token * 7168 + column, count=4) + offset = ((parity * self.tp + rank) * capacity + token) * self.shard + offset += column % self.shard + packed = tuple(no_neg_zero(w) for w in words) + store_relaxed(peers[column // self.shard] + Int64(offset) * 2, packed) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/tail_runtime.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/tail_runtime.py new file mode 100644 index 0000000000..a1aa106e27 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/tail_runtime.py @@ -0,0 +1,241 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Routed MoE tail: one fused kernel finalizes this rank's top-16 rows, multicasts the +latent partial and peer-stores the shared shard (three rotating parities), all-reduces +and RMS-normalizes per K-slice, runs the K-split up projection and multicasts the output +row into the next AttnRes mailbox. At TP8 B5+ and TP16 B9+ the experts publish the +latent partial, WaveMoeTail.push the shared shard and WaveMoeTail does the rest. At +``twoshot_batches`` the latent all-reduce is ``common.twoshot``: the experts push each +partial chunk to its column owner, whose WaveMoeTail CTAs sum and multicast it.""" + +from functools import cache + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass import BFloat16, Float32, Int32, Int64, Uint8 +from cutlass.cute.runtime import make_fake_stream + +from ..common.host import fake, prepare_shared_carveout, stream, tensor_maps, to_cute +from ..common.policy import batch_table +from ..common.twoshot import TwoShot +from ..oproj.prepared import OprojMailbox +from .tail import MoeTail, WaveMoeTail + +# MoE tail mailbox tokens (padded decode batches); small-batch front_poll G/W/S rows per +# parity. +CAP, POLL_CAP = 64, 8 +# Two-shot latent all-reduce batches per TP (B9+): 2-/4-node GB300, us/layer of RS push +# + experts + WaveMoeTail, max over ranks, multicast vs two-shot: TP8 B16 41.74/42.80, +# B32 52.20/53.18, B64 79.27/81.50; TP16 B16 33.98/34.49, B32 42.19/42.52, B64 +# 65.05/60.51. +TWOSHOT = {8: "", 16: "33-64:1"} +# WaveMoeTail batches that release the next AttnRes at entry. +EARLY = {8: "9-64:1", 16: "9-64:1"} +# WaveMoeTail batches whose up projection runs on tcgen05 (B9+; 1-GPU WaveMoeTail us +# mma.sync vs tcgen05, B16/32/64: TP8 4.53/5.80/9.05 vs 4.45/5.26/7.80, TP16 two-shot +# 4.50/5.39/7.75 vs 4.52/5.46/7.21). +UMMA = {8: "9-64:1", 16: "33-64:1"} + + +def flat(t): + return to_cute(t.reshape(-1), 16) + + +def published(tp, batch): + """The experts publish the latent and WaveMoeTail runs the tail (else MoeTail + front).""" + return batch > (4 if tp == 8 else 8) + + +def twoshot_batches(tp): + table = batch_table(TWOSHOT.get(tp, "")) + return {batch for batch, on in table.items() if on and batch > 8} + + +def early_batches(tp): + table = batch_table(EARLY.get(tp, "")) + return {batch for batch, on in table.items() if on and published(tp, batch)} + + +def umma_batches(tp): + table = batch_table(UMMA.get(tp, "")) + return {batch for batch, on in table.items() if on and batch > 8} + + +def shares(tp, batch, twoshot=False): + """WaveMoeTail clusters share reduced latent rows (all 112 CTAs must be + co-resident).""" + return published(tp, batch) and WaveMoeTail(batch, tp, twoshot).reduce + + +@cache +def resident_clusters(device): + """Co-resident 16-CTA clusters at one CTA per SM.""" + with torch.accelerator.device_index(device): + return cutlass.utils.HardwareInfo(device).get_max_active_clusters(16) + + +class MoeTailBox: + """Model-owned MoE tail mailboxes: latent [3, TP, CAP, 3584] (multicast), shared + reduce-scatter [3, TP, CAP, 7168/TP] (peer stores), output [1, CAP, 7168], the + local all-reduced latent rows [3, CAP, 3584] that WaveMoeTail clusters share and the + two-shot latent buffers (``twoshot_args``/``twoshot_tail`` replace args/tail).""" + + def __init__(self, group, device, tp, rank): + self.rank = rank + self.latent = OprojMailbox(group, device, tp, rank, (3, tp, CAP, 3584)) + self.rs = OprojMailbox(group, device, tp, rank, (3, tp, CAP, 7168 // tp)) + self.out = OprojMailbox(group, device, tp, rank, (1, CAP, 7168)) + ptrs = self.rs.handle.buffer_ptrs + self.peers = torch.tensor(ptrs, dtype=torch.int64, device=device) + self.args = (flat(self.latent.buffer), flat(self.rs.buffer)) + self.tail: tuple = (to_cute(self.peers, 16), flat(self.out.buffer)) + self.tail += (Int64(self.latent.mc), Int64(self.out.mc)) + self.reduced = torch.empty((3, CAP, 3584), dtype=torch.bfloat16, device=device) + self.reduced.view(torch.int32).fill_(-0x80000000) + self.twoshot = box = TwoShot(group, device, tp, rank, 3584, CAP, parities=3) + self.twoshot_args = (flat(box.recv), self.args[1]) + self.twoshot_tail = (*self.tail[:2], Int64(box.gather_mc), self.tail[3]) + # Experts' push: every parity's receive slots [3 x TP, CAP, 3584/TP] and the + # owners' buffer bases (slot parity x TP + rank). + self.push_box = ( + to_cute(box.recv.view(3 * tp, CAP, -1)), + to_cute(box.peers[0], 16), + ) + + +@cache +def compile_tail(tp, batch, device, mode="front", trigger=2, early=False, umma=False): + """front: rows complete at the PDL wait; front_poll: three 0xffff-armed G/W/S parity + slots polled per row and re-armed after use; publish (``published``): WaveMoeTail + over the published mailboxes, reading the weight shard directly, not by TMA; + twoshot: publish over the two-shot receive slots and gather.""" + m8, twoshot = mode in ("publish", "twoshot"), mode == "twoshot" + vec = lambda n, dtype=BFloat16, align=16: fake((n,), None, dtype, align) + slots = 3 * POLL_CAP if mode == "front_poll" else batch + latent = 3 * tp * CAP * (3584 // tp if twoshot else 3584) + tensors: tuple = (vec(latent), vec(3 * tp * CAP * (7168 // tp))) + weights = vec(7168 // tp * 3584) if m8 and not umma else vec(128, Uint8, 64) + tensors += (weights, vec(3584)) + tensors += (vec(slots * 16 * 3584), vec(slots * 16), vec(slots * 7168)) + tensors += (vec(tp, Int64), vec(CAP * 7168)) + scalars: tuple = (Int64(0), Int64(0), Int32(0), Int32(0), Int32(CAP), Float32(1e-6)) + scalars += (Int32(2), Int32(CAP)) + kernel: MoeTail + if m8: + kernel = wave = WaveMoeTail(batch, tp, twoshot, early, umma) + scalars += (vec(3 * CAP * 3584),) if wave.reduce else () + else: + kernel = MoeTail(tp, batch, mode=mode, trigger=trigger) + options = "" if m8 else "--ptxas-options=--maxrregcount=120" + with torch.accelerator.device_index(device): + args = (*tensors, *scalars, make_fake_stream()) + code = cute.compile(kernel, *args, options=options) + # MoE tail requests 100% shared carveout so the next large PDL kernel can + # coexist. + return prepare_shared_carveout(code, device, 100) + + +@cache +def compile_push(batch, device, tp=8): + args = (fake((batch, 7168)), fake((tp,), None, Int64), Int32(0), Int32(0)) + with torch.accelerator.device_index(device): + return cute.compile( + WaveMoeTail(batch, tp).push, *args, Int32(CAP), make_fake_stream() + ) + + +class MoeTailWeights: + """One layer's up-projection shard [7168/TP, 3584], RMS weight and TMA map.""" + + def __init__(self, tp, weight, gamma, eps, ordinal, count, batches, twoshots=()): + self.weight, self.gamma = weight, gamma + # Three parities rotate over the model's MoE layers; each launch re-arms the + # previous MoE layer's slots of its batch (wrapping across decode steps, where + # the first layer re-arms the largest batch the last layer may have published). + parity, rearm = ordinal % 3, (ordinal - 1) % count % 3 + if parity == rearm: # a launch must never re-arm the slots it publishes to + raise ValueError(f"MoE tail parities collide for {count} MoE layers") + self.parity, rearm = Int32(parity), Int32(rearm) + self.scalars = (Int32(CAP), Float32(eps), rearm) + self.wrap = 0 if ordinal else max(batches) + cols = 512 if tp == 8 else 448 + u16 = cuda.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT16 + dims = (cols // 2, 3584 // (cols // 2), 7168 // tp) + spec: tuple = (weight.data_ptr(), u16, dims, (cols, 7168), (cols // 2, 2, 32)) + spec += (cuda.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_NONE,) + spec += (cuda.CUtensorMapL2promotion.CU_TENSOR_MAP_L2_PROMOTION_NONE,) + self.descriptor = tensor_maps([spec], weight.device) + self.views = (to_cute(self.descriptor, 64), to_cute(gamma, 16)) + self.codes, device, early = {}, weight.device.index, early_batches(tp) + umma = umma_batches(tp) + if any(published(tp, b) and b in umma for b in batches): + # WaveMoeTail tcgen05: [64 x lines, 64] boxes of the shard, 128B swizzle + spec = (weight.data_ptr(), u16, (3584, 7168 // tp), (7168,)) + spec += ((64, 128 if tp == 8 else 64),) + spec += (cuda.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_128B,) + spec += (cuda.CUtensorMapL2promotion.CU_TENSOR_MAP_L2_PROMOTION_L2_256B,) + self.umma_map = tensor_maps([spec], weight.device) + self.umma_views = (to_cute(self.umma_map, 64), self.views[1]) + for b in batches: # WaveMoeTail reads the shard directly, not by TMA + twoshot = b in twoshots + mode = "twoshot" if twoshot else "publish" if published(tp, b) else "front" + views = self.views if mode == "front" else (flat(weight), self.views[1]) + on = mode != "front" and b in umma + views = self.umma_views if on else views + code = compile_tail(tp, b, device, mode, early=b in early, umma=on) + self.codes[b] = code, views, shares(tp, b, twoshot), twoshot + shared = any(shares(tp, b, b in twoshots) for b in batches) + if shared and resident_clusters(device) < 7: + raise RuntimeError("WaveMoeTail needs seven co-resident 16-CTA clusters") + + def run(self, box, expanded, weights, shared, code=None, batch=None): + batch = batch or weights.shape[0] + if code is None: + code, views, share, twoshot = self.codes[batch] + else: + views, share, twoshot = self.views, 0, False + args, tail = box.args, box.tail + if twoshot: + args, tail = box.twoshot_args, box.twoshot_tail + rows = (flat(expanded), flat(weights), flat(shared), *tail, self.parity) + rows += (Int32(box.rank), *self.scalars, Int32(self.wrap or batch)) + if share: + rows += (flat(box.twoshot.gather if twoshot else box.reduced),) + code(*args, *views, *rows, stream(expanded.device)) + + +class PreparedTail: + """The layer's MoE tail; placement = (model MoeTailBox, MoE ordinal, MoE layer + count). At WaveMoeTail batches the experts publish into latent slot ``publish[b]`` + (multicast, parity, rank; at ``twoshots`` the two-shot push box) and the side + stream pushes the shared-expert output (``push``).""" + + def __init__(self, op, norm, up, batches, placement): + tp, rank = op.contract.tp_size, op.rank + local_up = up.narrow(0, rank * (7168 // tp), 7168 // tp) + self.tail_box, ordinal, count = placement + eps = op.contract.rms_eps + m8 = [b for b in batches if published(tp, b)] + self.twoshots = twoshot_batches(tp) & set(m8) + args = (tp, local_up, norm, eps, ordinal, count, batches, self.twoshots) + self.tail_weights = MoeTailWeights(*args) + self.peers, _, mc, _ = self.tail_box.tail + self.ring = (self.tail_weights.parity, Int32(rank)) + self.pushers = {b: compile_push(b, norm.device.index, tp) for b in m8} + self.publish = {b: (mc, *self.ring) for b in m8} + slot = Int32(ordinal % 3 * tp + rank) + self.publish.update({b: (*self.tail_box.push_box, slot) for b in self.twoshots}) + + def push(self, shared): + if shared.shape[0] in self.pushers: + args = (to_cute(shared), self.peers, *self.ring, Int32(CAP)) + self.pushers[shared.shape[0]](*args, stream(shared.device)) + + def run(self, expanded, weights, shared, out, *, consume): + self.tail_weights.run(self.tail_box, expanded, weights, shared) + if consume(out) is not out: + raise RuntimeError("The tail output must be consumed from its mailbox ring") + return out diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/__init__.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/consumer.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/consumer.py new file mode 100644 index 0000000000..ad5f2a4cf9 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/consumer.py @@ -0,0 +1,423 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""O-projection all-reduce consumer (``groups`` 128-thread groups poll disjoint +rank slots, group 0 sums in rank order into the AttnRes) and TP16 fused producer + +consumer. With ``twoshot`` the mailbox is the one-slot gather row of +``common.twoshot``: the threads owning this rank's columns first reduce their +receive slots into it (args 11-13: receive view, gather multicast address, rank).""" + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Uint32 + +from ..attn_res.kernel import arm_exchanges, rstd +from ..common.mailbox import rearm +from ..common.primitives import ( + add_rn, + alloc, + asm, + at, + bits_float, + copy_g2s, + float_bits, + fma, + launch, + load_vector, + named_barrier, + pack_pair, + publish_hint, + published, + push_float, + smem_words, + store_relaxed, + store_shared_words, + store_words, + unpack_bf16, + zeros, +) +from ..common.primitives import wait_phase_zero_cluster as wait_phase_zero +from ..common.twoshot import reduce +from .producer import OprojProducer, pack8, tree_reduce + + +@asm((Float32,) * 2, pure=True) +def pair_op(a, b, c=(), *, op="add"): + values = (*a, *b, *c) + code = "{ .reg .b64 a,b,c,d; mov.b64 a,{$2,$3}; mov.b64 b,{$4,$5}; " + if op == "fma": + code += "mov.b64 c,{$6,$7}; fma.rn.f32x2 d,a,b,c; " + else: + code += f"{op}.rn.f32x2 d,a,b; " + constraints = ",".join(["=f", "=f"] + ["f"] * len(values)) + return code + "mov.b64 {$0,$1},d; }", constraints, tuple(map(Float32, values)) + + +@asm(Float32, pure=True) +def mul_rn(a, b): + return "mul.rn.f32 $0,$1,$2;", "=f,f,f", (*(Float32(v) for v in (a, b)),) + + +def two(values, j): + return values[2 * j], values[2 * j + 1] + + +@cute.jit +def dot8(a, b): + s0 = pair_op(two(a, 0), two(b, 0), op="mul") + s1 = pair_op(two(a, 1), two(b, 1), op="mul") + s0 = pair_op(two(a, 2), two(b, 2), s0, op="fma") + s1 = pair_op(two(a, 3), two(b, 3), s1, op="fma") + s = pair_op(s0, s1) + return add_rn(s[0], s[1]) + + +@cute.jit +def reduce_scatter(values, size: cutlass.Constexpr[int], lane): + """Warp-reduce ``size`` values; returns the index of the one this lane holds.""" + for shift in cutlass.range_constexpr(4, -1, -1): + live: cutlass.Constexpr = max(1, size >> (4 - shift)) + half: cutlass.Constexpr = max(1, live // 2) + if cutlass.const_expr(live == 1): + values[0] = add_rn( + values[0], cute.arch.shuffle_sync_bfly(values[0], 1 << shift) + ) + else: + upper = (lane & (1 << shift)) != 0 + for j in cutlass.range_constexpr(half): + send, keep = values[j + half], values[j] + if upper: + send, keep = values[j], values[j + half] + values[j] = add_rn(keep, cute.arch.shuffle_sync_bfly(send, 1 << shift)) + idx = cutlass.Int32(0) + for shift in cutlass.range_constexpr(4, -1, -1): + width: cutlass.Constexpr = size >> (4 - shift) + if cutlass.const_expr(width > 1): # noqa: SIM102 + if (lane & (1 << shift)) != 0: + idx += width // 2 + return idx + + +@cute.jit +def fetch(raw, rank, pointer): + words = load_vector(pointer, cache="volatile") + for j in cutlass.range_constexpr(4): + raw[rank, j] = words[j] + + +class OprojConsumer: + def __init__( + self, + nb, + *, + has_delta=True, + groups=4, + twoshot=False, + ): + if not 0 <= nb <= 8: + raise ValueError("Output-projection AttnRes supports zero to eight blocks") + self.nb, self.has_delta = nb, has_delta + self.groups, self.twoshot = groups, twoshot + self.kpre = 1 << max(0, (2 * nb - 1).bit_length()) + self.npost = 4 if nb > 0 else 1 + + @cute.jit + def wait_exchange(self, barrier): + wait_phase_zero(barrier) # peers complete it by st.async + + @cute.jit + def allocate(self): + smem = cutlass.utils.SmemAllocator() + wpre = alloc(smem, Float32, (4, 16)) + pre = alloc(smem, Float32, (7, 16)) + blk = alloc(smem, Float32, 16) + post = alloc(smem, Float32, (7, 4, 4)) + total = alloc(smem, Float32, 4) + bars = alloc(smem, cutlass.Uint64, 2, 8) + parts = alloc(smem, Float32, (max(1, self.groups - 1), 2, 128, 4)) + lean = self.groups == 1 # above B8: rows for this nb only (co-residency) + pf = alloc(smem, Uint32, (4 + self.nb if lean else 12, 128, 4)) + return wpre, pre, blk, post, total, bars, parts, pf + + @cute.jit + def setup(self, bars): + arm_exchanges(bars, (2 * self.nb, 4 * self.npost)) + cute.arch.cluster_arrive_relaxed() + + @cute.jit + def prefetch(self, args, pf, rank, token): + _, prefix, blocks, norm, query, out_norm = args[:6] + tid, _, _ = cute.arch.thread_idx() + h = (rank * 128 + tid) * 8 + if cutlass.const_expr(self.nb > 0): + copy_g2s(pf.iterator + tid * 4, norm.iterator + h) + copy_g2s(pf.iterator + 512 + tid * 4, query.iterator + h) + copy_g2s(pf.iterator + 1024 + tid * 4, out_norm.iterator + h) + if cutlass.const_expr(self.has_delta): + copy_g2s(pf.iterator + 1536 + tid * 4, at(prefix, token, h)) + for block in cutlass.range_constexpr(self.nb): + dst = pf.iterator + (4 + block) * 512 + tid * 4 + copy_g2s(dst, at(blocks, token, block, h)) + cute.arch.cp_async_commit_group() + + @cute.jit + def fragment(self, pf, row, tid): + words = smem_words(pf.iterator + (row * 128 + tid) * 4) + values = cute.make_rmem_tensor(8, Float32) + for j in cutlass.range_constexpr(4): + values[2 * j], values[2 * j + 1] = unpack_bf16(words[j]) + return values + + @cute.jit + def combined_weight(self, pf, tid): + a, b = self.fragment(pf, 0, tid), self.fragment(pf, 1, tid) + result = cute.make_rmem_tensor(8, Float32) + for j in cutlass.range_constexpr(4): + result[2 * j], result[2 * j + 1] = pair_op(two(a, j), two(b, j), op="mul") + return result + + @cute.jit + def pre(self, state, rank, eps, mix, wait_groups: cutlass.Constexpr[int] = 0): + wpre, pre, blk, _, _, bars, _, pf = state + tid, _, _ = cute.arch.thread_idx() + warp, lane = tid // 32, tid % 32 + maximum, exponent_sum, qpart = Float32(0), Float32(0), Float32(0) + cute.arch.cp_async_wait_group(wait_groups) + if cutlass.const_expr(self.nb > 0): + weight = self.combined_weight(pf, tid) + partial = zeros(self.kpre, Float32) + for block in cutlass.range_constexpr(self.nb): + value = self.fragment(pf, 4 + block, tid) + partial[block] = dot8(value, value) + partial[self.nb + block] = dot8(value, weight) + idx = reduce_scatter(partial, self.kpre, lane) + if ((lane & (32 // self.kpre - 1)) == 0) & (idx < 2 * self.nb): + wpre[warp, idx] = partial[0] + named_barrier() + cute.arch.cluster_wait() + if tid < 2 * self.nb: + pre_sum = Float32(0) + for w in cutlass.range_constexpr(4): + pre_sum = add_rn(pre_sum, wpre[w, tid]) + for peer in cutlass.range_constexpr(7): + push_float(at(pre, rank, tid), pre_sum, bars.iterator, peer) + self.wait_exchange(bars.iterator) + if tid < 2 * self.nb: + sums = [Float32(0), Float32(0)] + for peer in cutlass.range_constexpr(7): + sums[peer % 2] = add_rn(sums[peer % 2], pre[peer, tid]) + blk[tid] = add_rn(*sums) + named_barrier() + logits = [ + mul_rn(blk[self.nb + b], rstd(blk[b], eps)) for b in range(self.nb) + ] + maximum = tree_reduce(logits, max) + exp = [cute.math.exp(x - maximum, fastmath=True) for x in logits] + exponent_sum = tree_reduce(exp, add_rn) + for block in cutlass.range_constexpr(self.nb): + value = self.fragment(pf, 4 + block, tid) + for j in cutlass.range_constexpr(4): + mix[2 * j], mix[2 * j + 1] = pair_op( + (exp[block], exp[block]), two(value, j), two(mix, j), op="fma" + ) + qpart = dot8(mix, mix) + else: + cute.arch.cluster_wait() + return maximum, exponent_sum, qpart + + def slots(self, tp): + """Rank slots each group polls (at least four).""" + return max(4, (tp + self.groups - 1) // self.groups) + + @cute.jit + def poll_sum(self, mailbox, token, h, first, count): + n: cutlass.Constexpr = self.slots(mailbox.shape[0]) + raw = zeros((n, 4), Uint32) + pending = Int32(0) + for rank in cutlass.range_constexpr(n): + if rank < count: + fetch(raw, rank, at(mailbox, first + rank, token, h)) + for rank in cutlass.range_constexpr(n): + if rank < count: # noqa: SIM102 + if not published(raw[rank, j] for j in range(4)): + pending |= 1 << rank + while pending != 0: + for rank in cutlass.range_constexpr(n): + if (pending & (1 << rank)) != 0: + fetch(raw, rank, at(mailbox, first + rank, token, h)) + for rank in cutlass.range_constexpr(n): + if (pending & (1 << rank)) != 0: # noqa: SIM102 + if published(raw[rank, j] for j in range(4)): + pending &= ~(1 << rank) + values = cute.make_rmem_tensor(8, Float32) + for j in cutlass.range_constexpr(4): + pair = unpack_bf16(raw[0, j]) + for rank in cutlass.range_constexpr(1, n): + if rank < count: + pair = pair_op(pair, unpack_bf16(raw[rank, j])) + values[2 * j], values[2 * j + 1] = pair[0], pair[1] + return values + + @cute.jit + def publish(self, output, xmoe, hints, token, h, values): + tid, _, _ = cute.arch.thread_idx() + words = pack8(values) + store_words(at(output, token, h), words) + if cutlass.const_expr(xmoe is not None): + words = pack8(values, "latent") + for rep in cutlass.range_constexpr(xmoe.shape[0]): + store_relaxed(at(xmoe, rep, token, h), words, scope="gpu") + cute.arch.sync_warp() + if tid % 32 == 0: + publish_hint(hints.iterator) + + @cute.jit + def post(self, args, state, block_mix, rank, token): + mailbox, prefix, _, _, _, _, output, xmoe, hints, eps, out_eps = args[:11] + _, _, _, post, total, bars, parts, pf = state + mix, maximum, exponent_sum, qpart = block_mix + thread, _, _ = cute.arch.thread_idx() + group, tid = thread // 128, thread % 128 + warp, lane = tid // 32, tid % 32 + h = (rank * 128 + tid) * 8 + S: cutlass.Constexpr = (mailbox.shape[0] + self.groups - 1) // self.groups + everyone: cutlass.Constexpr = 128 * self.groups + if group > 0: + first = group * S + count = min(S, mailbox.shape[0] - first) + values = self.poll_sum(mailbox, token, h, first, count) + for half in cutlass.range_constexpr(2): + words = tuple(float_bits(values[half * 4 + i]) for i in range(4)) + store_shared_words(at(parts, group - 1, half, tid, 0), words) + named_barrier(2, everyone, arrive=True) + for r in cutlass.range_constexpr(self.slots(mailbox.shape[0])): + if r < count: + rearm(at(mailbox, first + r, token, h)) + else: + if cutlass.const_expr(self.twoshot): + recv, gather_mc, me = args[11:14] + if h // recv.shape[2] == me: + reduce(recv, mailbox, gather_mc, me, token, h) + if cutlass.const_expr(self.has_delta): + old = self.fragment(pf, 3, tid) + attn = self.poll_sum(mailbox, token, h, 0, min(S, mailbox.shape[0])) + if cutlass.const_expr(self.groups > 1): + named_barrier(2, everyone) + for other in cutlass.range_constexpr(self.groups - 1): + for half in cutlass.range_constexpr(2): + words = smem_words(at(parts, other, half, tid, 0)) + for j in cutlass.range_constexpr(2): + k: cutlass.Constexpr = half * 2 + j + part = tuple(map(bits_float, two(words, j))) + attn[2 * k], attn[2 * k + 1] = pair_op(two(attn, k), part) + prefix_value = cute.make_rmem_tensor(8, Float32) + for j in cutlass.range_constexpr(4): + rounded = unpack_bf16(pack_pair(*two(attn, j))) + if cutlass.const_expr(self.has_delta): + rounded = pair_op(two(old, j), rounded) + rounded = unpack_bf16(pack_pair(rounded[0], rounded[1])) + prefix_value[2 * j], prefix_value[2 * j + 1] = rounded[0], rounded[1] + partial = cute.make_rmem_tensor(self.npost, Float32) + if cutlass.const_expr(self.nb > 0): + weight = self.combined_weight(pf, tid) + partial[0], partial[1] = dot8(mix, prefix_value), qpart + partial[2] = dot8(prefix_value, prefix_value) + partial[3] = dot8(prefix_value, weight) + else: + partial[0] = dot8(prefix_value, prefix_value) + idx = reduce_scatter(partial, self.npost, lane) + if idx < self.npost: + for peer in range(lane & (32 // self.npost - 1), 7, 32 // self.npost): + dst = at(post, rank, warp, idx) + push_float(dst, partial[0], bars.iterator + 1, peer) + words = pack8(prefix_value) + store_words(at(prefix, token, h), words) + for r in cutlass.range_constexpr(min(S, mailbox.shape[0])): + rearm(at(mailbox, r, token, h)) + if cutlass.const_expr(xmoe is not None): + cute.arch.fence_acq_rel_gpu() + self.wait_exchange(bars.iterator + 1) + if tid < self.npost: + sums = [Float32(0)] * 4 + for peer in cutlass.range_constexpr(7): + for w in cutlass.range_constexpr(4): + sums[w] = add_rn(sums[w], post[peer, w, tid]) + total[tid] = tree_reduce(sums, add_rn) + named_barrier() + norm = self.fragment(pf, 2, tid) + final = cute.make_rmem_tensor(8, Float32) + if cutlass.const_expr(self.nb > 0): + lp = mul_rn(total[3], rstd(total[2], eps)) + mx = max(maximum, lp) + alpha = cute.math.exp(maximum - mx, fastmath=True) + ep = cute.math.exp(lp - mx, fastmath=True) + den = fma(alpha, exponent_sum, ep) + a2, t1 = mul_rn(alpha, alpha), total[1] + ae2 = mul_rn(mul_rn(2.0, alpha), ep) + quad = fma(a2, t1, fma(ae2, total[0], mul_rn(mul_rn(ep, ep), total[2]))) + scale = rstd(quad, mul_rn(mul_rn(out_eps, den), den)) + ca, cp = mul_rn(alpha, scale), mul_rn(ep, scale) + else: + cp = rstd(total[0], out_eps) + for j in cutlass.range_constexpr(4): + value = pair_op((cp, cp), two(prefix_value, j), op="mul") + if cutlass.const_expr(self.nb > 0): + value = pair_op((ca, ca), two(mix, j), value, op="fma") + value = pair_op(value, two(norm, j), op="mul") + final[2 * j], final[2 * j + 1] = value[0], value[1] + self.publish(output, xmoe, hints, token, h, final) + + @cute.jit + def __call__(self, *args): + launch(self.kernel, args, (7, args[1].shape[0], 1), 128 * self.groups, 7) + + @cute.kernel + def kernel(self, args): + rank, token, _ = cute.arch.block_idx() + tid, _, _ = cute.arch.thread_idx() + state = self.allocate() + self.setup(state[5]) + cute.arch.griddepcontrol_launch_dependents() + mix = zeros(8, Float32) + maximum, exponent_sum, qpart = Float32(0), Float32(0), Float32(0) + if tid < 128: + self.prefetch(args, state[7], rank, token) + maximum, exponent_sum, qpart = self.pre(state, rank, args[9], mix) + self.post(args, state, (mix, maximum, exponent_sum, qpart), rank, token) + cute.arch.griddepcontrol_wait() + + +class OprojFused: + def __init__(self, nb, *, has_delta=True): + self.producer = OprojProducer(768, staged_blocks=0) + # TP16 requires prefix/blocks to be ready at launch. + self.consumer = OprojConsumer(nb, has_delta=has_delta) + + @cute.jit + def __call__(self, *args): + if cutlass.const_expr(args[0].shape[0] > 8): + raise ValueError("The selected fused schedule supports B1..8") + launch(self.kernel, args, (8, 14, 1), 512, 8, min_blocks_per_mp=1) + + @cute.kernel + def kernel(self, args): + x, weight, rows = args[0], args[1], (args[2], *args[5:]) + rank, cluster, _ = cute.arch.block_idx() + tid, _, _ = cute.arch.thread_idx() + cta = cluster * 8 + rank + active = (cluster < x.shape[0]) & (rank < 7) + state = self.consumer.allocate() + smem = self.producer.allocate(x.shape[0]) + self.consumer.setup(state[5]) + cute.arch.griddepcontrol_launch_dependents() + group0 = active & (tid < 128) + if group0: + self.consumer.prefetch(rows, state[7], rank, cluster) + weights = self.producer.load_weight(weight, smem[2], cta) + mix = zeros(8, Float32) + maximum, exponent_sum, qpart = Float32(0), Float32(0), Float32(0) + if group0: + maximum, exponent_sum, qpart = self.consumer.pre(state, rank, rows[9], mix) + self.producer.compute(args, smem, weights, cta, False) + if active: + block_mix = (mix, maximum, exponent_sum, qpart) + self.consumer.post(rows, state, block_mix, rank, cluster) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/prepared.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/prepared.py new file mode 100644 index 0000000000..5f753fc868 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/prepared.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Output projection fused with its TP all-reduce and post-attention AttnRes: rank +partials go into Lamport slots of a multicast [TP, 64, 7168] mailbox (parity alternates +by layer) and the consumer sums and re-arms them. TP16 B<=8 fuses producer and consumer +(MLA layers below B5 keep them split); above B8 the producer runs 16-token groups. +At ``twoshot_batches``, given ``TwoShot`` buffers, the all-reduce is ``common.twoshot`` +(chunks pushed to their column owner, owner sums + multicasts, all Lamport-polled), so +each GPU receives ~2 x 14 KB per token instead of TP x 14 KB.""" + +from functools import cache +from typing import Any + +import cutlass +import cutlass.cute as cute +import torch +import torch.distributed._symmetric_memory as symm_mem +from cutlass.cute.runtime import make_fake_stream + +from ..common.host import dynamic_fake, dynamic_view, fake, stream, to_cute +from ..common.policy import batch_table +from .consumer import OprojConsumer, OprojFused +from .producer import OprojProducer, OprojUmma + + +class OprojMailbox: + """Sentinel-armed multicast mailbox; defaults to [parity, TP, 64, 7168] BF16.""" + + def __init__(self, group, device, tp, rank, shape=None): + self.tp, self.rank = tp, rank + shape = (2, tp, 64, 7168) if shape is None else shape + self.buffer = symm_mem.empty(shape, dtype=torch.bfloat16, device=device) + self.handle = handle = symm_mem.rendezvous(self.buffer, group) + self.mc = int(handle.multicast_ptr or 0) + if not self.mc: # slot() offsets this address + raise RuntimeError( + "Kimi K3 specialized o-projection needs multicast memory" + ) + self.buffer.view(torch.int32).fill_(-0x80000000) + torch.accelerator.synchronize(device) + torch.distributed.barrier(group=group) + + def slot(self, parity): + return self.buffer[parity], self.mc + parity * self.buffer[0].numel() * 2 + + +# Two-shot all-reduce batches per TP (2-/4-node GB300, 32 MLA layers, us/layer, max over +# ranks, multicast vs two-shot: TP8 B16 8.85/11.40, B32 10.63/11.32, B64 18.06/13.18; +# TP16 B16 9.78/12.28, B32 14.96/13.05, B64 28.06/14.07). +TWOSHOT = {8: "33-64:1", 16: "17-64:1"} + + +def twoshot_batches(tp): + table = batch_table(TWOSHOT.get(tp, "")) + return {batch for batch, on in table.items() if on} + + +def engine(batch, tp): + """Producer engine above B8 (GB300): grouped mma.sync for TP8 up to B16, else + tcgen05.""" + default = "mma" if tp == 8 and batch <= 16 else "umma" + return default + + +def producer(k, batch, tp, mla, twoshot=False, kind=None): + """B<=8: one mma.sync pass, TP8 stages 8 (MLA 4) of 12 K blocks per warp. Above + B8 (``kind`` overrides ``engine``): the tcgen05 engine, or 2-tile warps over 8 K + parts in 16-token groups, TP8 staging 2 of 6 K blocks per warp.""" + if batch <= 8: + staged = 0 if tp == 16 else 4 if mla else 8 + return OprojProducer(k, staged_blocks=staged, twoshot=twoshot) + if (kind or engine(batch, tp)) == "umma": + return OprojUmma(k, twoshot=twoshot) + grouped = dict(rows=2, group_tiles=2, twoshot=twoshot) + return OprojProducer(k, staged_blocks=0 if tp == 16 else 2, **grouped) + + +@cache +def compile_oproj( + kind, k, batch, tp, nb, has_delta, mla, front=False, cap=64, twoshot=False +): + """Producer, consumer or fused kernel over a [TP, cap, 7168] parity view; + ``front``: the consumer also publishes x_moe [1, 4, 7168] and the epoch hint. + ``twoshot``: the two-shot pair (arguments as ``TwoShot.views``).""" + box = fake((1 if twoshot else tp, cap, 7168), (cap * 7168, 7168, 1)) + x, w = fake((batch, k), (k, 1)), fake((7168, k), (k, 1)) + mc, slot, eps = cutlass.Int64(0), cutlass.Int32(0), cutlass.Float32(1e-6) + recv, peers = fake((tp, cap, 7168 // tp)), fake((tp,), None, torch.int64) + if kind.startswith("producer"): # producer_mma / producer_umma: engine A/B + kernel = producer(k, batch, tp, mla, twoshot, kind[9:] or None) + tail = (recv, peers, slot) if twoshot else (box, mc, slot) + return cute.compile(kernel, x, w, *tail, make_fake_stream()) + row, weight = dynamic_fake(2), fake((7168,), (1,)) + xmoe = fake((1, 4, 7168), (4 * 7168, 7168, 1)) if front else None + hints = fake((224,), (1,), torch.int32) if front else None + rows = (row, dynamic_fake(3), weight, weight, weight, row, xmoe, hints, eps, eps) + if kind == "consumer": # B64: one 128-thread group per CTA polls every rank + options = dict(has_delta=has_delta) + groups = 4 if batch <= 32 and not twoshot else 1 + kernel = OprojConsumer(nb, groups=groups, twoshot=twoshot, **options) + extra = (recv, mc, slot) if twoshot else () + return cute.compile(kernel, box, *rows, *extra, make_fake_stream()) + kernel = OprojFused(nb, has_delta=has_delta) + return cute.compile(kernel, x, w, box, mc, slot, *rows, make_fake_stream()) + + +class PreparedOproj: + """One decoder layer's o-projection + all-reduce + post-attention AttnRes.""" + + def __init__(self, decoder, mailbox, batches, mla, front=None, twoshot=None): + weight = decoder.self_attn.o_proj.weight + norm, out_norm = decoder.mlp_res_norm, decoder.post_attention_layernorm + weights = (norm.weight, decoder.mlp_res_proj.weight.squeeze(0), out_norm.weight) + self.views = tuple(to_cute(w.reshape(-1)) for w in weights) + self.eps = tuple(cutlass.Float32(n.variance_epsilon) for n in (norm, out_norm)) + parity, rank = decoder.layer_idx % 2, cutlass.Int32(mailbox.rank) + local_box, mc = mailbox.slot(parity) + self.device, self.box = weight.device, to_cute(local_box) + self.head = (self.box, cutlass.Int64(mc), rank) + self.twoshot: Any = None + if twoshot is not None: + recv, peers, gather, gather_mc = twoshot.views(parity) + recv = to_cute(recv) + self.twoshot = ( + (recv, to_cute(peers), rank), + to_cute(gather), + (recv, cutlass.Int64(gather_mc), rank), + ) + self.wview, tp, k = to_cute(weight), mailbox.tp, weight.shape[1] + nb = decoder.prev_valid_blocks + decoder.is_block_write_layer + delta = not decoder.is_block_write_layer + self.front = None if front is None else (*map(to_cute, front[:2]), front[2]) + self.codes, self.twoshots = {}, twoshot_batches(tp) if self.twoshot else set() + for b in batches: + publish = front is not None and b in front[2] + fused = tp == 16 and (not mla or b >= 5) and b <= 8 + kinds = ("fused",) if fused else ("producer", "consumer") + twoshot = b in self.twoshots and not fused + shape = (k, b, tp, nb, delta, mla) + self.codes[b] = tuple( + compile_oproj(kind, *shape, publish and kind != "producer", 64, twoshot) + for kind in kinds + ) + + def run(self, local, prefix, blocks, out): + batch, s = local.shape[0], stream(self.device) + codes, x = self.codes[batch], to_cute(local) + front = (None, None) + if self.front is not None and batch in self.front[2]: + front = self.front[:2] + rows = (dynamic_view(prefix), dynamic_view(blocks), *self.views) + rows += (dynamic_view(out), *front, *self.eps, s) + if len(codes) == 1: + codes[0](x, self.wview, *self.head, *rows) + elif batch in self.twoshots: + head, gather, tail = self.twoshot + codes[0](x, self.wview, *head, s) + codes[1](gather, *rows[:-1], *tail, s) + else: + codes[0](x, self.wview, *self.head, s) + codes[1](self.box, *rows) + return out, prefix, blocks diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/producer.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/producer.py new file mode 100644 index 0000000000..e48397ce4f --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/oproj/producer.py @@ -0,0 +1,443 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""O-projection producer for the split and fused AttnRes paths: 112 CTAs of 64 output +rows. Each warp keeps ``rows`` 16-row weight tiles of one of ``4 * rows`` K parts in +registers (TP8 stages some K blocks in shared memory), runs BF16 MMA chains and the K +parts are reduced in order. Weights load before the PDL wait; the split producer +triggers its consumer right after it. With ``group_tiles`` the batch runs in groups of +that many 8-token tiles: each group's x arrives through a double-buffered cp.async ring +during the previous group's MMAs, and is reduced and published during the next one's. +With ``twoshot`` the partial chunks go to their owners' receive slots instead (the +first shot of ``common.twoshot``; args 2-4: receive view, owners' bases, rank). +``OprojUmma`` is the tcgen05 engine of +the same producer for B>8.""" + +import cutlass +import cutlass.cute as cute +from cutlass import Constexpr, Float32, Int64, Uint8, Uint32, Uint64 + +from ..common.primitives import ( + add_rn, + alloc, + at, + bits_float, + copy_g2s, + fence_tmem, + launch, + load_tmem, + load_vector, + mma_bf16_block, + mma_into, + multicast_words, + named_barrier, + pack_pair, + prefetch_bytes, + release_tmem, + smem_desc, + smem_words, + stage_weight, + store_shared_words, + thread, + view, + wait, +) +from ..common.twoshot import push + + +def pack8(values, policy="plain"): + """Four BF16x2 words of ``values[0:8]``.""" + return tuple( + pack_pair(values[2 * j], values[2 * j + 1], policy=policy) for j in range(4) + ) + + +def tree_reduce(values, op): + tree = list(values) + for stride in (1, 2, 4): + for block in range(0, len(tree) - stride, 2 * stride): + tree[block] = op(tree[block], tree[block + stride]) + return tree[0] + + +class OprojProducer: + def __init__(self, k, *, staged_blocks, rows=1, group_tiles=0, twoshot=False): + self.k, self.staged, self.rows = k, staged_blocks, rows + self.twoshot = twoshot + self.parts = 4 * rows + self.kblocks = k // (32 * self.parts) + self.register_blocks = self.kblocks - self.staged + self.group_tiles = group_tiles + # Single pass: TP8 copies the batch to shared memory, TP16 loads registers. + self.shared_x = self.staged > 0 and not group_tiles + self.register_x = k == 768 and not group_tiles + self.x_stride = k // 8 + 4 + self.chains = max(1, 4 // (rows * max(1, group_tiles))) + + def schedule(self, batches): + """(8-token tiles, tiles per group, groups) of ``batches`` tokens.""" + tiles = (batches + 7) // 8 + per_group = self.group_tiles or tiles + return tiles, per_group, (tiles + per_group - 1) // per_group + + @cute.jit + def allocate(self, batches: cutlass.Constexpr[int]): + nt, gt, ng = self.schedule(batches) + smem = cutlass.utils.SmemAllocator() + red = alloc(smem, Float32, (gt, 4, self.parts, 8, 20)) + rows = nt * 8 if self.shared_x else 0 + if cutlass.const_expr(self.group_tiles): + rows = min(2, ng) * gt * 8 + xs = alloc(smem, Uint32, rows * self.x_stride * 4 if rows else 1) + staged = alloc(smem, Uint32, max(1, 16 * self.staged * self.rows * 2 * 32 * 4)) + return red, xs, staged + + @cute.jit + def load_weight(self, weight, staged, cta): + tid, _, _ = cute.arch.thread_idx() + lane, warp = tid % 32, tid // 32 + group, part = lane // 4, lane % 4 + tile = cta * 4 + warp // self.parts * self.rows + kq = warp % self.parts + base = (tile * 16 + group) * self.k + kq * self.kblocks * 32 + part * 8 + if cutlass.const_expr(self.staged > 0): + for j in cutlass.range_constexpr(self.staged): + for t in cutlass.range_constexpr(self.rows): + unit = (warp * self.staged + j) * self.rows + t + dst = staged.iterator + (unit * 2 * 32 + lane) * 4 + src = weight.iterator + base + t * 16 * self.k + src += (self.register_blocks + j) * 32 + stage_weight(dst, src) + stage_weight(dst + 128, src + 8 * self.k) + cute.arch.cp_async_commit_group() + wr = cute.make_rmem_tensor((self.register_blocks, self.rows, 2, 4), Uint32) + for j in cutlass.range_constexpr(self.register_blocks): + for t in cutlass.range_constexpr(self.rows): + for r in cutlass.range_constexpr(2): + src = weight.iterator + base + (t * 2 + r) * 8 * self.k + j * 32 + words = load_vector(src, cache="weight") + for v in cutlass.range_constexpr(4): + wr[j, t, r, v] = words[v] + return wr + + @cute.jit + def stage_x(self, x, xs, first: Constexpr[int], buffer: Constexpr[int]): + """One cp.async group: x tokens [first, first + group) into ring ``buffer``.""" + tid, _, _ = cute.arch.thread_idx() + chunks: Constexpr = self.k // 8 + count: Constexpr = min(x.shape[0] - first, self.group_tiles * 8) + rows: Constexpr = buffer * self.group_tiles * 8 + for i in cutlass.range_constexpr((count * chunks + 511) // 512): + chunk = tid + i * 512 + token, col = chunk // chunks, chunk % chunks + dst = xs.iterator + ((rows + token) * self.x_stride + col) * 4 + if cutlass.const_expr((i + 1) * 512 <= count * chunks): + copy_g2s(dst, at(x, first + token, col * 8)) + else: + if chunk < count * chunks: # noqa: SIM102 + copy_g2s(dst, at(x, first + token, col * 8)) + cute.arch.cp_async_commit_group() + + def x_words(self, xs, row, kofs, j): + return smem_words(xs.iterator + (row * self.x_stride + kofs // 8 + j * 4) * 4) + + @cute.jit + def multiply(self, smem, wr, xr, first: Constexpr, buffer: Constexpr, M: Constexpr): + """FP32 chains of tokens [first, first + group); x from ``xr`` (registers) or + shared-memory rows (ring ``buffer``).""" + red, xs, staged = smem + _, gt, _ = self.schedule(M) + tiles: Constexpr = min(gt, (M - first + 7) // 8) + tid, _, _ = cute.arch.thread_idx() + lane, warp = tid % 32, tid // 32 + group, part = lane // 4, lane % 4 + kofs = warp % self.parts * self.kblocks * 32 + part * 8 + acc = cute.make_rmem_tensor((gt, self.rows, self.chains, 4), Float32) + acc.fill(0) + for j in cutlass.range_constexpr(self.kblocks): + wa, wb = [], [] + for t in cutlass.range_constexpr(self.rows): + if cutlass.const_expr(j < self.register_blocks): + wa.append(tuple(wr[j, t, 0, v] for v in range(4))) + wb.append(tuple(wr[j, t, 1, v] for v in range(4))) + else: + sj: Constexpr = j - self.register_blocks + unit = (warp * self.staged + sj) * self.rows + t + ptr = staged.iterator + (unit * 2 * 32 + lane) * 4 + wa.append(smem_words(ptr)) + wb.append(smem_words(ptr + 128)) + for nt in cutlass.range_constexpr(tiles): + if cutlass.const_expr(self.register_x): + xv = tuple(xr[nt, j, v] for v in range(4)) + else: + row = buffer * gt * 8 + nt * 8 + if cutlass.const_expr(self.group_tiles): + xv = self.x_words(xs, row + group, kofs, j) + else: + xv = (Uint32(0), Uint32(0), Uint32(0), Uint32(0)) + if nt * 8 + group < M: + xv = self.x_words(xs, row + group, kofs, j) + for t in cutlass.range_constexpr(self.rows): + for half in cutlass.range_constexpr(2): + chain: Constexpr = (j * 2 + half) % self.chains + a: tuple = (wa[t][2 * half], wb[t][2 * half]) + a += (wa[t][2 * half + 1], wb[t][2 * half + 1]) + b = (xv[2 * half], xv[2 * half + 1]) + mma_into(acc[nt, t, chain, None], a, b, bf16=True) + return acc + + @cute.jit + def store_partials(self, red, acc, tiles: Constexpr): + tid, _, _ = cute.arch.thread_idx() + lane, warp = tid % 32, tid // 32 + group, part = lane // 4, lane % 4 + tile, kq = warp // self.parts * self.rows, warp % self.parts + for nt in cutlass.range_constexpr(tiles): + for t in cutlass.range_constexpr(self.rows): + for v in cutlass.range_constexpr(4): + chains = [acc[nt, t, c, v] for c in range(self.chains)] + token, row = 2 * part + v % 2, group + (v // 2) * 8 + red[nt, tile + t, kq, token, row] = tree_reduce(chains, add_rn) + + @cute.jit + def publish(self, args, red, cta, first: Constexpr[int]): + """Sum the K parts of tokens [first, first + group) in order; multicast them.""" + x, _, mailbox, mc, slot = args[:5] + _, gt, _ = self.schedule(x.shape[0]) + count: Constexpr = min(x.shape[0] - first, gt * 8) + tid, _, _ = cute.arch.thread_idx() + if tid < 8 * count: + reduced = cute.make_rmem_tensor(8, Float32) + if cutlass.const_expr(self.group_tiles): + # Token-major lanes: conflict-free 16-byte partial loads. + token, out_half = tid % count, tid // count % 2 + row_tile = tid // count // 2 + tile, row = token // 8, token % 8 + for i in cutlass.range_constexpr(2): + column = out_half * 8 + i * 4 + words = smem_words(at(red, tile, row_tile, 0, row, column)) + sums = [bits_float(w) for w in words] + for q in cutlass.range_constexpr(1, self.parts): + words = smem_words(at(red, tile, row_tile, q, row, column)) + sums = [add_rn(s, bits_float(w)) for s, w in zip(sums, words)] + for v in cutlass.range_constexpr(4): + reduced[i * 4 + v] = sums[v] + else: + out_half, row_tile, token = tid % 2, (tid // 2) % 4, tid // 8 + tile, row = token // 8, token % 8 + for v in cutlass.range_constexpr(8): + value = red[tile, row_tile, 0, row, out_half * 8 + v] + for q in cutlass.range_constexpr(1, self.parts): + part = red[tile, row_tile, q, row, out_half * 8 + v] + value = add_rn(value, part) + reduced[v] = value + words = pack8(reduced, "latent") + out_col = (cta * 4 + row_tile) * 16 + out_half * 8 + if cutlass.const_expr(self.twoshot): + push(mailbox, mc, slot, first + token, out_col, words) + else: + offset = cute.crd2idx((slot, first + token, out_col), mailbox.layout) + multicast_words(mc + Int64(offset) * 2, words) + + @cute.jit + def compute(self, args, smem, wr, cta, trigger_after_wait: cutlass.Constexpr[bool]): + x = args[0] + red, xs, staged = smem + M: Constexpr = x.shape[0] + NT, GT, NG = self.schedule(M) + tid, _, _ = cute.arch.thread_idx() + lane, warp = tid % 32, tid // 32 + group, part = lane // 4, lane % 4 + kofs = warp % self.parts * self.kblocks * 32 + part * 8 + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(trigger_after_wait): + cute.arch.griddepcontrol_launch_dependents() + xr = cute.make_rmem_tensor((NT, self.kblocks, 4), Uint32) + if cutlass.const_expr(self.group_tiles): + self.stage_x(x, xs, 0, 0) + if cutlass.const_expr(self.shared_x): + for chunk in range(tid, M * (self.k // 8), 512): + token, col = chunk // (self.k // 8), chunk % (self.k // 8) + words = load_vector(at(x, token, col * 8)) + store_shared_words( + xs.iterator + (token * self.x_stride + col) * 4, words + ) + cute.arch.sync_threads() + if cutlass.const_expr(self.register_x): + xr.fill(0) + for nt in cutlass.range_constexpr(NT): + token = nt * 8 + group + if token < M: + for j in cutlass.range_constexpr(self.kblocks): + words = load_vector(at(x, token, kofs + j * 32)) + for v in cutlass.range_constexpr(4): + xr[nt, j, v] = words[v] + if cutlass.const_expr(self.staged > 0 and not self.group_tiles): + cute.arch.cp_async_wait_group(0) + for g in cutlass.range_constexpr(NG): + if cutlass.const_expr(self.group_tiles): + # x(g) and group g - 1's partials are visible: stage x(g + 1), publish. + cute.arch.cp_async_wait_group(0) + cute.arch.sync_threads() + if cutlass.const_expr(g + 1 < NG): + self.stage_x(x, xs, (g + 1) * GT * 8, (g + 1) % 2) + if cutlass.const_expr(g > 0): + self.publish(args, red, cta, (g - 1) * GT * 8) + acc = self.multiply(smem, wr, xr, g * GT * 8, g % 2, M) + if cutlass.const_expr(self.group_tiles): + cute.arch.sync_threads() + self.store_partials(red, acc, min(GT, NT - g * GT)) + cute.arch.sync_threads() + self.publish(args, red, cta, (NG - 1) * GT * 8) + + @cute.jit + def __call__(self, *args): + launch(self.kernel, args, 112, 512, min_blocks_per_mp=1) + + @cute.kernel + def kernel(self, args): + cta, _, _ = cute.arch.block_idx() + smem = self.allocate(args[0].shape[0]) + wr = self.load_weight(args[1], smem[2], cta) + self.compute(args, smem, wr, cta, True) + + +class OprojUmma(OprojProducer): + """B>8 producer on tcgen05: 112 CTAs x 64 rows, tokens on N. Four loader warps + stream 128B-swizzled W [64, 64] and x [N, 64] K blocks through a cp.async ring + (W of the first stages before the PDL wait, the rest prefetched to L2), the MMA + warp accumulates D [64, N] FP32 in TMEM with M=64 kind::f16 MMAs, and the loader + warps round D to BF16 and publish it like ``OprojProducer``.""" + + def __init__(self, k, *, twoshot=False, budget=184 * 1024): + super().__init__(k, staged_blocks=0, twoshot=twoshot) + self.blocks, self.budget = k // 64, budget + + def geometry(self, batches): + """(N, ring stage bytes, stages, TMEM columns).""" + n = (batches + 7) // 8 * 8 + stage = 8192 + n * 128 + stages = min(self.blocks, self.budget // stage) + return n, stage, stages, max(32, 1 << (n - 1).bit_length()) + + @cute.jit + def load(self, source, ring, rows: Constexpr, valid: Constexpr, row0, column): + """cp.async ``rows`` x 64 BF16 of ``source`` from (row0, column) into a K-major + 128B-swizzled tile; rows >= ``valid`` are left alone. ``valid`` < 0: weights + (evict-first, every row).""" + tid, _, _ = cute.arch.thread_idx() + for i in cutlass.range_constexpr((rows * 8 + 127) // 128): + c = tid + 128 * i + row, chunk = c // 8, c % 8 + dst = ring + row // 8 * 1024 + row % 8 * 128 + (chunk ^ row % 8) * 16 + src = at(source, row0 + row, column + chunk * 8) + if cutlass.const_expr(valid < 0): + stage_weight(dst, src) + elif cutlass.const_expr((i + 1) * 128 <= valid * 8): + copy_g2s(dst, src) + else: + if c < valid * 8: # noqa: SIM102 + copy_g2s(dst, src) + + @cute.jit + def multiply(self, ring, bars, tmem, n: Constexpr, stage: Constexpr): + """The MMA warp: wait for each K block, issue its four K16 MMAs from the + elected lane and release its stage (the last block: the accumulator).""" + stages: Constexpr = self.geometry(n)[2] + # kind::f16 descriptor: FP32 D, BF16 A and B, both K-major, N, M = 64 + instruction: Constexpr = 1 << 4 | 1 << 7 | 1 << 10 | n >> 3 << 17 | 4 << 24 + for kb in cutlass.range_constexpr(self.blocks): + s: Constexpr = kb % stages + wait(bars.iterator + s, kb // stages % 2) + cute.arch.fence_proxy("async.shared", space="cta") + fence_tmem() + a = smem_desc(ring.iterator + s * stage, True) + b = smem_desc(ring.iterator + s * stage + 8192, True) + last: Constexpr = kb + 1 == self.blocks + release: Constexpr = last or kb + stages < self.blocks + barrier = bars.iterator + (2 * stages if last else stages + s) + mma_bf16_block(tmem, a, b, instruction, int(kb > 0), barrier, int(release)) + + @cute.jit + def epilogue(self, args, ring, tmem, cta, n: Constexpr): + """D rows 16 * warp + lane (lane < 16: M=64 TMEM layout) -> FP32 [token, row] + tile in the drained ring -> 8-row BF16 chunks to the mailbox / partials.""" + x, _, mailbox, mc, slot = args[:5] + M: Constexpr = x.shape[0] + tid, warp, lane, _ = thread() + values = load_tmem(tmem + (32 * warp << 16), self.geometry(M)[3]) + tile = view(ring.iterator, 0, Float32, (n, 68), (68, 1)) # [token, row] + if lane < 16: + for t in cutlass.range_constexpr(M): + tile[t, 16 * warp + lane] = values[t] + named_barrier(1, 128) + for i in cutlass.range_constexpr((M * 8 + 127) // 128): + c = tid + 128 * i + if c < M * 8: + token, part = c // 8, c % 8 + reduced = cute.make_rmem_tensor(8, Float32) + for h in cutlass.range_constexpr(2): + words = smem_words(at(tile, token, part * 8 + h * 4)) + for v in cutlass.range_constexpr(4): + reduced[h * 4 + v] = bits_float(words[v]) + words = pack8(reduced, "latent") + column = cta * 64 + part * 8 + if cutlass.const_expr(self.twoshot): + push(mailbox, mc, slot, token, column, words) + else: + offset = cute.crd2idx((slot, token, column), mailbox.layout) + multicast_words(mc + Int64(offset) * 2, words) + + @cute.jit + def __call__(self, *args): + launch(self.kernel, args, 112, 160, min_blocks_per_mp=1) + + @cute.kernel + def kernel(self, args): + x, weight = args[0], args[1] + n, stage, stages, columns = self.geometry(x.shape[0]) + tid, warp, _, cta = thread() + smem = cutlass.utils.SmemAllocator() + ring = alloc(smem, Uint8, stages * stage, 1024) + bars = alloc(smem, Uint64, 2 * stages + 1, 8) + tmem_slot = alloc(smem, Uint32, 1) + if tid == 0: + for s in cutlass.range_constexpr(stages): + cute.arch.mbarrier_init(bars.iterator + s, 256) + cute.arch.mbarrier_init(bars.iterator + stages + s, 1) + cute.arch.mbarrier_init(bars.iterator + 2 * stages, 1) + cute.arch.mbarrier_init_fence() + if warp == 4: + cute.arch.alloc_tmem(columns, tmem_slot.iterator) + cute.arch.relinquish_tmem_alloc_permit() + fence_tmem(before=True) + cute.arch.sync_threads() + fence_tmem() + tmem = tmem_slot[0] + row0 = cta * 64 + if warp < 4: # weights do not depend on the previous kernel + if cutlass.const_expr(stages < self.blocks) and tid < 64: + rest = at(weight, row0 + tid, stages * 64) + prefetch_bytes(rest, (self.blocks - stages) * 128) + for s in cutlass.range_constexpr(stages): + self.load(weight, ring.iterator + s * stage, 64, -1, row0, s * 64) + cute.arch.cp_async_mbarrier_arrive_noinc(bars.iterator + s) + cute.arch.griddepcontrol_wait() + cute.arch.griddepcontrol_launch_dependents() + if warp < 4: + for kb in cutlass.range_constexpr(self.blocks): + slot: Constexpr = kb % stages + tile = ring.iterator + slot * stage + if cutlass.const_expr(kb >= stages): + wait(bars.iterator + stages + slot, (kb // stages - 1) % 2) + self.load(weight, tile, 64, -1, row0, kb * 64) + cute.arch.cp_async_mbarrier_arrive_noinc(bars.iterator + slot) + self.load(x, tile + 8192, n, x.shape[0], 0, kb * 64) + cute.arch.cp_async_mbarrier_arrive_noinc(bars.iterator + slot) + wait(bars.iterator + 2 * stages, 0) + fence_tmem() + # The MMAs (async proxy) read the ring the epilogue now overwrites. + cute.arch.fence_proxy("async.shared", space="cta") + self.epilogue(args, ring, tmem, cta, n) + else: + self.multiply(ring, bars, tmem, n, stage) + release_tmem(tmem, columns, warp, 4) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/__init__.py b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/__init__.py new file mode 100644 index 0000000000..dccfb9a5cb --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/__init__.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kimi K3 specialized decode: a CuTe decode executor over the stock backbone, enabled +by ``additional_config={"decode_backend": "latency_specialized"}``.""" + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class DecodeConfig: + max_context: int + mla_core: str + dcp: int + # Tokens per request in a verify step (1 + draft tokens; 1 without speculation). + q: int + capture_sizes: tuple[int, ...] + + +def specialized_decode_config(vllm_config) -> DecodeConfig | None: + """The deployment's specialized decode settings, None for the stock decode. + Raises when the specialized decode is requested for an unsupported deployment.""" + extra = vllm_config.additional_config + backend = extra.get("decode_backend", "auto") if isinstance(extra, dict) else "auto" + if backend == "auto": + return None + if backend != "latency_specialized": + raise ValueError("Supported decode_backend values are auto/latency_specialized") + reason = unsupported(vllm_config) + if reason is not None: + raise ValueError(f"Kimi K3 specialized decode does not support {reason}") + spec = vllm_config.speculative_config + return DecodeConfig( + max_context=vllm_config.model_config.max_model_len, + mla_core=extra.get("mla_core", "auto"), + dcp=vllm_config.parallel_config.decode_context_parallel_size, + q=1 + spec.num_speculative_tokens if spec is not None else 1, + capture_sizes=tuple( + vllm_config.compilation_config.cudagraph_capture_sizes or () + ), + ) + + +def kda_dtypes(vllm_config): + """(conv, recurrent) KDA state dtypes the deployment's cache config selects.""" + from vllm.model_executor.layers.mamba.mamba_utils import MambaStateDtypeCalculator + + cache = vllm_config.cache_config + return MambaStateDtypeCalculator.kda_state_dtype( + vllm_config.model_config.dtype, + cache.mamba_cache_dtype, + cache.mamba_ssm_cache_dtype, + ) + + +def unsupported(vllm_config) -> str | None: + import torch + + from vllm import envs + from vllm.config import CompilationMode, CUDAGraphMode + from vllm.platforms import current_platform + + parallel, spec = vllm_config.parallel_config, vllm_config.speculative_config + text = vllm_config.model_config.hf_text_config + moe = tuple( + getattr(text, name, None) + for name in ( + "hidden_act", + "activation_situ_beta", + "activation_situ_linear_beta", + "num_experts", + "num_experts_per_token", + "num_expert_group", + "topk_group", + "moe_router_activation_func", + ) + ) + compilation = vllm_config.compilation_config + graphs = compilation.cudagraph_mode + checks = ( + (not current_platform.is_device_capability_family(100), "this GPU (SM100/103)"), + ( + moe != ("situ", 4.0, 25.0, 896, 16, 1, 1, "sigmoid"), + "MoE configs but the one-group sigmoid top-16 SiTU MoE of Kimi K3", + ), + (parallel.tensor_parallel_size not in (8, 16), "TP other than 8/16"), + (parallel.pipeline_parallel_size != 1, "pipeline parallelism"), + (parallel.enable_expert_parallel, "expert parallelism"), + (parallel.use_ubatching, "DBO"), + ( + parallel.decode_context_parallel_size + not in (1, parallel.tensor_parallel_size), + "decode context parallel other than 1 or TP", + ), + (vllm_config.model_config.enable_sleep_mode, "sleep mode"), + (vllm_config.lora_config is not None, "LoRA"), + # The runner's ModelState hooks build the executor and prepare the KV layouts. + (not vllm_config.use_v2_model_runner, "the V1 model runner"), + # The decode loop runs in forward; a compiled forward would trace it. + (compilation.mode == CompilationMode.STOCK_TORCH_COMPILE, "torch.compile"), + # The loop runs only in FULL decode graphs; eager and PIECEWISE steps run stock. + ( + graphs is None or graphs.decode_mode() != CUDAGraphMode.FULL, + "deployments without FULL decode CUDA graphs (e.g. enforce-eager)", + ), + ( + spec is not None and spec.method != "dspark", + "speculative methods but DSpark", + ), + ( + spec is not None and spec.enable_adaptive_verification, + "adaptive verification", + ), + # KdaSplit keeps a BF16 conv state and an FP32 recurrent state in place; it + # maintains neither other cache dtypes nor ReplaySSM / RecoverSSM records. + ( + kda_dtypes(vllm_config) != (torch.bfloat16, torch.float32), + "KDA cache dtypes other than a BF16 conv state and an FP32 recurrent " + "state (mamba cache dtype)", + ), + ( + vllm_config.cache_config.use_replayssm + or vllm_config.cache_config.use_kda_recoverssm, + "ReplaySSM / RecoverSSM (--use-replayssm)", + ), + # The verify kernels take at most 8 tokens per request (KdaSplit's barriers). + ( + spec is not None and spec.num_speculative_tokens > 7, + "more than 7 speculative tokens", + ), + ( + spec is not None and envs.VLLM_KIMI_K3_AUX_ATTN_RES_STREAM, + "the AttnRes aux stream", + ), + ) + return next((reason for failed, reason in checks if failed), None) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/decode.py b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/decode.py new file mode 100644 index 0000000000..e1563d38cb --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/decode.py @@ -0,0 +1,280 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The specialized decode executor over the stock Kimi K3 backbone: one routing +decision per step, then every layer's prepared providers in one loop.""" + +from functools import partial +from typing import Any, NamedTuple + +import torch + +from vllm.config import CUDAGraphMode +from vllm.distributed import get_tp_group +from vllm.distributed import tensor_model_parallel_all_reduce as all_reduce +from vllm.forward_context import get_forward_context, is_forward_context_available +from vllm.platforms import current_platform +from vllm.utils.math_utils import round_up + +from ..cutedsl.attn_res.prepared import PreparedAttnRes +from ..cutedsl.common.embedding import PreparedEmbedding +from ..cutedsl.common.twoshot import TwoShot +from ..cutedsl.moe.dense import PreparedDenseMLP +from ..cutedsl.moe.experts_runtime import PersistentExperts +from ..cutedsl.moe.front import SharedFront +from ..cutedsl.moe.prepared import MoEScratch, PreparedMoE +from ..cutedsl.moe.small_batch import chain as small_batch_chain +from ..cutedsl.moe.tail_runtime import MoeTailBox, PreparedTail +from ..cutedsl.oproj.prepared import OprojMailbox +from .prefetch import PrefetchController +from .providers import BATCHES, PREPARED_MAX, Attention, pack_front + + +class Step(NamedTuple): + q: int # tokens per request + metadata: list # per layer + slot_mapping: Any + + +def layer_name(module): + return getattr(module, "prefix", getattr(module, "layer_name", None)) + + +def decode_metadata(attention, context, q, b): + """One layer's metadata if this step is a pure decode (q = 1) or a uniform + verify step of b tokens, else None.""" + md = context.attn_metadata + md = md.get(layer_name(attention)) if isinstance(md, dict) else None + if md is None: + return None + if q > 1: + return md if is_uniform_verify(md, q, b) else None + if md.num_prefills or getattr(md, "num_spec_decodes", 0): + return None + return md + + +def is_uniform_verify(md, q, b): + """Whether every request of the b // q graph rows verifies q tokens (FULL graphs + pad whole requests: KDA stages NULL state rows, MLA zero-length decodes).""" + if md.num_prefills: + return False + if not hasattr(md, "num_spec_decodes"): + return md.num_decode_tokens == md.num_decodes * q == b + rows, tokens = md.spec_state_indices_tensor, md.num_spec_decode_tokens + full = rows is not None and rows.shape[0] * q == b + return full and not md.num_decodes and tokens == md.num_spec_decodes * q + + +def caches_of(attention): + cache = getattr(attention, "kv_cache", ()) + return [cache] if isinstance(cache, torch.Tensor) else list(cache) + + +class SpecializedDecode: + """Built once the weights are loaded; owns every prepared provider, the + mailboxes and the decode buffers. Reads the backbone, never modifies it.""" + + def __init__(self, backbone, config): + self.backbone, self.buffers = backbone, {} + q = config.q + # Verify steps decode q tokens per request at M = requests * q tokens; FULL + # verify graphs round each capture size up to whole requests. + rounded = {round_up(s, q) for s in config.capture_sizes} + verify = {m for m in rounded if m <= PREPARED_MAX} + self.batches, self.q = tuple(sorted({*BATCHES, *verify})), q + self.queries = (1, q) if q > 1 else (1,) + # The longest sequence the deployment serves bounds the MLA split count. + self.max_context, self.mla_core = config.max_context, config.mla_core + # Decode context parallel over the TP ranks: MLA KV is sequence-sharded. + self.dcp, self.dcp_box, self.mla_layers = config.dcp, None, 0 + group = get_tp_group() + self.tp, self.rank = group.world_size, group.rank_in_group + self.group = group.device_group + self.device = torch.device(current_platform.current_device()) + + layers = backbone.layers[backbone.start_layer : backbone.end_layer] + self.layers = layers = tuple(layers) + self.build_providers() + self.embedding = PreparedEmbedding(backbone.embed_tokens, self.tp, self.rank) + self.before = { + d.layer_idx: PreparedAttnRes.for_target(d, has_delta=d is not layers[0]) + for d in layers + } + # The final AttnRes also applies the backbone's final norm. + self.final = PreparedAttnRes.for_target(backbone, final=True) + # The next (or final) AttnRes consumes each MoE output from its MoE tail output + # row. + self.fused_before = {} + followers = [self.before[f.layer_idx] for f in layers[1:]] + [self.final] + for layer, consumer in zip(layers, followers): + moe = self.feed_forward[layer.layer_idx] + if isinstance(moe, PreparedMoE): + consumer.bind_ring(moe.tail.tail_box.out.buffer) + self.fused_before[layer.layer_idx] = consumer + # The AttnRes entering each aux layer (or the final one) also stores the + # updated prefix: the stock aux stream (prefix sum plus the last MLP output). + self.aux_layers = tuple(backbone.aux_hidden_state_layers) + end = layers[-1].layer_idx + 1 + for index in self.aux_layers: + (self.final if index == end else self.before[index]).bind_aux() + empty = partial(torch.empty, device=self.device, dtype=torch.bfloat16) + self.hidden = (empty((64, 7168)), empty((64, 7168))) + self.prefix = empty((64, 7168)) + self.blocks = empty((64, backbone.num_attn_res_blocks, 7168)) + + def build_providers(self): + """Every layer's attention and feed-forward providers, the model's mailboxes + and the prefetch side stream. All collective allocations, weight packing and + kernel compiles finish here, before KV binding and CUDA graph capture.""" + config, layers = self.backbone.config, self.layers + beta = config.activation_situ_beta + linear_beta = config.activation_situ_linear_beta + tp, rank, device, batches = self.tp, self.rank, self.device, self.batches + owner = self.group, device, tp, rank + packed = { + d.layer_idx: pack_front(d.mlp, rank, tp) for d in layers if d.is_moe_layer + } + prefetch = self.prefetch = PrefetchController(self, layers, packed) + state = MoEScratch(tp, device, batches) + mailbox, tail_box, moe_layers = ( + OprojMailbox(*owner), + MoeTailBox(*owner), + list(packed), + ) + twoshot = TwoShot(*owner) + front_box = OprojMailbox(*owner, (1, 64, 3584 + 2 * 896)) + small_batch = small_batch_chain.SmallBatchBox(*owner) + self.attention: dict[int, Attention] = {} + self.feed_forward: dict[int, Any] = {} + for layer in layers: + index, module = layer.layer_idx, layer.mlp + boxes = (mailbox, small_batch, twoshot) + self.attention[index] = Attention(layer, self, prefetch, *boxes) + if not layer.is_moe_layer: + reduce = all_reduce if module.down_proj.reduce_results else None + self.feed_forward[index] = PreparedDenseMLP(module, batches, reduce) + continue + args = (module, packed[index], front_box, state.front, batches, beta) + shared = dict(plan=state.plan, counters=state.shared, q=self.q) + front = SharedFront(*args, linear_beta, **shared) + routed = module.experts.routed_experts + w13, w2 = ( + w.view(torch.uint8) for w in (routed.w13_weight, routed.w2_weight) + ) + weights = w13, routed.w13_weight_scale, w2, routed.w2_weight_scale + experts = PersistentExperts(*weights, beta=beta, linear_beta=linear_beta) + norm = module.routed_expert_norm.weight + up = module.routed_expert_up_proj.weight + op = module.experts._k3_latent_moe_tail_op + placement = tail_box, moe_layers.index(index), len(moe_layers) + tail = PreparedTail(op, norm, up, batches, placement) + moe = PreparedMoE( + front, experts, tail, state, prefetch=prefetch, layer_idx=index + ) + moe.small_batch = small_batch_chain.SmallBatchChain( + small_batch, module, experts, tail, index, beta, linear_beta, prefetch + ) + self.feed_forward[index] = moe + + def prepare_kv(self): + """Compile each layer's cores for its bound KV cache layout (the profiling + cache, then the real one).""" + for layer in self.layers: + caches = caches_of(layer.self_attn) + self.attention[layer.layer_idx].prepare_layout(caches) + + def embed_input_ids(self, ids): + rows = ids.shape[0] if ids.ndim == 1 and ids.is_contiguous() else 0 + if not 1 <= rows <= PREPARED_MAX: + return self.backbone.embed_input_ids(ids) + table = self.backbone.embed_tokens.weight + return self.embedding.run(ids, table.new_empty((ids.shape[0], table.shape[1]))) + + def plan(self, positions, intermediate_tensors): + """This step's decode plan, or None to run the stock forward instead. + + Steps that run the stock forward: + - pipeline-parallel stage inputs (``intermediate_tensors``); + - token counts without prepared providers (above 64 or not in ``batches``); + - steps outside a CUDA graph capture (profiling, warmup and the warmup run + before each capture), so every stock kernel runs before its capture; + - PIECEWISE steps, keyed by token count alone and so captured as stock; + - steps with prompt rows (prefills, chunked prefills, fresh prompts); + - speculative steps whose requests do not all verify q tokens (partial + drafts, verify rows mixed with plain decode rows). + + Every predicate is the same on all TP ranks: the mailboxes of the specialized + decode need every rank to take the same path. + """ + b = positions.shape[0] + if intermediate_tensors is not None or b not in self.batches: + return None + if not torch.cuda.is_current_stream_capturing(): + return None + if not is_forward_context_available(): + return None + context = get_forward_context() + if context.cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE: + return None + by_layer = context.attn_metadata + rows: list[Any] = list(by_layer.values()) if isinstance(by_layer, dict) else [] + q = self.q if any(getattr(md, "num_spec_decodes", 0) for md in rows) else 1 + if b % q: + return None + metadata = [decode_metadata(d.self_attn, context, q, b) for d in self.layers] + if any(md is None for md in metadata): + return None + return Step(q, metadata, context.slot_mapping) + + def run(self, step, input_ids, inputs_embeds): + """Normalized hidden states (and the aux streams) of a planned step.""" + x = inputs_embeds + if x is None: + x = self.embed_input_ids(input_ids) + out = torch.empty_like(x) + # The stock aux streams: the AttnRes prefix entering each aux layer. + aux_out = [torch.empty_like(x) for _ in self.aux_layers] + aux = dict(zip(self.aux_layers, aux_out)) + b, q, slots = x.shape[0], step.q, step.slot_mapping + prefix, blocks = self.prefix[:b], self.blocks[:b] + first, second = self.hidden[0][:b], self.hidden[1][:b] + hidden, pre_ready = x, False + for layer, md in zip(self.layers, step.metadata): + index, caches = layer.layer_idx, caches_of(layer.self_attn) + if not pre_ready: + self.before[index].run(hidden, prefix, blocks, first, aux.get(index)) + self.attention[index].run( + first, second, prefix, blocks, caches, md, slots, q + ) + consumer = self.fused_before.get(index) + pre_ready = consumer is not None + if consumer is not None: + target = out if consumer is self.final else first + stream = aux.get(index + 1) + consume = partial(consumer.run_ring, prefix, blocks, aux=stream) + hidden = self.feed_forward[index].run(second, target, consume=consume) + else: + hidden = self.feed_forward[index].run(second, first) + first, second = second, first + if not pre_ready: + self.final.run(hidden, prefix, blocks, out, aux.get(index + 1)) + self.prefetch.join() + return (out, aux_out) if aux_out else out + + def normalize(self, output): + """A stock step's output with the backbone's final norm applied, as the + specialized decode's last AttnRes applies it; the aux streams stay pre-norm. + + Runs in chunks of at most 255 rows: the stock RMSNorm kernel reduces with + 1024-thread blocks below 256 rows and 256-thread blocks above, so chunking keeps + the reduction order ``compute_logits`` uses for the few sampled rows, and + prefill logits match stock. + """ + hidden, aux = output if isinstance(output, tuple) else (output, None) + norm = self.backbone.norm + if hidden.ndim == 2 and hidden.shape[0] >= 256: + chunks = range(0, hidden.shape[0], 255) + hidden = torch.cat([norm(hidden[i : i + 255], None) for i in chunks]) + else: + hidden = norm(hidden, None) + return hidden if aux is None else (hidden, aux) diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/prefetch.py b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/prefetch.py new file mode 100644 index 0000000000..8ca3e1e527 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/prefetch.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""L2 weight prefetch for the specialized decode. + +Small-batch decode is bound by HBM bandwidth, and each kernel that starts on weights +not yet in L2 stalls on DRAM. Some kernels leave bandwidth idle, such as the attention +core and the all-reduces. At fixed points of each layer, this module launches a small +kernel (``cutedsl/common/prefetch.py``) on a side stream. It pulls the byte ranges of +the weights the next kernels read into L2, so that they start warm. + +Prefetch points (``fork(point, layer, batch)``), each with its batches per TP: + +- ``oproj_down``: after the o-projection (B3-8), the MoE front's packed + router + latent-down rows. +- ``attention_moe``: at TP8 only, after the KDA input projection (B5-8) or the MLA + core (B8), the same front rows plus the shared expert's weights. +- ``next_attention``: once the routed experts are launched, the next layer's + attention weights, capped in MiB per batch (``NEXT_ATTENTION``). Not above B8, + where the HBM-bound experts lose more bandwidth than the attention gains. + +The front and shared-expert ranges are on the critical path. They launch with the +highest stream priority and the maximum shared-memory carveout, which needs a driver +launch because CuTe drops both attributes. + +Each fork records an event on the main stream that the side stream waits on, so a +prefetch starts only after the work it overlaps. ``join()`` at the end of the step +makes the main stream wait for the side stream, so a captured graph always rejoins +it. This lives in ``runtime/`` rather than ``cutedsl/`` because it is host-side +orchestration across layers: a stream, events, and per-(point, layer, batch) tables +built from the model's layers. +""" + +import ctypes +from functools import partial + +import cuda.bindings.driver as cuda +import torch + +from ...ops.cute_dsl.latent_moe_tail.primitives import to_cute +from ..cutedsl.common.host import _checked, resolve_kernel +from ..cutedsl.common.policy import batch_table +from ..cutedsl.common.prefetch import compile_kernel + +# Per-TP batch -> MiB cap (or cap.carveout) of the next layer's attention weights, +# prefetched once the routed experts finish. +# Not at B>8: the HBM-bound experts lose more bandwidth than the attention gains. +NEXT_ATTENTION = {8: "1:96.100,2-3:96,4-7:64,8:32", 16: "1-8:1024"} + + +def attention_weights(attention): + """Next-attention weights in their consumption order.""" + names = ( + *("in_proj_qkvgfab", "fused_qkv_a_g_proj", "fused_qkv_a_proj", "q_proj"), + *("kv_a_proj_with_mqa", "f_b_proj", "q_b_proj", "g_proj"), + ) + weights = [getattr(getattr(attention, n, None), "weight", None) for n in names] + weights += [getattr(attention, "W_UK_T", None), getattr(attention, "W_UV", None)] + weights = [w for w in weights if isinstance(w, torch.Tensor)] + return (*weights, attention.o_proj.weight) + + +class PreparedWeightRanges: + """One owner retains weights, immutable range table, executable and launch ABI.""" + + def __init__(self, weights, *, byte_cap=0, grid=32, carveout=None, priority=False): + self.weights = weights = tuple(weights) + device = weights[0].device + ranges, remaining = [], byte_cap + for weight in weights: + if weight.numel() == 0: + continue + if weight.data_ptr() % 16: # the kernel walks 16-byte aligned chunks + raise ValueError("L2 prefetch ranges must start 16-byte aligned") + span = 1 + sum((n - 1) * s for n, s in zip(weight.shape, weight.stride())) + size = span * weight.element_size() + if byte_cap: + size = min(size, remaining) // 16 * 16 + remaining -= size + if size: + ranges.append((weight.data_ptr(), size)) + if byte_cap and remaining < 16: + break + self.table = torch.tensor(ranges, dtype=torch.int64, device=device) + self.compiled = compile_kernel(len(ranges), grid, carveout, device.index) + self.config = None + if not priority and carveout is None: + self.views = to_cute(self.table, 16) + return + # Launch priority / carveout need a driver launch (CuTe drops both). + self.executor, self.function = resolve_kernel(self.compiled, device.index) + ids, self.attributes = cuda.CUlaunchAttributeID, [] + if priority: # the greatest stream priority + top = cuda.CUlaunchAttribute() + top.id = ids.CU_LAUNCH_ATTRIBUTE_PRIORITY + top.value.priority = _checked(cuda.cuCtxGetStreamPriorityRange())[1] + self.attributes.append(top) + if carveout is not None: + smem = cuda.CUlaunchAttribute() + smem.id = ids.CU_LAUNCH_ATTRIBUTE_PREFERRED_SHARED_MEMORY_CARVEOUT + smem.value.sharedMemCarveout = carveout + self.attributes.append(smem) + config = self.config = cuda.CUlaunchConfig() + config.gridDimX, config.gridDimY, config.gridDimZ = grid, 1, 1 + config.blockDimX, config.blockDimY, config.blockDimZ = 32, 1, 1 + config.sharedMemBytes, config.attrs = 0, self.attributes + config.numAttrs = len(self.attributes) + self.pointer = ctypes.c_void_p(self.table.data_ptr()) + self.arguments = (ctypes.c_void_p * 1)(ctypes.addressof(self.pointer)) + self.address = ctypes.addressof(self.arguments) + + def launch(self, stream): + if self.config is None: + self.compiled(self.views, stream) + else: + self.config.hStream = stream + _checked(cuda.cuLaunchKernelEx(self.config, self.function, self.address, 0)) + + +class PrefetchController: + """Prefetch launches per (point, layer, batch) on one side stream.""" + + def __init__(self, runtime, layers, packed): + self.device, tp = runtime.device, runtime.tp + self.stream = torch.cuda.Stream(device=self.device) + self.events, self.pending, self.tables = {}, False, {} + urgent = partial(PreparedWeightRanges, carveout=100, priority=True) + following = batch_table(NEXT_ATTENTION[tp]) + for position, layer in enumerate(layers): + if not layer.is_moe_layer: + continue + index, moe, attention = layer.layer_idx, layer.mlp, layer.self_attn + # The front GEMM reads [router shard; down shard] from the packed copy. + ranges = urgent((packed[index],), grid=8) + self._add("oproj_down", index, range(3, 9), ranges) + if tp == 8: + shared = moe.shared_experts + shared = shared.gate_up_proj.weight, shared.down_proj.weight + ranges = urgent((packed[index], *shared), grid=32) + batches = range(5, 9) if hasattr(attention, "in_proj_qkvgfab") else (8,) + self._add("attention_moe", index, batches, ranges) + if position + 1 < len(layers): + weights = attention_weights(layers[position + 1].self_attn) + self._add_capped("next_attention", index, weights, following) + + def _add_capped(self, point, layer, weights, table): + groups: dict[tuple, list[int]] = {} + for batch, spec in table.items(): + cap, carveout = spec if isinstance(spec, tuple) else (spec, None) + if cap: + groups.setdefault((cap, carveout), []).append(batch) + for (cap, carveout), batches in groups.items(): + ranges = PreparedWeightRanges( + weights, byte_cap=cap * 1024**2, grid=32, carveout=carveout + ) + self._add(point, layer, batches, ranges) + + def _add(self, point, layer, batches, ranges): + key = point, layer + if key not in self.events: + event = torch.cuda.Event() + event.record(torch.cuda.current_stream(self.device)) + self.events[key] = event + for batch in batches: + self.tables[point, layer, batch] = ranges + + def fork(self, point, layer, batch): + ranges = self.tables.get((point, layer, batch)) + if ranges is None: + return False + event = self.events[point, layer] + event.record(torch.cuda.current_stream(self.device)) + self.stream.wait_event(event) + self.pending = True + ranges.launch(cuda.CUstream(self.stream.cuda_stream)) + return True + + def join(self): + if self.pending: + torch.cuda.current_stream(self.device).wait_stream(self.stream) + self.pending = False diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/providers.py b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/providers.py new file mode 100644 index 0000000000..63e96bfb79 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/providers.py @@ -0,0 +1,172 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Per-layer attention providers (input projection, KDA/MLA core, o-projection) and +the helpers ``SpecializedDecode.build_providers`` uses.""" + +from typing import Any + +import torch + +from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import ( + SkinnyGemmConfig, + shape_dynamic_skinny_gemm, +) +from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first + +from ..cutedsl.gemm.projection import PreparedProjection +from ..cutedsl.kda.prepared import PreparedKda +from ..cutedsl.mla.dcp import DcpBox +from ..cutedsl.mla.prepared import PreparedMla +from ..cutedsl.oproj.prepared import PreparedOproj + +# Decode batches (tokens) served by the prepared providers: every batch up to 8, the +# padded CUDA-graph batches and each further captured size; others run stock. +BATCHES = (*range(1, 9), 16, 32, 64) +PREPARED_MAX = BATCHES[-1] + + +def projection_policies(tp, rows, batches): + """Split-K (M tile, N tile, split, stages) per execution rows' next tested bin.""" + policies = {} + for batch in batches: + n = next(b for b in (8, 16, 32, 64) if batch <= b) + if rows == 6288 and (tp == 8 or n > 8): + tactic = (64, 64, 1, 6) if n == 64 else (128, n, 2, 5 if n == 32 else 6) + else: # TP16 KDA B64: a sixth stage, -0.7 us per layer in the KDA chain + stages = {8: 8, 16: 10, 32: 8, 64: 6 if rows == 3216 else 5}[n] + tactic = (64, n, 4 if n == 8 else 2, stages) + policies[batch] = tactic + return policies + + +class SkinnyProjection: + """B1/B2 skinny GEMM; static K loads two weight tiles before the PDL wait.""" + + def __init__(self, weight, fallback): + self.weight, self.fallback, k = weight, fallback, weight.shape[1] + tactics = {1: (128, 4, 1), 2: (64, 4, 2)} + self.configs = {m: SkinnyGemmConfig(m, *t, 8, k) for m, t in tactics.items()} + for m, cfg in self.configs.items(): # compile before capture + shape_dynamic_skinny_gemm(weight.new_zeros(m, k), weight, cfg, None) + + def run(self, x, out): + cfg = self.configs.get(x.shape[0]) + if cfg is None: + return self.fallback.run(x, out) + return shape_dynamic_skinny_gemm(x, self.weight, cfg, None) + + +def scratch(runtime, name, shape, dtype=torch.bfloat16): + if name not in runtime.buffers: + runtime.buffers[name] = torch.empty(shape, device=runtime.device, dtype=dtype) + return runtime.buffers[name] + + +class Attention: + """Input projection -> KDA/MLA core -> o-projection with its fused AttnRes.""" + + def __init__(self, decoder, runtime, prefetch, mailbox, small_batch, twoshot=None): + attention, tp = decoder.self_attn, runtime.tp + self.kind = kind = "kda" if hasattr(attention, "in_proj_qkvgfab") else "mla" + name = "in_proj_qkvgfab" if kind == "kda" else "fused_qkv_a_g_proj" + weight = getattr(attention, name).weight + pitch = weight.shape[0] + policies = projection_policies(tp, pitch, runtime.batches) + split_k = PreparedProjection(weight, policies).prepare() + self.projection = SkinnyProjection(weight, split_k) + self.core: Any # PreparedKda or PreparedMla + if kind == "kda": + mode = lambda b: (1 if b == 1 else 2 if b <= 8 else 0) if tp == 8 else 0 + policies = {b: (pitch, mode(b)) for b in runtime.batches} + q = runtime.q # verify steps: B requests x q tokens + verify = [b // q for b in runtime.batches if q > 1 and b % q == 0] + policies.update({(b, q): (pitch, int(tp == 8)) for b in verify}) + self.core = PreparedKda(attention, policies, is_conv_state_dim_first()) + else: + heads = attention.num_local_heads + query = scratch(runtime, "mla.query", (64, heads, 576), torch.uint8) + policy = dict(max_context=runtime.max_context, core=runtime.mla_core) + policy.update(queries=runtime.queries) + if runtime.dcp > 1: + policy.update( + dcp=dcp_box(runtime, heads), parity=runtime.mla_layers % 2 + ) + runtime.mla_layers += 1 + pitches = {b: pitch for b in runtime.batches} + self.core = PreparedMla(attention, pitches, query, **policy) + parity = decoder.layer_idx % 2 + front: tuple | None = None + if decoder.is_moe_layer: + front = small_batch.front_views[parity], small_batch.hint_views[parity] + front += (small_batch.batches,) + self.output = PreparedOproj( + decoder, mailbox, runtime.batches, kind == "mla", front, twoshot + ) + projected = scratch(runtime, kind + ".projected", (64, pitch)) + local = scratch(runtime, "attention.local", (64, 12288 // tp)) + self.prefetch, self.layer = prefetch, decoder + self.storage = {b: (projected[:b], local[:b]) for b in range(1, 65)} + + def prepare_layout(self, caches): + if self.kind == "kda": + self.core.prepare_caches(caches) + else: + self.core.prepare_cache(caches[0]) + + def run(self, x, out, prefix, blocks, caches, metadata, slot_mapping, q=1): + b = x.shape[0] + projected, local = self.storage[b] + projected = self.projection.run(x, projected) + if self.kind == "kda": + self.prefetch.fork("attention_moe", self.layer.layer_idx, b) + if q > 1: # uniform q-token verify rows + slots = metadata.spec_state_indices_tensor + accepted = metadata.num_accepted_tokens + self.core.run( + projected, local, caches=caches, slots=slots, accepted=accepted + ) + else: + slots = metadata.non_spec_state_indices_tensor + self.core.run(projected, local, caches=caches, slots=slots) + else: + decode = metadata.decode + tables, lengths = decode.block_table, decode.seq_lens + # Request metadata shorter than b: the core runs the prepared batch b + # with the missing rows as empty requests. + rows = min(b, q * min(lengths.numel(), tables.shape[0])) + if rows < b: + local[rows:].zero_() + slots = slot_mapping[self.layer.self_attn.layer_name] + mla = dict(cache=caches[0], block_table=tables, lengths=lengths) + # DCP: lengths are the local shard's, totals the full sequences'. + mla.update(totals=decode.dcp_tot_seq_lens) + self.core.run(projected, local, slots=slots, q=q, **mla) + self.prefetch.fork("attention_moe", self.layer.layer_idx, b) + self.output.run(local, prefix, blocks, out) + self.prefetch.fork("oproj_down", self.layer.layer_idx, b) + return out + + +def dcp_box(runtime, heads): + """The model's DCP mailboxes (one set for every MLA layer).""" + if runtime.dcp_box is None: + from vllm.distributed import get_dcp_group + + group = get_dcp_group() + assert group.world_size == runtime.dcp == runtime.tp + from vllm.config import get_current_vllm_config + + args = group.device_group, runtime.device, group.world_size, group.rank_in_group + parallel = get_current_vllm_config().parallel_config + interleave = parallel.cp_kv_cache_interleave_size + runtime.dcp_box = DcpBox(*args, heads, interleave=interleave) + return runtime.dcp_box + + +def pack_front(moe, rank, tp): + """[this rank's router rows (padded to 16); this rank's latent down shard].""" + router, down = moe.gate.weight, moe.routed_expert_down_proj.weight + rows, shard = 896 // tp, 3584 // tp + router = router[rank * rows : (rank + 1) * rows] + down = down[rank * shard : (rank + 1) * shard] + return torch.cat((router, router.new_zeros((-rows % 16, 7168)), down)).detach() diff --git a/vllm/models/qwen4_exp/amd/model.py b/vllm/models/qwen4_exp/amd/model.py index cf8357ee6c..e7144693d7 100644 --- a/vllm/models/qwen4_exp/amd/model.py +++ b/vllm/models/qwen4_exp/amd/model.py @@ -37,6 +37,7 @@ from vllm.model_executor.models.interfaces import ( IsHybrid, MixtureOfExperts, MultiModalEmbeddings, + SupportsHCSP, SupportsLoRA, SupportsMRoPE, SupportsPP, @@ -67,6 +68,7 @@ from vllm.model_executor.models.utils import ( maybe_fuse_shared_experts, maybe_prefix, ) +from vllm.models.common.ops.sequence_parallel import sp_all_gather from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.sequence import IntermediateTensors @@ -74,9 +76,11 @@ from vllm.tokenizers.registry import cached_tokenizer_from_config from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.kv_cache_interface import MambaSpec +from ..common.hyperconnection import is_hc_sequence_parallel_enabled from ..common.mtp import make_mtp_hidden_buffer from .hyperconnection import GatedResidual, HyperConnectionConfig from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm +from .ops.sequence_parallel import sp_reduce_scatter, sp_shard from .ple_layer import Qwen4ExpPLELayer from .qsa import Qwen4ExpQSAAttention @@ -162,13 +166,19 @@ _EXTRA_WEIGHTS_MAPPER = WeightsMapper( class Qwen4ExpSparseMoeBlock(Qwen3NextSparseMoeBlock): """Qwen3Next MoE with Qwen4Exp HC validation.""" - def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: - parallel_config = vllm_config.parallel_config - if parallel_config.use_sequence_parallel_moe: - raise NotImplementedError( - "Qwen4Exp HC does not support sequence-parallel MoE" - ) - super().__init__(vllm_config=vllm_config, prefix=prefix) + def __init__( + self, + vllm_config: VllmConfig, + prefix: str = "", + reduce_results: bool = True, + ) -> None: + super().__init__( + vllm_config=vllm_config, prefix=prefix, reduce_results=reduce_results + ) + if self.replicate_shared_expert and not self.is_sequence_parallel: + # Reduce routed outputs before adding replicated shared outputs. + # SP then slices the result to avoid summing shared outputs again via RS. + self.experts.moe_config.skip_final_all_reduce = False config = vllm_config.model_config.hf_text_config self.n_shared_experts = int(config.shared_expert_intermediate_size > 0) @@ -186,11 +196,12 @@ class Qwen4ExpDecoderLayer(nn.Module): self.config = config self.layer_type = layer_type + # MoE SP also enables HC SP, while HC SP can be enabled independently. + self.use_hc_sequence_parallel = is_hc_sequence_parallel_enabled(vllm_config) self.layer_idx = extract_layer_index(prefix) - if vllm_config.parallel_config.use_sequence_parallel_moe: - raise NotImplementedError( - "Qwen4Exp HC does not support sequence-parallel MoE" - ) + self.use_sequence_parallel = ( + vllm_config.parallel_config.use_sequence_parallel_moe + ) self.ple: Qwen4ExpPLELayer | None = None ple_layer_ids = config.ple_layer_ids if (self.layer_idx + 1) in ple_layer_ids: @@ -205,6 +216,7 @@ class Qwen4ExpDecoderLayer(nn.Module): layer_idx=self.layer_idx, ple_dense_layer_id=ple_dense_layer_id, prefix=f"{prefix}.ple", + use_sequence_parallel=self.use_hc_sequence_parallel, ) if layer_type == "linear_attention": @@ -213,6 +225,7 @@ class Qwen4ExpDecoderLayer(nn.Module): vllm_config=vllm_config, prefix=f"{prefix}.linear_attn", gqa_interleaved_layout=False, + reduce_results=not self.use_hc_sequence_parallel, ) elif layer_type in _QSA_LAYER_TYPES: use_qsa = getattr(config, "indexer_n_heads", None) is not None @@ -223,6 +236,7 @@ class Qwen4ExpDecoderLayer(nn.Module): cache_config=vllm_config.cache_config, quant_config=quant_config, prefix=f"{prefix}.self_attn", + reduce_results=not self.use_hc_sequence_parallel, ) else: self.self_attn = Qwen4ExpQSAAttention( @@ -231,12 +245,15 @@ class Qwen4ExpDecoderLayer(nn.Module): layer_id=self.layer_idx, quant_config=quant_config, prefix=f"{prefix}.self_attn", + reduce_results=not self.use_hc_sequence_parallel, ) else: raise ValueError(f"Invalid layer_type {layer_type}") self.mlp = Qwen4ExpSparseMoeBlock( - vllm_config=vllm_config, prefix=f"{prefix}.mlp" + vllm_config=vllm_config, + prefix=f"{prefix}.mlp", + reduce_results=not self.use_hc_sequence_parallel, ) hc_config = HyperConnectionConfig( @@ -294,6 +311,9 @@ class Qwen4ExpDecoderLayer(nn.Module): else: hidden_states, block_input, injection = attn_hc.mix(hidden_states) + if self.use_hc_sequence_parallel: + # Positions retain the full token axis for both text and MRoPE inputs. + block_input = sp_all_gather(block_input)[: positions.shape[-1]] if self.layer_type == "linear_attention": attn_out = self.linear_attn(hidden_states=block_input) elif self.layer_type in _QSA_LAYER_TYPES: @@ -305,10 +325,25 @@ class Qwen4ExpDecoderLayer(nn.Module): raise ValueError("Invalid layer_type") mlp_hc = self.mlp_hyper_connection + if self.use_hc_sequence_parallel: + attn_out = sp_reduce_scatter(attn_out) hidden_states, block_input, injection = mlp_hc.combine_and_mix( hidden_states, attn_out, injection ) - mlp_out = self.mlp(block_input) + if self.use_sequence_parallel: + # MoE consumes local tokens and returns complete local outputs. + mlp_out = self.mlp(block_input, already_sequence_parallel=True) + else: + if self.use_hc_sequence_parallel: + block_input = sp_all_gather(block_input)[: positions.shape[-1]] + mlp_out = self.mlp(block_input) + if self.use_hc_sequence_parallel: + assert isinstance(self.mlp, Qwen4ExpSparseMoeBlock) + reduce_output = self.mlp.experts.moe_config.skip_final_all_reduce + # Only sum partial outputs. Reduced outputs already include all ranks. + mlp_out = ( + sp_reduce_scatter(mlp_out) if reduce_output else sp_shard(mlp_out) + ) return hidden_states, mlp_out, injection @@ -383,6 +418,8 @@ class Qwen4ExpModel(nn.Module): self.num_redundant_experts = ( vllm_config.parallel_config.eplb_config.num_redundant_experts ) + # MoE SP also enables HC SP, while HC SP can be enabled independently. + self.use_hc_sequence_parallel = is_hc_sequence_parallel_enabled(vllm_config) self.vocab_size = config.vocab_size self._qsa_layer_ids = frozenset( layer_idx @@ -457,12 +494,18 @@ class Qwen4ExpModel(nn.Module): if input_ids is None: raise ValueError("input_ids or inputs_embeds is required") hidden_states = self.embed_input_ids(input_ids) + if self.use_hc_sequence_parallel: + # The MoE padding mask stays unsharded: under torch.compile its + # update would land only after the graph has run. + hidden_states = sp_shard(hidden_states) + # Expand only the local token rows when SP is enabled. hidden_states = hidden_states.repeat(1, self.config.hc_count) else: if intermediate_tensors is None: raise ValueError("pipeline stage requires intermediate tensors") hidden_states = intermediate_tensors["hidden_states"] + full_num_tokens = positions.shape[-1] block_output = None injection = None last_layer = None @@ -485,6 +528,8 @@ class Qwen4ExpModel(nn.Module): deepstack_embed = deepstack_input_embeds[ f"deepstack_input_embeds_{layer_idx}" ] + if self.use_hc_sequence_parallel: + deepstack_embed = sp_shard(deepstack_embed) deepstack_embed = ( deepstack_embed.unsqueeze(-2) .expand( @@ -519,6 +564,24 @@ class Qwen4ExpModel(nn.Module): multi_hidden, sample_hidden_states, _ = final_mixer.combine_and_mix( hidden_states, block_output, injection ) + if self.use_hc_sequence_parallel: + if self._mtp_hidden_buffer is not None: + # Gather LM-head and MTP states together when both are needed. + hidden_size = self.config.hidden_size + packed_hidden_states = torch.cat( + [sample_hidden_states, multi_hidden], dim=-1 + ) + packed_hidden_states = sp_all_gather(packed_hidden_states)[ + :full_num_tokens + ] + sample_hidden_states, multi_hidden = packed_hidden_states.split( + [hidden_size, self.config.hc_count * hidden_size], dim=-1 + ) + sample_hidden_states = sample_hidden_states.contiguous() + else: + sample_hidden_states = sp_all_gather(sample_hidden_states)[ + :full_num_tokens + ] if self._mtp_hidden_buffer is not None: # Capture the pre-final-mixer multi-stream hidden state # [T, hc_count*H] for the MTP drafter (zero extra compute: @@ -574,6 +637,7 @@ class Qwen4ExpModel(nn.Module): class Qwen4ExpForCausalLM( nn.Module, HasInnerState, + SupportsHCSP, SupportsLoRA, SupportsMRoPE, SupportsPP, @@ -806,6 +870,7 @@ class Qwen4ExpProcessingInfo(Qwen3VLProcessingInfo): class Qwen4ExpForConditionalGeneration( Qwen3_5ForConditionalGeneration, HasInnerState, + SupportsHCSP, Qwen4ExpMixtureOfExperts, ): """Qwen3-VL vision tower backed by the Qwen4Exp language model.""" diff --git a/vllm/models/qwen4_exp/amd/mtp.py b/vllm/models/qwen4_exp/amd/mtp.py index 98ee4318fb..73c9833335 100644 --- a/vllm/models/qwen4_exp/amd/mtp.py +++ b/vllm/models/qwen4_exp/amd/mtp.py @@ -34,7 +34,11 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) from vllm.model_executor.model_loader.utils import configure_quant_config -from vllm.model_executor.models.interfaces import LocalArgmaxMixin, SupportsPP +from vllm.model_executor.models.interfaces import ( + LocalArgmaxMixin, + SupportsHCSP, + SupportsPP, +) from vllm.model_executor.models.qwen3_5 import Qwen3_5Model from vllm.model_executor.models.utils import ( AutoWeightsLoader, @@ -45,8 +49,10 @@ from vllm.model_executor.models.utils import ( maybe_fuse_shared_experts, maybe_prefix, ) +from vllm.models.common.ops.sequence_parallel import sp_all_gather, sp_shard from vllm.sequence import IntermediateTensors +from ..common.hyperconnection import is_hc_sequence_parallel_enabled from .hyperconnection import GatedResidual, HyperConnectionConfig from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm from .model import ( @@ -181,6 +187,8 @@ class Qwen4ExpMultiTokenPredictor(nn.Module): config: Qwen4ExpTextConfig = model_config.hf_text_config self.config = config + # MoE SP also enables HC SP, while HC SP can be enabled independently. + self.use_hc_sequence_parallel = is_hc_sequence_parallel_enabled(vllm_config) self.vocab_size = config.vocab_size self.mtp_start_layer_idx = config.num_hidden_layers @@ -321,6 +329,9 @@ class Qwen4ExpMultiTokenPredictor(nn.Module): hidden_states = intermediate_tensors["hidden_states"] current_step_idx = spec_step_idx % self.num_mtp_layers + full_num_tokens = positions.shape[-1] + if self.use_hc_sequence_parallel: + hidden_states = sp_shard(hidden_states) layer = self.layers[current_step_idx] hidden_states, block_output, injection = layer( hidden_states=hidden_states, @@ -348,6 +359,17 @@ class Qwen4ExpMultiTokenPredictor(nn.Module): hidden_states, block_output, injection ) ) + if self.use_hc_sequence_parallel: + # Gather both outputs together, then restore their contiguous layouts. + packed_hidden_states = torch.cat( + [sample_hidden_states, multi_hidden], dim=-1 + ) + packed_hidden_states = sp_all_gather(packed_hidden_states)[:full_num_tokens] + sample_hidden_states, multi_hidden = packed_hidden_states.split( + [hidden_size, hc_count * hidden_size], dim=-1 + ) + sample_hidden_states = sample_hidden_states.contiguous() + multi_hidden = multi_hidden.contiguous() return sample_hidden_states, multi_hidden def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: @@ -377,7 +399,13 @@ class Qwen4ExpMultiTokenPredictor(nn.Module): "hidden_states": 0, } ) -class Qwen4ExpMTP(LocalArgmaxMixin, nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts): +class Qwen4ExpMTP( + LocalArgmaxMixin, + nn.Module, + SupportsHCSP, + SupportsPP, + Qwen4ExpMixtureOfExperts, +): packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], "gate_up_proj": ["gate_proj", "up_proj"], diff --git a/vllm/models/qwen4_exp/amd/ops/sequence_parallel.py b/vllm/models/qwen4_exp/amd/ops/sequence_parallel.py new file mode 100644 index 0000000000..2efcd56d54 --- /dev/null +++ b/vllm/models/qwen4_exp/amd/ops/sequence_parallel.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.models.common.ops import sequence_parallel as common_sp +from vllm.utils.math_utils import cdiv +from vllm.utils.torch_utils import direct_register_custom_op + + +def _sp_shard_op(x: torch.Tensor) -> torch.Tensor: + output = common_sp.sp_shard(x) + # A functional custom op must not return a view of its input. + unpadded = x.shape[0] % get_tensor_model_parallel_world_size() == 0 + return output.clone() if unpadded else output + + +def _sp_local_fake(x: torch.Tensor) -> torch.Tensor: + tp_size = get_tensor_model_parallel_world_size() + return x.new_empty((cdiv(x.shape[0], tp_size), *x.shape[1:])) + + +# Dynamo would specialize the padding branch on the traced token count, and +# vLLM drops the resulting guards. Keep the branch opaque when compiling. +direct_register_custom_op( + op_name="sp_reduce_scatter", + op_func=common_sp.sp_reduce_scatter, + fake_impl=_sp_local_fake, +) +direct_register_custom_op( + op_name="sp_shard", + op_func=_sp_shard_op, + fake_impl=_sp_local_fake, +) + + +def sp_reduce_scatter(x: torch.Tensor) -> torch.Tensor: + if torch.compiler.is_compiling(): + return torch.ops.vllm.sp_reduce_scatter(x) + return common_sp.sp_reduce_scatter(x) + + +def sp_shard(x: torch.Tensor) -> torch.Tensor: + if torch.compiler.is_compiling(): + return torch.ops.vllm.sp_shard(x) + return common_sp.sp_shard(x) diff --git a/vllm/models/qwen4_exp/amd/ple_layer.py b/vllm/models/qwen4_exp/amd/ple_layer.py index 30595df32d..c1cfc0f246 100644 --- a/vllm/models/qwen4_exp/amd/ple_layer.py +++ b/vllm/models/qwen4_exp/amd/ple_layer.py @@ -9,6 +9,7 @@ from torch import nn from transformers import Qwen4ExpTextConfig from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config +from vllm.distributed import get_etp_group, get_tp_group from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.layers.linear import MergedColumnParallelLinear @@ -21,7 +22,13 @@ from vllm.model_executor.layers.mamba.mamba_utils import ( from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, ) +from vllm.model_executor.layers.quantization.utils.fp8_utils import is_fp8 from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm.models.common.ops.sequence_parallel import ( + sp_all_gather, + sp_reduce_scatter, + sp_shard, +) from vllm.utils.torch_utils import direct_register_custom_op from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.attention.backends.short_conv_attn import ( @@ -182,8 +189,10 @@ class Qwen4ExpNGramEmbedding(nn.Module): data_parallel_rank: int = 0, quant_config: QuantizationConfig | None = None, params_dtype: torch.dtype | None = None, + use_sequence_parallel: bool = False, ) -> None: super().__init__() + self.use_sequence_parallel = use_sequence_parallel self.embedding_dim = embedding_dim self.layer_name = layer_name self.ngram_size = int(config.ngram_size) @@ -253,6 +262,11 @@ class Qwen4ExpNGramEmbedding(nn.Module): if engram_config is not None and engram_config.cpu_offload else Qwen4ExpPLEDeviceEmbedding ) + # A TP reduce-scatter covers all vocabulary owners only when ETP=TP. + self.use_reduce_scatter = ( + self.use_sequence_parallel + and get_etp_group().world_size == get_tp_group().world_size + ) self.ngram_embedding = embedding_cls( padded_vocab_size, self.head_dim, @@ -263,6 +277,7 @@ class Qwen4ExpNGramEmbedding(nn.Module): num_ngram_heads=self.ngram_heads, max_total_tokens=max_total_tokens, data_parallel_rank=data_parallel_rank, + reduce_results=not self.use_reduce_scatter, ) logger.info( "Initialized AMD PLE embedding %s: quantization_method=%s, " @@ -274,12 +289,25 @@ class Qwen4ExpNGramEmbedding(nn.Module): self.ngram_embedding.weight.is_pinned(), ) + def _sp_shard_embeddings(self, embeddings: torch.Tensor) -> torch.Tensor: + """Shard embedding token rows, optionally reducing vocabulary contributions.""" + # Each embedding element has one vocabulary owner, including FP8 bytes. + output_dtype = embeddings.dtype + if is_fp8(embeddings): + embeddings = embeddings.view(torch.int8) + if self.use_reduce_scatter: + embeddings = sp_reduce_scatter(embeddings) + else: + embeddings = sp_shard(embeddings) + return embeddings.view(output_dtype) + def forward( self, input_ids: torch.Tensor, query_start_loc: torch.Tensor, ngram_context: torch.Tensor, ) -> torch.Tensor: + """Return full embeddings or a token shard after combining vocabulary owners.""" input_ids = input_ids.reshape(-1) ngram_ids = input_ids.new_empty( (input_ids.shape[0], self.ngram_heads), dtype=torch.int64 @@ -302,17 +330,19 @@ class Qwen4ExpNGramEmbedding(nn.Module): output, self.layer_name, ) + else: + output = ngram_ids.new_empty( + (ngram_ids.shape[0], self.embedding_dim), + dtype=embedding.params_dtype, + ) + torch.ops.vllm.qwen4_exp_amd_ple_ngram_embedding( + ngram_ids, + output, + self.layer_name, + ) + if not self.use_sequence_parallel: return output - output = ngram_ids.new_empty( - (ngram_ids.shape[0], self.embedding_dim), - dtype=embedding.params_dtype, - ) - torch.ops.vllm.qwen4_exp_amd_ple_ngram_embedding( - ngram_ids, - output, - self.layer_name, - ) - return output + return self._sp_shard_embeddings(output) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load hash buffers and checkpoint-split embedding rows.""" @@ -388,8 +418,10 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): layer_idx: int = 0, ple_dense_layer_id: int | None = None, prefix: str = "", + use_sequence_parallel: bool = False, ) -> None: super().__init__() + self.use_sequence_parallel = use_sequence_parallel model_config = vllm_config.model_config cache_config = vllm_config.cache_config quant_config = vllm_config.quant_config @@ -420,8 +452,9 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): data_parallel_rank=vllm_config.parallel_config.data_parallel_rank, quant_config=quant_config, params_dtype=model_config.dtype, + use_sequence_parallel=self.use_sequence_parallel, ) - # The PLE cache is TP-replicated, so this merged projection is too. + # Projection weights stay replicated even when their token rows are sharded. self.kv_proj = MergedColumnParallelLinear( int(config.ple_embed_dim), [self.hc_hidden_size, self.hidden_size], @@ -486,7 +519,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): self, inputs: torch.Tensor, residual: torch.Tensor, - outer_residual: torch.Tensor, + outer_residual: torch.Tensor | None, metadata: PleShortConvAttentionMetadata, conv_state: torch.Tensor, conv_weights: torch.Tensor, @@ -501,7 +534,8 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): has_non_spec = has_prefill or has_decode inputs = inputs[: metadata.num_actual_tokens] residual = residual[: metadata.num_actual_tokens] - outer_residual = outer_residual[: metadata.num_actual_tokens] + if outer_residual is not None: + outer_residual = outer_residual[: metadata.num_actual_tokens] spec_token_indices = None non_spec_token_indices = None @@ -553,11 +587,14 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): residual_d, residual_p = torch.split( residual, [num_decode_tokens, num_prefill_tokens], dim=0 ) - outer_residual_d, outer_residual_p = torch.split( - outer_residual, - [num_decode_tokens, num_prefill_tokens], - dim=0, - ) + outer_residual_d: torch.Tensor | None = None + outer_residual_p: torch.Tensor | None = None + if outer_residual is not None: + outer_residual_d, outer_residual_p = torch.split( + outer_residual, + [num_decode_tokens, num_prefill_tokens], + dim=0, + ) token_indices_d = None token_indices_p = None else: @@ -628,14 +665,16 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): self, inputs: torch.Tensor, residual: torch.Tensor, - outer_residual: torch.Tensor, + outer_residual: torch.Tensor | None, ) -> None: + """Update convolution state, optionally fusing both residual additions.""" forward_context = get_forward_context() attn_metadata = forward_context.attn_metadata # Profiling omits all metadata or this Mamba entry. Short convolution # is a no-op there, but preserve the outer residual addition. if attn_metadata is None: - residual.add_(outer_residual) + if outer_residual is not None: + residual.add_(outer_residual) return if not isinstance(attn_metadata, dict): @@ -646,7 +685,8 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): layer_attn_metadata = attn_metadata.get(self.prefix) if layer_attn_metadata is None: - residual.add_(outer_residual) + if outer_residual is not None: + residual.add_(outer_residual) return if not isinstance(layer_attn_metadata, PleShortConvAttentionMetadata): raise TypeError( @@ -687,7 +727,10 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): ngram_context: torch.Tensor, ) -> torch.Tensor: input_ids = input_ids.reshape(-1) - if input_ids.shape[0] != hidden_states.shape[0]: + if ( + not self.use_sequence_parallel + and input_ids.shape[0] != hidden_states.shape[0] + ): raise ValueError( "PLE expects input_ids and hidden_states to have the same " f"token length, got {input_ids.shape[0]} and " @@ -711,12 +754,23 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): # The short-conv op is a piecewise-graph splitting op: state routing # reads runtime request metadata. It accumulates the convolution and # the outer residual into gated_output. - torch.ops.vllm.qwen4_exp_ple_short_conv( - conv_input, - gated_output, - hidden_states, - self.prefix, - ) + if self.use_sequence_parallel: + # Only convolution needs full tokens. Add residuals on local shards. + conv_input = sp_all_gather(conv_input)[: input_ids.shape[0]] + # Metadata may omit padding tokens or all tokens during profiling. + conv_output = torch.zeros_like(conv_input) + torch.ops.vllm.qwen4_exp_ple_short_conv( + conv_input, conv_output, None, self.prefix + ) + gated_output.add_(sp_shard(conv_output)) + gated_output.add_(hidden_states) + else: + torch.ops.vllm.qwen4_exp_ple_short_conv( + conv_input, + gated_output, + hidden_states, + self.prefix, + ) return gated_output @@ -787,7 +841,7 @@ def qwen4_exp_amd_ple_ngram_embedding_pinned( def qwen4_exp_ple_short_conv( inputs: torch.Tensor, residual: torch.Tensor, - outer_residual: torch.Tensor, + outer_residual: torch.Tensor | None, layer_name: str, ) -> None: layer = get_forward_context().no_compile_layers[layer_name] diff --git a/vllm/models/qwen4_exp/common/hyperconnection.py b/vllm/models/qwen4_exp/common/hyperconnection.py index 5f8bde0479..42c3fba4ff 100644 --- a/vllm/models/qwen4_exp/common/hyperconnection.py +++ b/vllm/models/qwen4_exp/common/hyperconnection.py @@ -35,6 +35,16 @@ import torch import torch.nn.functional as F from torch import nn +from vllm.config import VllmConfig + + +def is_hc_sequence_parallel_enabled(vllm_config: VllmConfig) -> bool: + """Enable GR/PLE SP explicitly or automatically with sequence-parallel MoE.""" + parallel_config = vllm_config.parallel_config + return parallel_config.tensor_parallel_size > 1 and ( + parallel_config.enable_hc_sp or parallel_config.use_sequence_parallel_moe + ) + # --------------------------------------------------------------------------- # Config diff --git a/vllm/models/qwen4_exp/common/mtp.py b/vllm/models/qwen4_exp/common/mtp.py index 457e8bc064..38e6c2540f 100644 --- a/vllm/models/qwen4_exp/common/mtp.py +++ b/vllm/models/qwen4_exp/common/mtp.py @@ -3,9 +3,9 @@ """Shared MTP drafter state for Qwen4Exp.""" import torch +from transformers import Qwen4ExpTextConfig from vllm.config import VllmConfig -from vllm.transformers_utils.configs.qwen4_exp import Qwen4ExpTextConfig def make_mtp_hidden_buffer( diff --git a/vllm/models/qwen4_exp/common/ngram_embedding.py b/vllm/models/qwen4_exp/common/ngram_embedding.py index d4a68d1338..0d875d015c 100644 --- a/vllm/models/qwen4_exp/common/ngram_embedding.py +++ b/vllm/models/qwen4_exp/common/ngram_embedding.py @@ -719,7 +719,11 @@ class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding): slot_size, slot_offset = self._get_dp_gather_slot(ngram_ids.shape[0]) gathered_ids = self._gather_dp_ids(ngram_ids, slot_size) embeddings = self._lookup(gathered_ids) - embeddings = self._reduce_etp_embeddings(embeddings) + if self.reduce_results: + embeddings = self._reduce_etp_embeddings(embeddings) + else: + # The outer TP reduce-scatter requires ETP=TP. + assert self.etp_data_parallel_size == 1 return self._select_embeddings( embeddings, ngram_ids.shape[0], diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index b9ad113e98..408ed07a38 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -84,6 +84,7 @@ from vllm.tokenizers.registry import cached_tokenizer_from_config from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.kv_cache_interface import MambaSpec +from ..common.hyperconnection import is_hc_sequence_parallel_enabled from ..common.mtp import make_mtp_hidden_buffer from .hyperconnection import GatedResidual, HyperConnectionConfig from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm @@ -96,14 +97,6 @@ from .qsa import Qwen4ExpQSAAttention _QSA_LAYER_TYPES = ("qwen_sparse_attention", "indexed_attention") -def is_hc_sequence_parallel_enabled(vllm_config: VllmConfig) -> bool: - """Enable GR/PLE SP explicitly or automatically with sequence-parallel MoE.""" - parallel_config = vllm_config.parallel_config - return parallel_config.tensor_parallel_size > 1 and ( - parallel_config.enable_hc_sp or parallel_config.use_sequence_parallel_moe - ) - - def without_modelopt_fp4( quant_config: QuantizationConfig | None, ) -> QuantizationConfig | None: diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py index df82e2d1da..659ae74bff 100644 --- a/vllm/models/qwen4_exp/nvidia/mtp.py +++ b/vllm/models/qwen4_exp/nvidia/mtp.py @@ -58,6 +58,7 @@ from vllm.models.common.ops.sequence_parallel import ( ) from vllm.sequence import IntermediateTensors +from ..common.hyperconnection import is_hc_sequence_parallel_enabled from .hyperconnection import GatedResidual, HyperConnectionConfig from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm from .model import ( @@ -66,7 +67,6 @@ from .model import ( Qwen4ExpDecoderLayer, Qwen4ExpMixtureOfExperts, Qwen4ExpSparseMoeBlock, - is_hc_sequence_parallel_enabled, ) diff --git a/vllm/parser/engine/adapters.py b/vllm/parser/engine/adapters.py index cd7aaa7297..a4a307a43a 100644 --- a/vllm/parser/engine/adapters.py +++ b/vllm/parser/engine/adapters.py @@ -115,6 +115,10 @@ class ParserEngineReasoningAdapter(ReasoningParser): def reasoning_end_token_ids(self) -> frozenset[int]: return self._parser_engine.reasoning_end_token_ids + @property + def implicit_reasoning_end_token_ids(self) -> frozenset[int]: + return self._parser_engine.implicit_reasoning_end_token_ids + def adjust_request( self, request: ChatCompletionRequest | ResponsesRequest, diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index 95df03eab5..bbb3605ed7 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -678,6 +678,11 @@ class ParserEngine(Parser): def reasoning_end_token_ids(self) -> frozenset[int]: return self._reasoning_end_token_ids + @property + def implicit_reasoning_end_token_ids(self) -> frozenset[int]: + """Reasoning-end tokens that are also content, e.g. a tool-call start.""" + return self._reasoning_end_token_ids - {self._reasoning_end_token_id} + def find_reasoning_end(self, token_ids: Sequence[int]) -> ReasoningEnd | None: """First reasoning end in `token_ids`, and if its token is content.""" end_ids = self._reasoning_end_token_ids diff --git a/vllm/renderers/embed_utils.py b/vllm/renderers/embed_utils.py index 981aeee3fb..69ff920fa8 100644 --- a/vllm/renderers/embed_utils.py +++ b/vllm/renderers/embed_utils.py @@ -91,11 +91,17 @@ def safe_load_prompt_embeds( parameter="prompt_embeds", ) - # Pin each tensor to the model's hidden_size. Validating here + # Pin each tensor to the model's text embedding width. Validating here # also transitively guarantees cross-tensor consistency for requests that # include multiple `prompt_embeds` parts, which is required by downstream # concatenation in `_build_mixed_prompt_embeds`. - expected_hidden_size = model_config.get_hidden_size() + # OPT projects token embeddings to hidden_size inside its decoder. CLIP + # and SigLIP still accept text embeddings at hidden_size, not projection_dim. + expected_hidden_size = getattr( + model_config.hf_text_config, + "word_embed_proj_dim", + model_config.get_hidden_size(), + ) if tensor.shape[1] != expected_hidden_size: raise VLLMValidationError( f"`prompt_embeds` hidden_size {tensor.shape[1]} does not match " diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py index 09ea33857c..0c9412816a 100644 --- a/vllm/renderers/online_derenderer.py +++ b/vllm/renderers/online_derenderer.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Sequence +from functools import cached_property from typing import Any from vllm.config import ModelConfig @@ -8,7 +9,11 @@ from vllm.entrypoints.generate.base.protocol import ( DeltaMessage, ToolCall, ) -from vllm.entrypoints.generate.base.serving import decode_token_ids +from vllm.entrypoints.generate.base.serving import ( + chat_top_logprobs_limit, + completion_top_logprobs_limit, + decode_token_ids, +) from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionLogProb, ChatCompletionLogProbs, @@ -130,6 +135,43 @@ class OnlineDerenderer: generate_response, chat_request, prompt_token_ids ) + @cached_property + def _model_eos_token_ids(self) -> frozenset[int]: + """EOS ids the engine stops on: the tokenizer's and generation_config's.""" + eos_ids: set[int] = set() + if (eos_token_id := self.renderer.get_eos_token_id()) is not None: + eos_ids.add(eos_token_id) + config_eos = self.model_config.try_get_generation_config().get("eos_token_id") + if config_eos is not None: + eos_ids.update([config_eos] if isinstance(config_eos, int) else config_eos) + return frozenset(eos_ids) + + def _text_token_ids( + self, + token_ids: list[int], + finish_reason: str | None, + stop_reason: int | str | None, + request: ChatCompletionRequest | None, + ) -> list[int]: + """The ids whose text the engine returns. Like its detokenizer, leave + out the token that stopped generation (EOS when `stop_reason` is None, + a `stop_token_ids` entry when it is an int) unless + `include_stop_str_in_output` is set. The last-id check guards against + `finish_reason` defaulting to "stop" when the client sends neither.""" + if ( + finish_reason != "stop" + or isinstance(stop_reason, str) + or not token_ids + or (request is not None and request.include_stop_str_in_output) + ): + return token_ids + last = token_ids[-1] + if stop_reason is None: + stopped = last in self._model_eos_token_ids + else: + stopped = last == stop_reason + return token_ids[:-1] if stopped else token_ids + def _derender_chat( self, generate_response: GenerateTokensResponse, @@ -166,12 +208,17 @@ class OnlineDerenderer: and chat_request.return_token_ids and (not has_parser or chat_request.include_reasoning) ) + top_limit = _chat_top_logprobs_limit(chat_request) for choice in generate_response.choices: if not choice.token_ids: raise ValueError(f"choice {choice.index} has empty or null token_ids") resolved_logprobs = ( - _resolve_logprobs(choice.logprobs, tokenizer) + _resolve_logprobs( + choice.logprobs, + tokenizer, + top_limit=top_limit, + ) if choice.logprobs is not None else None ) @@ -182,7 +229,12 @@ class OnlineDerenderer: # (default True when no request was given). decoded_text, _ = self._detokenize_delta( tokenizer, - choice.token_ids, + self._text_token_ids( + choice.token_ids, + choice.finish_reason, + choice.stop_reason, + chat_request, + ), seed_state, skip_special_tokens=skip_special, spaces_between_special_tokens=spaces_between, @@ -443,11 +495,14 @@ class OnlineDerenderer: updated_state = state return_token_ids = chat_request is not None and chat_request.return_token_ids + top_limit = _chat_top_logprobs_limit(chat_request) for choice in generate_chunk.choices: delta_tids = choice.token_ids or [] new_text, updated_state = await self._detokenize_delta_async( tokenizer, - delta_tids, + self._text_token_ids( + delta_tids, choice.finish_reason, choice.stop_reason, chat_request + ), updated_state, skip_special_tokens=skip_special, spaces_between_special_tokens=spaces_between, @@ -466,6 +521,7 @@ class OnlineDerenderer: choice.logprobs, tokenizer, initial_context_token_ids=state.logprob_context_token_ids, + top_limit=top_limit, ) include_role = not updated_state.role_sent @@ -617,7 +673,12 @@ class OnlineDerenderer: # in this method's docstring. text, detok_state = self._detokenize_delta( tokenizer, - delta_tids, + self._text_token_ids( + delta_tids, + choice.finish_reason, + choice.stop_reason, + chat_request, + ), detok_state, skip_special_tokens=skip_special, spaces_between_special_tokens=spaces_between, @@ -744,6 +805,7 @@ class OnlineDerenderer: total_prompt_tokens = 0 total_completion_tokens = 0 index = 0 + top_limit = _completion_top_logprobs_limit(completion_request) for gen, pt, seed_ids_override in zip( generate_responses, prompt_tokens_list, prompt_token_ids_list @@ -772,9 +834,13 @@ class OnlineDerenderer: ) completion_logprobs = None if choice.logprobs is not None: - resolved = _resolve_logprobs(choice.logprobs, tokenizer) + resolved = _resolve_logprobs( + choice.logprobs, + tokenizer, + top_limit=top_limit, + ) completion_logprobs = _convert_chat_logprobs_to_completion_logprobs( - resolved + resolved, choice.logprobs ) choices.append( CompletionResponseChoice( @@ -847,6 +913,7 @@ class OnlineDerenderer: stream_choices: list[CompletionResponseStreamChoice] = [] updated_state = state + top_limit = _completion_top_logprobs_limit(completion_request) for choice in generate_chunk.choices: delta_tids = choice.token_ids or [] new_text, updated_state = await self._detokenize_delta_async( @@ -863,9 +930,12 @@ class OnlineDerenderer: choice.logprobs, tokenizer, initial_context_token_ids=state.logprob_context_token_ids, + top_limit=top_limit, ) completion_logprobs = _convert_chat_logprobs_to_completion_logprobs( - resolved, initial_text_offset=state.logprob_text_offset + resolved, + choice.logprobs, + initial_text_offset=state.logprob_text_offset, ) updated_state = updated_state.model_copy( @@ -1017,10 +1087,38 @@ def _correct_decoded_token( return "" +def _chat_top_logprobs_limit(chat_request: ChatCompletionRequest | None) -> int | None: + """The cut `/v1/chat/completions` applies (`chat_top_logprobs_limit`, the + same helper its `_get_top_logprobs` uses). None keeps every candidate, which + is also what happens without a `chat_request` or with one that does not ask + for logprobs: it then says nothing about the cut, so the candidates + generate returned are passed through.""" + if chat_request is None or not chat_request.logprobs: + return None + return chat_top_logprobs_limit( + chat_request.top_logprobs, bool(chat_request.logprob_token_ids) + ) + + +def _completion_top_logprobs_limit( + completion_request: CompletionRequest | None, +) -> int | None: + """The cut `/v1/completions` applies (`completion_top_logprobs_limit`, the + same helper `_create_completion_logprobs` uses). None keeps every candidate, + which is also what happens without a `completion_request` or with one that + does not ask for logprobs (`logprobs` unset).""" + if completion_request is None: + return None + return completion_top_logprobs_limit( + completion_request.logprobs, bool(completion_request.logprob_token_ids) + ) + + def _resolve_logprobs( logprobs: GenerateLogProbs, tokenizer: TokenizerLike, initial_context_token_ids: Sequence[int] = (), + top_limit: int | None = None, ) -> ChatCompletionLogProbs: """Convert generate's integer-id logprobs to the OpenAI chat shape. @@ -1030,6 +1128,11 @@ def _resolve_logprobs( ids as context). ``initial_context_token_ids`` seeds that context with sampled IDs from preceding chunks (streaming), so multi-byte characters split across chunk boundaries still resolve. + + Generate returns every candidate the engine produced for a position (the + sampled token first, then the top k). ``top_limit`` applies the target + endpoint's cut before decoding, so each derender endpoint returns what its + coupled counterpart would; None keeps all of them. """ if logprobs.content is None: return ChatCompletionLogProbs() @@ -1038,9 +1141,12 @@ def _resolve_logprobs( resolved_content = [] for entry in logprobs.content: - # One batch per position: the sampled id and its top-k ids. + top_entries = ( + entry.top_logprobs if top_limit is None else entry.top_logprobs[:top_limit] + ) + # One batch per position: the sampled id and its kept top-k ids. (token_str, token_bytes), *top_decoded = decode_token_ids( - [entry.token_id, *(top.token_id for top in entry.top_logprobs)], + [entry.token_id, *(top.token_id for top in top_entries)], tokenizer, ) @@ -1051,7 +1157,7 @@ def _resolve_logprobs( token_bytes = list(token_str.encode("utf-8")) resolved_top = [] - for top, (top_str, top_bytes) in zip(entry.top_logprobs, top_decoded): + for top, (top_str, top_bytes) in zip(top_entries, top_decoded): if top_str.endswith("\ufffd"): top_str = _correct_decoded_token( top.token_id, context_token_ids, tokenizer @@ -1081,11 +1187,17 @@ def _resolve_logprobs( def _convert_chat_logprobs_to_completion_logprobs( logprobs: ChatCompletionLogProbs, + generate_logprobs: GenerateLogProbs, initial_text_offset: int = 0, ) -> CompletionLogProbs: """Convert ChatCompletionLogProbs (per-token objects) to CompletionLogProbs (parallel flat lists) as required by the /v1/completions response schema. + ``generate_logprobs`` is what ``logprobs`` was resolved from. A position + with no candidates there had no logprobs in the engine output, which + /v1/completions reports as ``None``; every other position is a dict of + the candidates the cut kept. + ``initial_text_offset`` keeps ``text_offset`` absolute across streaming chunks, mirroring the generate streaming path. """ @@ -1098,13 +1210,13 @@ def _convert_chat_logprobs_to_completion_logprobs( text_offset: list[int] = [] offset = initial_text_offset - for entry in logprobs.content: + for entry, source in zip(logprobs.content, generate_logprobs.content or ()): text_offset.append(offset) tokens.append(entry.token) token_logprobs.append(entry.logprob) top_logprobs_list.append( {t.token: t.logprob for t in entry.top_logprobs} - if entry.top_logprobs + if source.top_logprobs else None ) offset += len(entry.token) diff --git a/vllm/tokenizers/registry.py b/vllm/tokenizers/registry.py index 972300cad1..a52bf18401 100644 --- a/vllm/tokenizers/registry.py +++ b/vllm/tokenizers/registry.py @@ -174,6 +174,10 @@ def resolve_tokenizer_args( token=kwargs.get("token"), ) + if tokenizer_mode == "mistral": + # MistralCommonBackend rejects HF-only keyword arguments, even None. + kwargs.pop("code_revision", None) + return tokenizer_mode, tokenizer_name, args, kwargs @@ -181,6 +185,10 @@ cached_resolve_tokenizer_args = lru_cache(resolve_tokenizer_args) def tokenizer_args_from_config(config: "ModelConfig", **kwargs): + if config.tokenizer == config.model: + code_revision = getattr(config, "code_revision", None) + if code_revision is not None: + kwargs.setdefault("code_revision", code_revision) return cached_resolve_tokenizer_args( config.tokenizer, runner_type=config.runner_type, @@ -241,6 +249,7 @@ def get_tokenizer( tokenizer_name, trust_remote_code=trust_remote_code, revision=revision, + code_revision=kwargs.get("code_revision"), config_format=config_format, token=kwargs.get("token"), ) @@ -288,6 +297,10 @@ def cached_tokenizer_from_config(model_config: "ModelConfig", **kwargs): return None _maybe_register_hf_config(getattr(model_config, "hf_config", None)) + if model_config.tokenizer == model_config.model: + code_revision = getattr(model_config, "code_revision", None) + if code_revision is not None: + kwargs.setdefault("code_revision", code_revision) return cached_get_tokenizer( model_config.tokenizer, diff --git a/vllm/tool_parsers/apertus_tool_parser.py b/vllm/tool_parsers/apertus_tool_parser.py index af9654c0cb..35e65ccb7b 100644 --- a/vllm/tool_parsers/apertus_tool_parser.py +++ b/vllm/tool_parsers/apertus_tool_parser.py @@ -417,7 +417,9 @@ class ApertusToolParser(ToolParser): """ try: - parsed, _ = partial_json_loads(json_str, Allow.ALL) + # An exponent can change a number's already-streamed representation. + flags = Allow.ALL if is_final else Allow.ALL & ~Allow.NUM + parsed, _ = partial_json_loads(json_str, flags) if not isinstance(parsed, list): parsed = [parsed] if parsed else [] except Exception: diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 5badccf6b5..b995d19d89 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -948,7 +948,7 @@ def fp8_fp4_paged_sparse_mqa_logits( Args: q: ``(q_values, q_scale)``; q_values is [num_q_tokens, 1, H, D]. kv_cache: Fused paged cache [num_pages, page_kv, 1, D' + 4] uint8, - page stride 512B-aligned. + page stride 256B-aligned. weights: [num_q_tokens, H] ``torch.bfloat16``. metadata: From `get_paged_sparse_mqa_logits_metadata`. num_max_sparse_blocks: Candidate blocks per row. diff --git a/vllm/utils/humming.py b/vllm/utils/humming.py index 8004812ea4..0cd80a9be5 100644 --- a/vllm/utils/humming.py +++ b/vllm/utils/humming.py @@ -109,10 +109,70 @@ def create_humming_locks(device: torch.device | str) -> torch.Tensor: ) +_patched = False + + +def _patch_humming() -> None: + """Make humming's per-call ``DeviceInfo(device)`` lookups traceable. + + The forward path (``LayerConfig.check_device``, PDL heuristics) builds a + ``DeviceInfo`` on every call, which breaks vLLM's fullgraph compile: + Dynamo recurses forever in ``DeviceInfo.__getattr__`` + (vllm-project/humming#147) and, with that fixed (#148), still cannot + trace the extension's ``_DeviceInfo`` constructor. Hand out one cached + instance per device instead; ``prepare_layer_config`` creates it eagerly + during weight processing. + + TODO: remove after bumping humming to a release with + vllm-project/humming#149. + """ + global _patched + if _patched: + return + _patched = True + + import humming.config.config + import humming.transform + import humming.tune + import humming.tune.process_input + import humming.utils.smem + from humming.device import DeviceInfo, get_device_index + + def __getattr__(self, name: str) -> Any: + if name.startswith("_"): + raise AttributeError( + f"{type(self).__name__!r} object has no attribute {name!r}" + ) + return getattr(self._get_value(), name) + + DeviceInfo.__getattr__ = __getattr__ + + device_infos: dict[int, DeviceInfo] = {} + + def cached_device_info(index: int | torch.device | None = None) -> DeviceInfo: + if isinstance(index, torch.device) and index.type != "cuda": + return DeviceInfo(index) + device_index = get_device_index(index) + if device_index not in device_infos: + device_infos[device_index] = DeviceInfo(device_index) + return device_infos[device_index] + + for module in ( + humming.config.config, + humming.transform, + humming.tune, + humming.tune.process_input, + humming.utils.smem, + ): + if getattr(module, "DeviceInfo", None) is DeviceInfo: + module.DeviceInfo = cached_device_info + + def __getattr__(name: str) -> Any: spec = _EXPORTS.get(name) if spec is None: raise AttributeError(f"module 'vllm.utils.humming' has no attribute {name!r}") + _patch_humming() if ":" in spec: mod_path, attr = spec.split(":", 1) obj = getattr(importlib.import_module(mod_path), attr) diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 466bba586d..d3c8d7ea24 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -950,6 +950,16 @@ class AttentionImpl(AttentionImplBase[T], Generic[T]): """Whether this implementation supports fused QKNorm+MRoPE+KVCache.""" return False + def set_interleaved_v_cache(self): + """Opt this implementation into the interleaved V-cache layout for all of + its cache writers and readers. QkNormRopeKvCacheFusionPass calls this + for every eligible layer when it is constructed, i.e. before any KV + cache is written, so fused and unfused paths stay consistent whether or + not the pattern matches. Default is a no-op; backends whose kernels + read an interleaved V-cache under this fusion (e.g. ROCM_ATTN) override. + """ + return False + def fused_rope_kvcache_supported(self): """Does this attention implementation support RoPE+KVCache fusion. This is used by the RopeKVCacheFusionPass to only fuse the RoPE ops diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index d673868975..605f551bc2 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -35,6 +35,7 @@ from vllm.v1.attention.backends.mla.sparse_utils import ( flat_kv_row_view, prepare_sparse_mla_safe_lengths, ) +from vllm.v1.attention.ops.metadata import compute_token_to_req_indices from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec if TYPE_CHECKING: @@ -297,6 +298,11 @@ class FlashInferMLASparseTRTLLMMetadataBuilder(FlashInferMLASparseMetadataBuilde device: torch.device, ) -> None: super().__init__(kv_cache_spec, layer_names, vllm_config, device) + self.supports_draft_decode_metadata_update = ( + self.dcp_world_size == 1 + and self.pcp_world_size == 1 + and vllm_config.attention_config.hisparse_config is None + ) # Under DCP the workspace must hold the trtllm-gen softmax-stats slab. # The buffer address is baked into CUDA graphs, so it has to reach its @@ -314,6 +320,18 @@ class FlashInferMLASparseTRTLLMMetadataBuilder(FlashInferMLASparseMetadataBuilde ), ) + def update_draft_decode_metadata( + self, metadata: FlashInferMLASparseMetadata + ) -> None: + num_tokens = metadata.num_decode_tokens + if num_tokens: + compute_token_to_req_indices( + metadata.query_start_loc, + metadata.req_id_per_token, + num_tokens, + num_tokens, + ) + # Global workspace buffer (lazily initialized) _fi_sparse_workspace: torch.Tensor | None = None diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index fd4b562f53..bbbb111d8b 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -103,10 +103,11 @@ For DeepSeek V4.1, each token's KV cache is 528 Bytes, structured as: - **Last 16 bytes:** Scale factors, containing 16 `ue8m0` values, one per 32 consecutive `float8_e4m3` values (i.e. MXFP8). -The V4 and V4.1 records are not laid out token-by-token within a page: a page -holds all its data rows first and all its scale rows after, so a page is -`block_size * bytes_per_token` bytes rounded up to the decode kernel's TMA -stride (576 B for V4, 512 B for V4.1). +The V4 record is not laid out token-by-token within a page: a page holds all +its data rows first and all its scale rows after, so a page is +`block_size * bytes_per_token` bytes rounded up to the decode kernel's 576 B +TMA stride. The V4.1 records are interleaved: each token's scales directly +follow its own data, and the decode kernel's TMA stride is the record width. In the "nvfp4_ds_mla" format (SM100 only, DeepSeek V3.2 geometry), each token's KV cache is 352 Bytes, structured as: diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 4f18173343..468040257e 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -205,19 +205,16 @@ class PrepareUniformDecodeKernel( ) -# Below this many rows per rank, stay replicated (TP4 engages at 64K rows). -MIN_TP_SHARD_ROWS_PER_RANK = 16_384 - - def balanced_prefill_row_shard( seq_lens_cpu: torch.Tensor, query_lens_cpu: torch.Tensor, compress_ratio: int, tp_size: int, + min_rows: int, ) -> list[int] | None: """Return contiguous TP row counts balanced by indexer MQA work.""" num_rows = int(query_lens_cpu.sum()) - if tp_size < 2 or num_rows < MIN_TP_SHARD_ROWS_PER_RANK * tp_size: + if tp_size < 2 or num_rows < max(min_rows, tp_size): return None query_lens = query_lens_cpu.to(torch.int64) @@ -1584,6 +1581,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): prefill_query_lens_cpu, self.compress_ratio, self.tp_size, + self.vllm_config.attention_config.indexer_tp_shard_min_rows, ) prefill_metadata = DeepseekV32IndexerPrefillMetadata( chunks, diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index 65df2fcee4..b6c090aa3f 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -85,6 +85,7 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase): block_size: int = 64, packed_bytes_per_token: int = 584, packed_page_alignment: int = 576, + packed_block_stride_alignment: int | None = None, bounded_replay: bool = False, ): super().__init__() @@ -99,6 +100,7 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase): self.dtype = dtype self.packed_bytes_per_token = packed_bytes_per_token self.packed_page_alignment = packed_page_alignment + self.packed_block_stride_alignment = packed_block_stride_alignment compilation_config = get_current_vllm_config().compilation_config if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") @@ -135,8 +137,11 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase): state_content_bytes=( self.packed_bytes_per_token if uses_fp8_ds_mla_layout else None ), - # FlashMLA packing stride; 512B for FlashInfer sparse (#44577). + # Packed page alignment; 512B for FlashInfer sparse (#44577). alignment=self.packed_page_alignment if uses_fp8_ds_mla_layout else 512, + block_stride_alignment=( + self.packed_block_stride_alignment if uses_fp8_ds_mla_layout else None + ), model_version="deepseek_v4", kv_quant_mode=get_kv_quant_mode(self.cache_config.cache_dtype), ) @@ -211,6 +216,9 @@ class DeepseekSparseSWAMetadata: # window attention may read (SWA bounded replay); # zeros when nothing replays. replay_start: torch.Tensor | None = None + # PCP cache writes use the gathered slots; queries use rank-local slots. + cache_slot_mapping: torch.Tensor | None = None + cache_positions: torch.Tensor | None = None # Number of decode/prefill requests/tokens (batch is reordered: decodes first) num_decodes: int = 0 diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index e1c4c8d067..d01f5b09a2 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -117,6 +117,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): "vllm.models.deepseek_v41.nvidia.flashinfer_sparse." "DeepseekV4FlashInferMLASparseBackend" ) + ROCM_FLASHMLA_SPARSE_DSV41 = ( + "vllm.models.deepseek_v41.amd.rocm.DeepseekV41ROCMAiterMLASparseBackend" + ) FLASHMLA_MEGA_ATTN_DSV41 = ( "vllm.models.deepseek_v41.sparse_mla.FlashMLAMegaAttnBackend" ) diff --git a/vllm/v1/attention/backends/rocm_aiter_fa.py b/vllm/v1/attention/backends/rocm_aiter_fa.py index 0ec0d5dd36..04c614a1d2 100644 --- a/vllm/v1/attention/backends/rocm_aiter_fa.py +++ b/vllm/v1/attention/backends/rocm_aiter_fa.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Attention layer with AiterFlashAttention.""" -from dataclasses import dataclass, replace +from dataclasses import dataclass, field, replace from typing import ClassVar import torch @@ -152,8 +152,8 @@ if current_platform.is_rocm(): k_reg = tl.load(key_cache_ptr_offset + k_reg_offset) v_reg = tl.load(value_cache_ptr_offset + v_reg_offset) if DEQUANT: - k_scale = 1.0 - v_scale = 1.0 + k_scale = tl.load(k_scale_ptr) + v_scale = tl.load(v_scale_ptr) k_reg = k_reg.to(tl.float32) * k_scale v_reg = v_reg.to(tl.float32) * v_scale tl.store(key_ptr_offset + col_offsets, k_reg) @@ -270,8 +270,8 @@ if current_platform.is_rocm(): k_val = tl.load(key_ptr + src_offset_k + offset) v_val = tl.load(value_ptr + src_offset_v + offset) if QUANT: - k_scale = 1.0 - v_scale = 1.0 + k_scale = tl.load(k_scale_ptr) + v_scale = tl.load(v_scale_ptr) k_dtype = key_cache_ptr.type.element_ty v_dtype = value_cache_ptr.type.element_ty k_val = (k_val.to(tl.float32) / k_scale).to(k_dtype) @@ -410,11 +410,48 @@ class AiterFlashAttentionMetadata: kv_sharing_metadata: AiterKVSharingMetadata | None = None +@dataclass +class AsmKvScaleTables: + """Dense per-slot scale tables the asm decode kernel reads for an fp8 cache. + + One table per distinct (device, shape, scale), shared by every layer and + every metadata builder in the process, so the target and draft models with + equal scales reuse a single allocation. Tables are constants: never write + to one in place; a per-token or dynamic scale needs its own buffer. + """ + + tables: dict[tuple[str, int, int, int, float], torch.Tensor] = field( + default_factory=dict + ) + + def get( + self, + device: torch.device, + num_blocks: int, + num_kv_heads: int, + block_size: int, + scale: float, + ) -> torch.Tensor: + key = (str(device), num_blocks, num_kv_heads, block_size, scale) + table = self.tables.get(key) + if table is None: + table = torch.full( + (num_blocks, num_kv_heads, block_size), + scale, + dtype=torch.float32, + device=device, + ) + self.tables[key] = table + return table + + class AiterFlashAttentionMetadataBuilder( AttentionMetadataBuilder[AiterFlashAttentionMetadata] ): _cudagraph_support = AttentionCGSupport.UNIFORM_BATCH + asm_kv_scale_tables: ClassVar[AsmKvScaleTables] = AsmKvScaleTables() + def __init__( self, kv_cache_spec: AttentionSpec, @@ -466,7 +503,8 @@ class AiterFlashAttentionMetadataBuilder( dtype=self.model_config.dtype, device=device, ) - self.scale = torch.tensor([1.0], dtype=torch.float, device=self.device) + self.k_scale: dict[str, torch.Tensor] | None = None + self.v_scale: dict[str, torch.Tensor] | None = None self.kv_sharing_workspace = ( torch.empty( ( @@ -488,6 +526,38 @@ class AiterFlashAttentionMetadataBuilder( common_prefix_len=0, common_attn_metadata=common_attn_metadata ) + def _build_asm_kv_scale_tables(self) -> None: + # Only for fp8 shuffle kv cache: the asm decode kernel reads a dense + # [num_blocks, num_kv_heads, block_size] scale table per layer. Map each + # layer to the shared table for its scale once the KV cache is bound. + if ( + self.k_scale is not None + or not rocm_aiter_ops.is_shuffle_kv_cache_enabled() + or not is_quantized_kv_cache(self.vllm_config.cache_config.cache_dtype) + ): + return + attn_layers = self.vllm_config.compilation_config.static_forward_context + k_scale_tables: dict[str, torch.Tensor] = {} + v_scale_tables: dict[str, torch.Tensor] = {} + for layer_name in self.layer_names: + layer = attn_layers[layer_name] + num_blocks = layer.kv_cache.shape[0] + k_scale_tables[layer_name] = self.asm_kv_scale_tables.get( + self.device, + num_blocks, + self.num_heads_kv, + self.block_size, + layer._k_scale_float, + ) + v_scale_tables[layer_name] = self.asm_kv_scale_tables.get( + self.device, + num_blocks, + self.num_heads_kv, + self.block_size, + layer._v_scale_float, + ) + self.k_scale, self.v_scale = k_scale_tables, v_scale_tables + def build( self, common_prefix_len: int, @@ -499,24 +569,7 @@ class AiterFlashAttentionMetadataBuilder( common_attn_metadata, decode_threshold=self.reorder_batch_threshold, ) - # Allocate scales for fp8 shuffle kv cache with shuffle_kv_cache enabled - if ( - rocm_aiter_ops.is_shuffle_kv_cache_enabled() - and self.scale.numel() == 1 - and is_quantized_kv_cache(self.vllm_config.cache_config.cache_dtype) - ): - # Size the scales from a layer this builder owns. The draft model - # runs its own builder over its own KV cache, so the first layer of - # the whole config can carry an unrelated block count. - kv_cache_shape = self.vllm_config.compilation_config.static_forward_context[ - self.layer_names[0] - ].kv_cache.shape - num_blocks = kv_cache_shape[0] - self.scale = torch.ones( - [num_blocks, self.num_heads_kv, self.block_size], - dtype=torch.float32, - device=self.device, - ) + self._build_asm_kv_scale_tables() ( num_decodes, num_extends, @@ -709,8 +762,8 @@ class AiterFlashAttentionMetadataBuilder( decode_metadata=decode_metadata, prefill_metadata=prefill_metadata, extend_metadata=extend_metadata, - k_scale=self.scale, - v_scale=self.scale, + k_scale=self.k_scale, + v_scale=self.v_scale, kv_sharing_metadata=kv_sharing_metadata, ) return attn_metadata @@ -726,6 +779,7 @@ class AiterFlashAttentionMetadataBuilder( skip split_decodes_prefills_and_extends() and avoid all .cpu() / .item() calls that would otherwise break CUDA graph capture. """ + self._build_asm_kv_scale_tables() num_reqs = common_attn_metadata.num_reqs num_tokens = common_attn_metadata.num_actual_tokens @@ -749,8 +803,8 @@ class AiterFlashAttentionMetadataBuilder( decode_metadata=decode_metadata, prefill_metadata=None, extend_metadata=None, - k_scale=self.scale, - v_scale=self.scale, + k_scale=self.k_scale, + v_scale=self.v_scale, ) @@ -1249,9 +1303,6 @@ class AiterFlashAttentionImpl(AttentionImpl): extend_outputs = output[extend_tokens_slice] k_scale = layer._k_scale v_scale = layer._v_scale - if rocm_aiter_ops.is_shuffle_kv_cache_enabled(): - k_scale = attn_metadata.k_scale - v_scale = attn_metadata.v_scale self.extend_forward( attn_metadata=attn_metadata, query=extend_queries, @@ -1428,16 +1479,13 @@ class AiterFlashAttentionImpl(AttentionImpl): device=query.device, ) max_logits = torch.empty_like(exp_sums) - k_qscale = ( - layer._k_scale - if attn_metadata.k_scale is None - else attn_metadata.k_scale - ) - v_qscale = ( - layer._v_scale - if attn_metadata.v_scale is None - else attn_metadata.v_scale - ) + # asm indexes this layer's dense scale buffer from build(); + # the HIP kernel reads the per-tensor scale directly. + k_qscale_asm = v_qscale_asm = None + if attn_metadata.k_scale is not None: + assert attn_metadata.v_scale is not None + k_qscale_asm = attn_metadata.k_scale[layer.layer_name] + v_qscale_asm = attn_metadata.v_scale[layer.layer_name] rocm_aiter_ops.paged_attention_common( Q=query[:num_decode_tokens], K=new_key_cache, @@ -1452,10 +1500,10 @@ class AiterFlashAttentionImpl(AttentionImpl): 0 ), scale=self.scale, - K_QScale_hip=k_qscale, - V_QScale_hip=v_qscale, - K_QScale_asm=k_qscale, - V_QScale_asm=v_qscale, + K_QScale_hip=layer._k_scale, + V_QScale_hip=layer._v_scale, + K_QScale_asm=k_qscale_asm, + V_QScale_asm=v_qscale_asm, out_=output[:num_decode_tokens], kv_cache_dtype=self.kv_cache_dtype, ) @@ -1542,9 +1590,6 @@ class AiterFlashAttentionImpl(AttentionImpl): # the reshape_and_cache_flash op uses the slot_mapping's shape # to determine the number of actual tokens. if rocm_aiter_ops.is_shuffle_kv_cache_enabled(): - # We may calculate per token quant scale in - # reshape_and_cache_shuffle_triton which might differ from - # vllm's style when shuffle layout is used. k_scale = layer._k_scale v_scale = layer._v_scale assert k_scale is not None and v_scale is not None, ( @@ -1581,12 +1626,7 @@ class AiterFlashAttentionImpl(AttentionImpl): ) def fused_qk_norm_rope_kvcache_supported(self): - # Only fuse when shuffle layout is off; the shuffle write path uses a - # dedicated cache update, mirroring fused_rope_kvcache_supported. - return ( - rocm_aiter_ops.is_enabled() - and not rocm_aiter_ops.is_shuffle_kv_cache_enabled() - ) + return rocm_aiter_ops.is_enabled() def do_qk_norm_rope_kvcache_update( self, @@ -1624,6 +1664,7 @@ class AiterFlashAttentionImpl(AttentionImpl): v_scale=layer._v_scale_cpu, kv_cache_dtype=self.kv_cache_dtype, use_shuffle_layout=rocm_aiter_ops.is_shuffle_kv_cache_enabled(), + return_kv=True, ) def do_rope_and_kv_cache_update( diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py index 9f8e9e9ab1..b05482c007 100644 --- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py +++ b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py @@ -15,6 +15,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kFp8StaticTensorSym, ) +from vllm.platforms.rocm import on_gfx1250 from vllm.utils.torch_utils import is_quantized_kv_cache from vllm.v1.attention.backend import ( AttentionLayer, @@ -28,7 +29,12 @@ from vllm.v1.attention.backends.rocm_attn import ( RocmAttentionMetadata, RocmAttentionMetadataBuilder, ) -from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + KVCacheLayout, + KVQuantMode, + get_kv_quant_mode, +) logger = init_logger(__name__) @@ -186,12 +192,34 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl): # RDNA3 has no low precision MMA, preserve query precision self.supports_quant_query_input = not on_gfx11() + # Layout flag: must mirror get_kv_cache_shape, which is attn_type + # agnostic because the allocation is shared across layers. + self._shuffle_kv_layout = ( + rocm_aiter_ops.is_shuffle_kv_cache_enabled() and on_gfx1250() + ) + # Kernel flag: the shuffled reader/writer is causal decoder only. + self._use_shuffle_kv = ( + self._shuffle_kv_layout and self.attn_type != AttentionType.ENCODER_DECODER + ) + def _split_kv_cache( self, kv_cache: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: + if self._shuffle_kv_layout: + # (2, B, N, H, hs) -> ((B, N, H, hs), (B, N, H, hs)) + return kv_cache.unbind(0) # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs)) return kv_cache.transpose(1, 2).split(self.head_size, dim=-1) + def _shuffle_kv_views( + self, key_cache: torch.Tensor, value_cache: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + nb, bs, nkv, hd = key_cache.shape + x = 16 // key_cache.element_size() + kc = key_cache.reshape(nb, nkv, hd // x, bs, x) + vc = value_cache.reshape(nb, nkv, bs // x, hd, x) + return kc, vc + def forward( self, layer: torch.nn.Module, @@ -264,6 +292,9 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl): key_cache = key_cache.view(self.fp8_dtype) value_cache = value_cache.view(self.fp8_dtype) + if self._use_shuffle_kv: + key_cache, value_cache = self._shuffle_kv_views(key_cache, value_cache) + cu_seqlens_q = attn_metadata.query_start_loc seqused_k = attn_metadata.seq_lens max_seqlen_q = attn_metadata.max_query_len @@ -291,6 +322,7 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl): v_descale=layer._v_scale, sinks=self.sinks, output_scale=output_scale, + shuffled_kv_cache=self._use_shuffle_kv, ) else: # The aiter kernel is causal-only. Non-causal cross-attention @@ -339,6 +371,27 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl): return key_cache, value_cache = self._split_kv_cache(kv_cache) + if self._use_shuffle_kv: + from vllm.v1.attention.backends.rocm_aiter_fa import ( + reshape_and_cache_shuffle_triton, + ) + + if get_kv_quant_mode(self.kv_cache_dtype) == KVQuantMode.FP8_PER_TENSOR: + key_cache = key_cache.view(self.fp8_dtype) + value_cache = value_cache.view(self.fp8_dtype) + _dummy = torch.empty(1, dtype=torch.float32, device=key.device) + reshape_and_cache_shuffle_triton( + key, + value, + key_cache, + value_cache, + slot_mapping, + self.kv_cache_dtype, + _dummy, + _dummy, + ) + return + # Reshape the input keys and values and store them in the cache. ops.reshape_and_cache_flash( key, @@ -360,6 +413,11 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl): def fused_qk_norm_mrope_kvcache_supported(self) -> bool: return IS_AITER_FOUND + def set_interleaved_v_cache(self): + # No-op override of RocmAttentionImpl: this backend + # decodes with the AITER triton unified-attention kernel + pass + def do_qk_norm_rope_kvcache_update( self, layer: AttentionLayer, @@ -395,7 +453,8 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl): k_scale=layer._k_scale_cpu, v_scale=layer._v_scale_cpu, kv_cache_dtype=self.kv_cache_dtype, - use_shuffle_layout=False, + use_shuffle_layout=self._use_shuffle_kv, + return_kv=False, ) def do_qk_norm_mrope_kvcache_update( @@ -456,13 +515,16 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl): # we use direct Q, K, V tensors without caching return key_cache, value_cache = self._split_kv_cache(kv_cache) - flash_layout = True is_fp8_kv_cache = is_quantized_kv_cache(self.kv_cache_dtype) if is_fp8_kv_cache: key_cache = key_cache.view(self.fp8_dtype) value_cache = value_cache.view(self.fp8_dtype) + if self._use_shuffle_kv: + key_cache, value_cache = self._shuffle_kv_views(key_cache, value_cache) + flash_layout = not self._use_shuffle_kv + rocm_aiter_ops.triton_rope_and_cache( query, key, diff --git a/vllm/v1/attention/backends/rocm_attn.py b/vllm/v1/attention/backends/rocm_attn.py index 471bb9f058..16b4c38903 100644 --- a/vllm/v1/attention/backends/rocm_attn.py +++ b/vllm/v1/attention/backends/rocm_attn.py @@ -292,6 +292,8 @@ class RocmAttentionImpl(AttentionImpl): f"num_heads: {num_heads}." ) + self._use_interleaved_v_cache = False + def _forward_encoder_attention( self, query: torch.Tensor, @@ -450,6 +452,7 @@ class RocmAttentionImpl(AttentionImpl): sm_scale=self.scale, output_scale=output_scale, sinks=self.sinks, + use_interleaved_v_cache=self._use_interleaved_v_cache, causal=attn_metadata.causal, ) @@ -475,6 +478,12 @@ class RocmAttentionImpl(AttentionImpl): block_size = value_cache.shape[3] has_native_layout = has_native_kv_cache_layout(key_cache, value_cache) + # Match the fused writer's interleaved V-cache when fusion is active. + interleaved_v = self._use_interleaved_v_cache + interleaved_v_pack_factor = ( + (16 // value_cache.element_size()) if interleaved_v else 0 + ) + if block_size in (16, 32) and has_native_layout: # Normal 16, 32 with contiguous blocks: use vLLM native HIP C++ logic. PagedAttention.write_to_paged_cache( @@ -486,6 +495,7 @@ class RocmAttentionImpl(AttentionImpl): self.kv_cache_dtype, layer._k_scale, layer._v_scale, + interleaved_v, ) else: # Non-standard blocks and hybrid attention/Mamba layouts need the @@ -501,8 +511,61 @@ class RocmAttentionImpl(AttentionImpl): self.kv_cache_dtype, layer._k_scale, layer._v_scale, + interleaved_v_pack_factor, ) + def fused_qk_norm_rope_kvcache_supported(self): + return rocm_aiter_ops.is_enabled() + + def set_interleaved_v_cache(self): + # Decode/prefill kernels must read the interleaved V-cache the AITER + # fused write emits when use_shuffle_layout=True. + self._use_interleaved_v_cache = True + + def do_qk_norm_rope_kvcache_update( + self, + layer: AttentionLayer, + qkv: torch.Tensor, + q_out: torch.Tensor, + k_out: torch.Tensor, + positions: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + rms_norm_eps: float, + cos_sin_cache: torch.Tensor, + is_neox: bool, + kv_cache: torch.Tensor, + layer_slot_mapping: torch.Tensor, + ): + # kv_cache is logical [num_blocks, 2, block_size, num_kv_heads * head_size] + # (see forward); the AITER shuffle writer takes one (B, N, H, hs) per side. + key_cache, value_cache = kv_cache.unflatten( + -1, (self.num_kv_heads, self.head_size) + ).unbind(1) + use_shuffle_layout = self._use_interleaved_v_cache + rocm_aiter_ops.do_qk_norm_rope_kvcache_update( + qkv=qkv, + q_weight=q_weight, + k_weight=k_weight, + cos_sin_cache=cos_sin_cache, + positions=positions, + num_heads_q=self.num_heads, + num_heads_k=self.num_kv_heads, + head_dim=self.head_size, + is_neox=is_neox, + rms_norm_eps=rms_norm_eps, + q_out=q_out, + k_out=k_out, + key_cache=key_cache, + value_cache=value_cache, + slot_mapping=layer_slot_mapping, + k_scale=layer._k_scale_cpu, + v_scale=layer._v_scale_cpu, + kv_cache_dtype=self.kv_cache_dtype, + use_shuffle_layout=use_shuffle_layout, + return_kv=True, + ) + def fused_rope_kvcache_supported(self): return rocm_aiter_ops.is_enabled() diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index f93156d53f..7e0a7fb644 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/vllm/v1/attention/backends/utils.py @@ -18,6 +18,7 @@ from typing_extensions import runtime_checkable from vllm.config import CacheConfig, VllmConfig, get_layers_from_vllm_config from vllm.config.cache import _layout_from_name +from vllm.triton_utils import tl, triton from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d, np_to_pinned_tensor @@ -1160,6 +1161,75 @@ def get_dcp_local_seq_lens( return dcp_local_seq_lens +@triton.jit( + do_not_specialize_on_alignment=["out_ptr", "block_table_ptr", "seq_lens_ptr"] +) +def _mamba_align_block_table_kernel( + out_ptr, + block_table_ptr, + block_table_stride_0: tl.int64, + block_table_stride_1, + seq_lens_ptr, + seq_lens_stride, + block_size, + n_out, + BLOCK: tl.constexpr, +): + """Gather the last ``n_out`` block-table entries of each request. + + One program per request. This replaces a five-statement PyTorch expression + that expanded to seven kernels over a ``(num_reqs,)`` tensor -- one element + in a single-request decode -- so almost all of its cost was launch + overhead rather than work. + + Both block-table strides are passed explicitly. The expression this replaces + was ``torch.gather(block_table, 1, ...)``, which honours any layout, and + this is a shared utility -- ``gdn_attn``, ``mamba_attn``, ``short_conv_attn`` + and ``linear_attn`` all call it -- so assuming a column stride of 1 would + silently read the wrong slots for a caller holding a non-contiguous view. + ``kimi_k3/nvidia/kda_metadata.py`` passes both strides for the same reason. + + Two annotations above are load-bearing and easy to drop by accident: + + ``block_table_stride_0: tl.int64`` + Without it ``req_idx * block_table_stride_0`` is computed in i32 -- + ``program_id`` is i32 and Triton specializes a stride below 2^31 to + i32 -- so the product wraps once a row offset reaches 2^31 elements + (a 4096 x 625000 table faults from row 3436 on MI355X). + ``torch.gather``, which this replaces, indexes in 64-bit and has no + such limit, so narrowing it would be a regression. + ``mamba_utils.py:43`` annotates its row stride the same way. + + ``do_not_specialize_on_alignment`` on the three pointers + Triton puts pointer alignment in the specialization key, and the MRV1 + DBO path hands this kernel ``seq_lens[request_slice]`` / + ``block_table_tensor[request_slice]`` VIEWS (ubatch_utils.py:278/316). + When a ubatch starts at a request index that is not a multiple of 4 + the view is differently aligned and a fresh variant JIT-compiles + mid-step -- 25-28 ms cold, ~1.5 ms warm, and an error under + ``jit_monitor_mode=error``. The torch expression this replaced never + compiled anything, so that recompile would be a regression introduced + by this change. + """ + req_idx = tl.program_id(0) + seq_len = tl.load(seq_lens_ptr + req_idx * seq_lens_stride) + # A zero-length request in a CUDA graph has an invalid block table; clamp + # to 0 so it reads a valid row, exactly as the previous clamp_ did. + start = tl.maximum((seq_len - 1) // block_size, 0) + offs = tl.arange(0, BLOCK) + mask = offs < n_out + vals = tl.load( + block_table_ptr + + req_idx * block_table_stride_0 + + (start + offs) * block_table_stride_1, + mask=mask, + other=0, + ) + # `out` is freshly allocated by the caller and therefore contiguous, so its + # row stride is exactly n_out. + tl.store(out_ptr + req_idx * n_out + offs, vals, mask=mask) + + def mamba_get_block_table_tensor( block_table: torch.Tensor, seq_lens: torch.Tensor, @@ -1180,16 +1250,58 @@ def mamba_get_block_table_tensor( return block_table else: assert isinstance(kv_cache_spec, MambaSpec) - # NOTE: For 0-length requests in CUDA graph, use a start_index of 0 - # to handle the invalid block table. - start_indices = (seq_lens - 1) // kv_cache_spec.block_size - start_indices.clamp_(min=0) - # Use int32 for arithmetic to avoid dtype promotion overhead, - # then convert to int64 for gather (which requires Long indices) - offsets = torch.arange( - 1 + kv_cache_spec.num_speculative_blocks, - device=block_table.device, - dtype=torch.int32, + # The arithmetic here is trivial -- a divide, a clamp, an arange, an + # add, a cast and a gather -- but as PyTorch ops it was seven kernel + # launches over a (num_reqs,) tensor, which holds a single element in a + # one-request decode. One Triton kernel does the same work in one + # launch. Zero-length requests in a CUDA graph still clamp to row 0. + num_reqs = block_table.shape[0] + n_out = 1 + kv_cache_spec.num_speculative_blocks + # The expression this replaced took its row count from `seq_lens` (a + # `torch.gather` whose index tensor is built from it), so a mismatch + # raised and a cross-device pair raised. Taking `num_reqs` from + # `block_table` instead silently changed all of that: with 4 block-table + # rows and 2 `seq_lens`, the old code returned 2 correct rows and the new + # kernel returns 4, two of them read from memory past the end of + # `seq_lens`; a pinned-CPU `seq_lens` with a GPU `block_table` used to + # raise and would now be read zero-copy from host memory. + # + # No in-tree caller does any of that today, but this helper is shared by + # gdn_attn, mamba_attn, short_conv_attn and linear_attn. Both are host + # metadata, so this costs no synchronisation, and it keeps the CPU + # branch below and the GPU kernel agreeing on the same inputs. + assert seq_lens.shape[0] == num_reqs, ( + f"seq_lens has {seq_lens.shape[0]} rows, block_table has {num_reqs}" ) - indices_to_gather = (start_indices.unsqueeze(1) + offsets).to(torch.int64) - return torch.gather(block_table, 1, indices_to_gather) + assert seq_lens.device == block_table.device, ( + f"seq_lens is on {seq_lens.device}, block_table on {block_table.device}" + ) + if not block_table.is_cuda: + # Triton runs on the accelerator only, and this helper IS called + # with CPU tensors: the GDN and linear-attention metadata builders + # are unit-tested on CPU (`DEVICE = torch.device("cpu")`). The + # expression this replaced was plain torch and handled that, so + # keep it for the CPU path rather than narrowing the contract. + # kimi_k3/nvidia/kda_metadata.py asserts `is_cuda` instead; a + # shared utility cannot, because its callers legitimately pass CPU. + start_indices = (seq_lens - 1) // kv_cache_spec.block_size + start_indices = start_indices.clamp(min=0) + offsets = torch.arange(n_out, device=block_table.device, dtype=torch.int32) + return torch.gather( + block_table, 1, (start_indices.unsqueeze(1) + offsets).to(torch.int64) + ) + out = block_table.new_empty((num_reqs, n_out)) + if num_reqs == 0: + return out + _mamba_align_block_table_kernel[(num_reqs,)]( + out, + block_table, + block_table.stride(0), + block_table.stride(1), + seq_lens, + seq_lens.stride(0), + kv_cache_spec.block_size, + n_out, + BLOCK=triton.next_power_of_2(n_out), + ) + return out diff --git a/vllm/v1/attention/ops/chunked_prefill_paged_decode.py b/vllm/v1/attention/ops/chunked_prefill_paged_decode.py index f48018e8fc..5f9e28d7ea 100644 --- a/vllm/v1/attention/ops/chunked_prefill_paged_decode.py +++ b/vllm/v1/attention/ops/chunked_prefill_paged_decode.py @@ -84,6 +84,7 @@ def kernel_paged_attention_2d( query_start_len_ptr, # [num_seqs+1] USE_SINKS: tl.constexpr, # bool USE_FP8: tl.constexpr, + INTERLEAVED_V_PACK_FACTOR: tl.constexpr, FP8_MIN: tl.constexpr = float8_info.min, FP8_MAX: tl.constexpr = float8_info.max, ): @@ -170,13 +171,24 @@ def kernel_paged_attention_2d( + (offs_d[:, None] % x) * stride_k_cache_4 ) - # 4D addressing logic of V (Slot is innermost) - v_offset = ( - p_block_idx[:, None] * stride_v_cache_0 - + kv_head_idx * stride_v_cache_1 - + offs_d[None, :] * stride_v_cache_2 - + internal_offsets[:, None] * stride_v_cache_3 - ) + # V addressing: standard 4D or interleaved 5D + if INTERLEAVED_V_PACK_FACTOR > 0: + v_offset = ( + p_block_idx[:, None] * stride_v_cache_0 + + kv_head_idx * stride_v_cache_1 + + (internal_offsets[:, None] // INTERLEAVED_V_PACK_FACTOR) + * HEAD_SIZE + * INTERLEAVED_V_PACK_FACTOR + + offs_d[None, :] * INTERLEAVED_V_PACK_FACTOR + + (internal_offsets[:, None] % INTERLEAVED_V_PACK_FACTOR) + ) + else: + v_offset = ( + p_block_idx[:, None] * stride_v_cache_0 + + kv_head_idx * stride_v_cache_1 + + offs_d[None, :] * stride_v_cache_2 + + internal_offsets[:, None] * stride_v_cache_3 + ) # Only the final tile can straddle seq_len. Slots >= seq_len are # unwritten KV cache that may hold NaN/garbage; they are score-masked @@ -304,6 +316,7 @@ def chunked_prefill_paged_decode( # Optional tensor for sinks sinks=None, is_block_table_ptr: bool = False, + use_interleaved_v_cache: bool = False, causal: bool = True, ): if sm_scale is None: @@ -314,6 +327,10 @@ def chunked_prefill_paged_decode( if sliding_window is None or sliding_window <= 0: sliding_window = 0 + interleaved_v_pack_factor = ( + (16 // value_cache.element_size()) if use_interleaved_v_cache else 0 + ) + if max_query_len > 1: context_attention_fwd( q=query, @@ -336,6 +353,7 @@ def chunked_prefill_paged_decode( skip_decode=True, fp8_out_scale=output_scale, sinks=sinks, + interleaved_v_pack_factor=interleaved_v_pack_factor, causal=causal, ) @@ -428,6 +446,7 @@ def chunked_prefill_paged_decode( k_scale=k_scale, v_scale=v_scale, fp8_out_scale=output_scale, + use_interleaved_v_cache=use_interleaved_v_cache, ) else: logger.warning_once( @@ -504,4 +523,5 @@ def chunked_prefill_paged_decode( query_start_len_ptr=query_start_loc, USE_SINKS=sinks is not None, USE_FP8=output_scale is not None, + INTERLEAVED_V_PACK_FACTOR=interleaved_v_pack_factor, ) diff --git a/vllm/v1/attention/ops/paged_attn.py b/vllm/v1/attention/ops/paged_attn.py index 896e929b54..3215a2f04b 100644 --- a/vllm/v1/attention/ops/paged_attn.py +++ b/vllm/v1/attention/ops/paged_attn.py @@ -38,6 +38,7 @@ class PagedAttention: kv_cache_dtype: str, k_scale: torch.Tensor, v_scale: torch.Tensor, + interleaved_v: bool, ) -> None: ops.reshape_and_cache( key, @@ -48,4 +49,5 @@ class PagedAttention: kv_cache_dtype, k_scale, v_scale, + interleaved_v, ) diff --git a/vllm/v1/attention/ops/prefix_prefill.py b/vllm/v1/attention/ops/prefix_prefill.py index 2577f52ca5..2fd9df983a 100644 --- a/vllm/v1/attention/ops/prefix_prefill.py +++ b/vllm/v1/attention/ops/prefix_prefill.py @@ -56,6 +56,8 @@ def _paged_kv_cache_offsets( stride_v_cache_d, stride_v_cache_bl, PHYSICAL_BLOCK_SIZE: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + INTERLEAVED_V_PACK_FACTOR: tl.constexpr = 0, MASK_BLOCK_TABLE: tl.constexpr = False, ): """Compute paged K/V cache element offsets for a tile of token positions. @@ -65,7 +67,9 @@ def _paged_kv_cache_offsets( tiles), then the physical element offsets are built with the same layout the kernel uses everywhere: K cache: [num_blocks, num_kv_heads, head_size // x, block_size, x] - V cache: [num_blocks, num_kv_heads, head_size, block_size] + V cache: [num_blocks, num_kv_heads, head_size, block_size], or the + interleaved layout [num_blocks, num_kv_heads, block_size // x, + head_size, x] when `INTERLEAVED_V_PACK_FACTOR` (= x) > 0. Returns `(off_k, off_v)` with shapes [D, N] and [N, D] respectively. When `MASK_BLOCK_TABLE` is set, the block-table load is masked with @@ -94,12 +98,25 @@ def _paged_kv_cache_offsets( + internal[None, :] * stride_k_cache_bl + (offs_d[:, None] % x) * stride_k_cache_x ) - off_v = ( - bn[:, None] * stride_v_cache_bs - + cur_kv_head * stride_v_cache_h - + offs_d[None, :] * stride_v_cache_d - + internal[:, None] * stride_v_cache_bl - ) + if INTERLEAVED_V_PACK_FACTOR > 0: + # Interleaved V-cache: [num_blocks, num_kv_heads, block_size/x, + # head_size, x] flattened onto the 4D stride layout. + off_v = ( + bn[:, None] * stride_v_cache_bs + + cur_kv_head * stride_v_cache_h + + (internal[:, None] // INTERLEAVED_V_PACK_FACTOR) + * BLOCK_DMODEL + * INTERLEAVED_V_PACK_FACTOR + + offs_d[None, :] * INTERLEAVED_V_PACK_FACTOR + + (internal[:, None] % INTERLEAVED_V_PACK_FACTOR) + ) + else: + off_v = ( + bn[:, None] * stride_v_cache_bs + + cur_kv_head * stride_v_cache_h + + offs_d[None, :] * stride_v_cache_d + + internal[:, None] * stride_v_cache_bl + ) return off_k, off_v @@ -159,6 +176,7 @@ def _fwd_kernel( USE_FP8: tl.constexpr, CAUSAL: tl.constexpr = True, MAX_Q_LEN: tl.constexpr = 0, + INTERLEAVED_V_PACK_FACTOR: tl.constexpr = 0, MAX_CTX_LEN: tl.constexpr = 0, FP8_MIN: tl.constexpr = float8_info.min, FP8_MAX: tl.constexpr = float8_info.max, @@ -263,6 +281,8 @@ def _fwd_kernel( stride_v_cache_d, stride_v_cache_bl, PHYSICAL_BLOCK_SIZE, + BLOCK_DMODEL, + INTERLEAVED_V_PACK_FACTOR, ) if ( @@ -403,6 +423,8 @@ def _fwd_kernel( stride_v_cache_d, stride_v_cache_bl, PHYSICAL_BLOCK_SIZE, + BLOCK_DMODEL, + INTERLEAVED_V_PACK_FACTOR, MASK_BLOCK_TABLE=True, ) k_cur_load = tl.load( @@ -540,6 +562,7 @@ def _fwd_kernel_alibi( BLOCK_DMODEL_PADDED: tl.constexpr, # head size padded to a power of 2 BLOCK_N: tl.constexpr, SKIP_DECODE: tl.constexpr, + INTERLEAVED_V_PACK_FACTOR: tl.constexpr, ): # attn_bias[] cur_batch = tl.program_id(0) @@ -608,12 +631,24 @@ def _fwd_kernel_alibi( + ((start_n + offs_n[None, :]) % block_size) * stride_k_cache_bl + (offs_d[:, None] % x) * stride_k_cache_x ) - off_v = ( - bn[:, None] * stride_v_cache_bs - + cur_kv_head * stride_v_cache_h - + offs_d[None, :] * stride_v_cache_d - + (start_n + offs_n[:, None]) % block_size * stride_v_cache_bl - ) + if INTERLEAVED_V_PACK_FACTOR > 0: + v_internal = (start_n + offs_n[:, None]) % block_size + off_v = ( + bn[:, None] * stride_v_cache_bs + + cur_kv_head * stride_v_cache_h + + (v_internal // INTERLEAVED_V_PACK_FACTOR) + * BLOCK_DMODEL + * INTERLEAVED_V_PACK_FACTOR + + offs_d[None, :] * INTERLEAVED_V_PACK_FACTOR + + (v_internal % INTERLEAVED_V_PACK_FACTOR) + ) + else: + off_v = ( + bn[:, None] * stride_v_cache_bs + + cur_kv_head * stride_v_cache_h + + offs_d[None, :] * stride_v_cache_d + + (start_n + offs_n[:, None]) % block_size * stride_v_cache_bl + ) k_load = tl.load( K_cache + off_k, mask=dim_mask[:, None] & ((start_n + offs_n[None, :]) < cur_batch_ctx_len), @@ -796,6 +831,7 @@ def context_attention_fwd( fp8_out_scale=None, sinks=None, is_block_table_ptr: bool = False, + interleaved_v_pack_factor: int = 0, causal: bool = True, ): q_dtype_is_f32 = q.dtype is torch.float32 @@ -937,6 +973,7 @@ def context_attention_fwd( BLOCK_DMODEL_PADDED=Lk_padded, BLOCK_N=BLOCK, SKIP_DECODE=skip_decode, + INTERLEAVED_V_PACK_FACTOR=interleaved_v_pack_factor, num_warps=NUM_WARPS, num_stages=1, ) @@ -1010,6 +1047,7 @@ def context_attention_fwd( num_warps=4, num_stages=1, USE_SINKS=sinks is not None, + INTERLEAVED_V_PACK_FACTOR=interleaved_v_pack_factor, CAUSAL=causal, KV_FROM_CACHE=kv_from_cache, **extra_kargs, diff --git a/vllm/v1/attention/ops/triton_reshape_and_cache_flash.py b/vllm/v1/attention/ops/triton_reshape_and_cache_flash.py index 882b6787ae..47261bbc73 100644 --- a/vllm/v1/attention/ops/triton_reshape_and_cache_flash.py +++ b/vllm/v1/attention/ops/triton_reshape_and_cache_flash.py @@ -52,6 +52,9 @@ def reshape_and_cache_kernel_flash( block_size: tl.constexpr, x: tl.constexpr, USE_HEAD_MAJOR_LAYOUT: tl.constexpr, + # >0 writes V interleaved as [Block, Head, Slot//kx, Dim, kx] (head-major + # layout only); 0 keeps the standard [Block, Head, Dim, Slot] V layout. + INTERLEAVED_V_PACK_FACTOR: tl.constexpr, # FP8 flags FP8_KV_CACHE: tl.constexpr, # tune parameters @@ -76,13 +79,25 @@ def reshape_and_cache_kernel_flash( # Decompose the tile index back into head and dim coordinates. cur_head = tile_pos // head_size cur_dim = tile_pos % head_size - # Value addressing (4D): [Block, Head, Dim, Slot] - tgt_idx_v = ( - block_idx * block_stride - + cur_head * head_stride - + cur_dim * dim_stride_v - + block_offset * 1 - ) + # Value addressing: standard 4D [Block, Head, Dim, Slot], or the + # interleaved 5D [Block, Head, Slot//kx, Dim, kx] the fused AITER + # writer emits and the interleaved read path expects. + if INTERLEAVED_V_PACK_FACTOR > 0: + tgt_idx_v = ( + block_idx * block_stride + + cur_head * head_stride + + (block_offset // INTERLEAVED_V_PACK_FACTOR) + * (head_size * INTERLEAVED_V_PACK_FACTOR) + + cur_dim * INTERLEAVED_V_PACK_FACTOR + + (block_offset % INTERLEAVED_V_PACK_FACTOR) + ) + else: + tgt_idx_v = ( + block_idx * block_stride + + cur_head * head_stride + + cur_dim * dim_stride_v + + block_offset * 1 + ) # Key addressing (5D): [Block, Head, Dim//8, Slot, 8] tgt_idx_k = ( block_idx * block_stride @@ -371,11 +386,18 @@ def triton_reshape_and_cache_flash( kv_cache_dtype: str, # "auto", "fp8" k_scale: torch.Tensor, # float32 v_scale: torch.Tensor, # float32 + interleaved_v_pack_factor: int = 0, # >0: interleave V (head-major); 0: standard ): num_heads = key.shape[1] head_size = key.shape[2] use_head_major_layout = key_cache.ndim == 5 + if interleaved_v_pack_factor and not use_head_major_layout: + raise NotImplementedError( + "Interleaved V-cache write is only implemented for the head-major " + "(5D) cache layout; got a non-head-major cache with " + f"interleaved_v_pack_factor={interleaved_v_pack_factor}." + ) if use_head_major_layout: block_size = key_cache.shape[3] x = key_cache.shape[4] @@ -452,6 +474,7 @@ def triton_reshape_and_cache_flash( block_size=block_size, x=x, USE_HEAD_MAJOR_LAYOUT=use_head_major_layout, + INTERLEAVED_V_PACK_FACTOR=interleaved_v_pack_factor, FP8_KV_CACHE=FP8_KV_CACHE, # autotune parameters TILE_SIZE=TILE_SIZE, diff --git a/vllm/v1/core/block_pool.py b/vllm/v1/core/block_pool.py index 3b0102fe84..41839bd1f7 100644 --- a/vllm/v1/core/block_pool.py +++ b/vllm/v1/core/block_pool.py @@ -759,9 +759,11 @@ class BlockPool: """ for block in blocks: + if block.is_null: + continue # ref_cnt=0 means this block is in the free list (i.e. eviction # candidate), so remove it. - if block.ref_cnt == 0 and not block.is_null: + if block.ref_cnt == 0: self.free_block_queue.remove(block) block.ref_cnt += 1 if self.metrics_collector: @@ -785,11 +787,13 @@ class BlockPool: blocks_to_evict_first = [] other_pools: dict[BlockPool, list[KVCacheBlock]] = {} for block in ordered_blocks: + if block.is_null: # never pooled; its ref_cnt is not maintained + continue if block.pool is not None and block.pool is not self: other_pools.setdefault(block.pool, []).append(block) continue block.ref_cnt -= 1 - if block.ref_cnt == 0 and not block.is_null: + if block.ref_cnt == 0: if block.block_hash is None or not self.enable_caching: # LIFO reuse of non-cached blocks for better GPU locality. blocks_to_evict_first.append(block) diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 4f472170a8..0ac1c38818 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -14,7 +14,7 @@ from vllm.v1.core.kv_cache_coordinator import ( get_kv_cache_coordinator, ) from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector -from vllm.v1.core.kv_cache_utils import KVCacheBlock, KVCacheBlockCopy +from vllm.v1.core.kv_cache_utils import KVCacheBlock, KVCacheBlockCopy, trim_null_head from vllm.v1.core.single_type_kv_cache_manager import MambaManager from vllm.v1.kv_cache_interface import ( AttentionSpec, @@ -117,7 +117,7 @@ class KVCacheBlocks: return [ [ block.block_id - for block in group + for block in trim_null_head(group) if block.block_hash is None and not block.is_null ] for group in self.blocks diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 5c51a8d131..aadc9cda57 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -8,7 +8,15 @@ import math import os from collections.abc import Callable, Iterable, Iterator, Sequence from dataclasses import dataclass, field, replace -from typing import TYPE_CHECKING, Any, NamedTuple, NewType, TypeAlias, overload +from typing import ( + TYPE_CHECKING, + Any, + NamedTuple, + NewType, + TypeAlias, + TypeVar, + overload, +) from vllm import envs from vllm.config import VllmConfig @@ -216,6 +224,32 @@ class KVCacheBlock: ) +def null_head_len(blocks: Sequence[KVCacheBlock]) -> int: + """Number of leading null blocks, e.g. a sliding-window table's skipped prefix.""" + if not blocks or not blocks[0].is_null: + return 0 + lo, hi = 1, len(blocks) + while lo < hi: + mid = (lo + hi) // 2 + if blocks[mid].is_null: + lo = mid + 1 + else: + hi = mid + # All null entries are one shared object, so count() checks the run at C speed. + if blocks[:lo].count(blocks[0]) == lo: + return lo + return next(i for i, b in enumerate(blocks) if not b.is_null) + + +_BlockSeq = TypeVar("_BlockSeq", list[KVCacheBlock], Sequence[KVCacheBlock]) + + +def trim_null_head(blocks: _BlockSeq) -> _BlockSeq: + """``blocks`` without its leading null blocks; ``blocks`` itself if none.""" + nhl = null_head_len(blocks) + return blocks[nhl:] if nhl else blocks + + class KVCacheBlockCopy(NamedTuple): src_block_id: int dst_block_id: int diff --git a/vllm/v1/core/sched/async_scheduler.py b/vllm/v1/core/sched/async_scheduler.py index 290204d58c..0b3ab67b0f 100644 --- a/vllm/v1/core/sched/async_scheduler.py +++ b/vllm/v1/core/sched/async_scheduler.py @@ -60,14 +60,16 @@ class AsyncScheduler(Scheduler): self, request: Request, new_token_ids: list[int], is_stale: bool = False ) -> tuple[list[int], bool]: status_before_update = request.status + # Tokens trimmed at a stop resolve their placeholders too. + num_new_tokens = len(new_token_ids) new_token_ids, stopped = super()._update_request_with_output( - request, new_token_ids + request, new_token_ids, is_stale ) # Placeholders were zeroed at preemption; a stale delivery must not # decrement them (it would underflow). if not is_stale: - request.num_output_placeholders -= len(new_token_ids) + request.num_output_placeholders -= num_new_tokens assert request.num_output_placeholders >= 0 # Cache the new tokens. Preempted requests should be skipped. diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index a6ca6074df..c7619a2f66 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -368,6 +368,36 @@ class Scheduler(SchedulerInterface): # prefill batch fully drained the waiting queue. Prefill throttling # is disabled in this case. self.prefill_capacity_bound = False + # The DP engine core drives this cadence from a counter shared across + # ranks; an engine without one never throttles, so apply the interval + # locally, timed from the last prefill-carrying step. Three regimes: + # 1. no DP -- local cadence applies; + # 2. dense DP rank -- `reconfigure_for_independent_dp_rank` already + # set data_parallel_size=1, so each rank gets its own cadence. + # Intended, but a behaviour change: this used to be a no-op; + # 3. MoE DP rank (size > 1) -- the engine core drives it, no-op here. + self.local_prefill_interval = ( + self.scheduler_config.prefill_schedule_interval + if self.parallel_config.data_parallel_size == 1 + else 1 + ) + if self.local_prefill_interval > 1: + # Neither the hold nor its release is otherwise visible in logs or + # stats, which makes a misapplied cadence hard to diagnose. + logger.info_once( + "Prefill cadence: applying prefill_schedule_interval=%d " + "locally (this engine is not DP-coordinated); prefill is held " + "until %d executed steps have passed since the last " + "prefill-carrying step.", + self.local_prefill_interval, + self.local_prefill_interval, + ) + # A full interval back, so the first step is never gated. + self.last_prefill_step = -self.local_prefill_interval + # Steps that actually scheduled tokens. `current_step` counts every + # schedule() call including empty ones, so it would expire the hold + # early; it keeps its own meaning (next_decode_eligible_step). + self.num_executed_steps = 0 self.scheduler_reserve_full_isl = ( self.scheduler_config.scheduler_reserve_full_isl ) @@ -623,6 +653,9 @@ class Scheduler(SchedulerInterface): scheduled_spec_decode_tokens: dict[str, list[int]] = {} # Whether the running batch contains any prefill requests. prefill_scheduled = False + # Whether a waiting request was admitted with prompt left to compute. + # Separate from `prefill_scheduled`, which covers only the running loop. + prefill_admitted = False # Whether any scheduled request has a synchronous connector KV load. has_sync_kv_loads = False @@ -654,6 +687,28 @@ class Scheduler(SchedulerInterface): long_prefill_token_threshold, input_budget // num_eligible_reqs ) + # Local prefill cadence: hold prefill until the interval since the last + # prefill-carrying step has elapsed, so decode runs in longer + # uninterrupted stretches. Mirrors the DP gate above, including its + # capacity escape: once a release step cannot drain the waiting queue, + # throttling stops (#44558, test_throttle_capacity_bound_guard_admits). + # Without that escape a backed-up engine holds both admissions and + # in-flight chunks for interval-1 of every interval steps, capping + # prefill at roughly one budget per interval and growing the waiting + # queue without bound. + if ( + not defer_prefills + and self.local_prefill_interval > 1 + and not self.prefill_capacity_bound + # This call becomes executed step num_executed_steps + 1 if it + # schedules anything, and last_prefill_step is stamped on the same + # convention below. + and (self.num_executed_steps + 1) - self.last_prefill_step + < self.local_prefill_interval + and any(not r.is_prefill_chunk for r in self.running) + ): + defer_prefills = True + # First, schedule the RUNNING requests. req_index = 0 while req_index < len(self.running) and token_budget > 0: @@ -1318,6 +1373,12 @@ class Scheduler(SchedulerInterface): ) if load_kv_async: + # This path deliberately does not reach `prefill_admitted` + # below, so it does not restart the prefill cadence: + # `num_new_tokens` was set to 0 above, so the step carries + # no prefill compute for this request. The interval is + # timed from steps that actually prefilled; refreshing it + # here would start the cadence from a step that did none. # If loading async, allocate memory and put request # into the WAITING_FOR_REMOTE_KV state. request.status = RequestStatus.WAITING_FOR_REMOTE_KVS @@ -1373,6 +1434,14 @@ class Scheduler(SchedulerInterface): scheduled_spec_decode_tokens[request_id] = [ -1 ] * self.num_spec_tokens + # Must use the SAME predicate as the deferral gate in the + # waiting loop (`num_computed_tokens < num_tokens - 1`). + # Stamping on `< num_prompt_tokens` instead let a request with + # exactly one token left pass the gate on a held step and still + # restart the interval -- which is every full prefix-cache hit, + # and every remote-KV resume, since those set + # num_computed_tokens = num_tokens - 1. + prefill_admitted |= num_computed_tokens < request.num_tokens - 1 # Only track requests that will still be prefilling after this chunk. if num_computed_tokens + num_new_tokens < request.num_tokens: self._inflight_prefills.add(request) @@ -1403,6 +1472,15 @@ class Scheduler(SchedulerInterface): # Check if the scheduling constraints are satisfied. total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) assert total_num_scheduled_tokens <= self.max_num_scheduled_tokens + # Count this step before stamping the cadence below, so that + # last_prefill_step includes the step that carried the prefill. + if total_num_scheduled_tokens > 0: + self.num_executed_steps += 1 + + # Restart the cadence interval from any step that carried prefill, + # whether a new admission or a running chunk. + if self.local_prefill_interval > 1 and (prefill_scheduled or prefill_admitted): + self.last_prefill_step = self.num_executed_steps assert token_budget >= 0 assert input_budget >= 0 @@ -2496,7 +2574,6 @@ class Scheduler(SchedulerInterface): def _update_request_with_output( self, request: Request, new_token_ids: list[int], is_stale: bool = False ) -> tuple[list[int], bool]: - # is_stale is only used by the AsyncScheduler override. # Append generated tokens and check for stop. Note that if # a request is still being prefilled, we expect the model runner # to return empty token ids for the request. @@ -2508,6 +2585,10 @@ class Scheduler(SchedulerInterface): # This must be called before we make the EngineCoreOutput. stopped = check_stop(request, self.max_model_len) if stopped: + # num_computed_tokens covers all new tokens but the last one. + # Keep that true after trimming: resumable requests resume from it. + if not is_stale: + request.num_computed_tokens -= len(new_token_ids) - num_new del new_token_ids[num_new:] # Trim new tokens if needed. break return new_token_ids, stopped diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index e3c3ad5a24..bb6e01e7e4 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -16,6 +16,7 @@ from vllm.v1.core.kv_cache_utils import ( BlockHashWithGroupId, KVCacheBlock, resolve_block_hashes, + trim_null_head, ) from vllm.v1.hisparse.block_pool import SharedEventQueueBlockPool from vllm.v1.kv_cache_interface import ( @@ -569,14 +570,15 @@ class SingleTypeKVCacheManager(ABC): request_id: The request ID. Returns: - The request's blocks in allocation order. + The request's blocks in allocation order, without the leading + null blocks (e.g. a sliding-window table's skipped prefix). """ # Default to [] in case a request is freed (aborted) before alloc. req_blocks = self.req_to_blocks.pop(request_id, []) self.num_cached_block.pop(request_id, None) self._partial_hit_reqs.pop(request_id, None) - return req_blocks + return trim_null_head(req_blocks) def free(self, request_id: str) -> None: """Free the blocks for the request. @@ -2547,7 +2549,7 @@ class HiSparseResidentManager(_HiSparseAuxiliaryManager): blocks = self.get_residency_row(request_id) self.residency_changes.pop(request_id, None) super().pop_blocks_for_free(request_id) - return blocks + return trim_null_head(blocks) def adopt_resident_page( self, request_id: str, block_idx: int, block: KVCacheBlock diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index 1b14b08246..38aee44720 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -104,6 +104,9 @@ class EngineCoreReadyResponse: supports_draft_weight_updates: bool = False # Full-attention block size in tokens after initialization, or unavailable. effective_attention_block_size: int | None = None + coord_store_port: int = 0 + coordinator_input_address: str | None = None + coordinator_output_address: str | None = None class EngineCoreRequest( diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index f2142af2b3..8b7f913b43 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -22,7 +22,11 @@ from vllm.distributed.weight_transfer.base import ( ) from vllm.engine.arg_utils import AsyncEngineArgs from vllm.engine.protocol import EngineClient, StreamingInput -from vllm.entrypoints.serve.elastic_ep.middleware import set_scaling_elastic_ep +from vllm.entrypoints.serve.elastic_ep.middleware import ( + get_elastic_ep_rank_retired, + set_elastic_ep_rank_retired, + set_scaling_elastic_ep, +) from vllm.exceptions import ( GracefulHTTPError, MaxQueuedTokensError, @@ -53,7 +57,11 @@ from vllm.v1.engine.admission_control import SharedAdmissionStats from vllm.v1.engine.core_client import EngineCoreClient from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm.v1.engine.input_processor import InputProcessor -from vllm.v1.engine.output_processor import OutputProcessor, RequestOutputCollector +from vllm.v1.engine.output_processor import ( + OutputProcessor, + RequestOutputCollector, + StreamingInputOutputCollector, +) from vllm.v1.engine.parallel_sampling import ParentRequest from vllm.v1.executor import Executor from vllm.v1.fault_tolerance.utils import FaultToleranceRequest, FaultToleranceResult @@ -620,7 +628,9 @@ class AsyncLLM(EngineClient): self.input_processor.assign_request_id(final_req) internal_req_id = final_req.request_id - queue = RequestOutputCollector(sampling_params.output_kind, internal_req_id) + queue = StreamingInputOutputCollector( + sampling_params.output_kind, internal_req_id + ) async def handle_inputs(): cancelled = False @@ -628,6 +638,13 @@ class AsyncLLM(EngineClient): any_added = False try: async for input_chunk in input_stream: + # Stop once the session has finished (aborted or errored). + # There is no await between this check and the ADD send, + # so any chunk already sent precedes the engine abort. + if any_added and not self.output_processor.has_request( + internal_req_id + ): + break sp = input_chunk.sampling_params if sp: self._validate_streaming_input_sampling_params(sp) @@ -663,8 +680,10 @@ class AsyncLLM(EngineClient): if not cancelled: if any_added: # Send empty final request to indicate that inputs have - # finished. Don't send if cancelled (session was aborted). - await self._add_request(final_req, None, None, 0, queue) + # finished. Don't send if cancelled or if the session + # already finished (aborted or errored). + if self.output_processor.has_request(internal_req_id): + await self._add_request(final_req, None, None, 0, queue) elif not errored: queue.put(STREAM_FINISHED) @@ -903,6 +922,12 @@ class AsyncLLM(EngineClient): mm_cache_stats=renderer.stat_mm_cache(), ) except Exception as e: + if get_elastic_ep_rank_retired() and isinstance(e, EngineDeadError): + logger.info( + "[Elastic EP] Retired rank stopped receiving EngineCore outputs" + ) + output_processor.propagate_error(e) + return logger.exception("AsyncLLM output_handler failed.") output_processor.propagate_error(e) @@ -1088,6 +1113,8 @@ class AsyncLLM(EngineClient): async def check_health(self) -> None: logger.debug("Called check_health.") + if get_elastic_ep_rank_retired(): + return if self.errored: raise self.dead_error @@ -1237,8 +1264,13 @@ class AsyncLLM(EngineClient): "waiting for requests to drain before scaling" ) await self.wait_for_requests_to_drain(drain_timeout) - except BaseException: - set_scaling_elastic_ep(False) + except BaseException as e: + if self.vllm_config.parallel_config.data_parallel_external_lb: + await self.engine_core.abort_external_elastic_ep( + str(e) or type(e).__name__ + ) + else: + set_scaling_elastic_ep(False) raise async def scale_elastic_ep( @@ -1251,7 +1283,8 @@ class AsyncLLM(EngineClient): async def _scale_elastic_ep( self, new_data_parallel_size: int, drain_timeout: int ) -> None: - old_data_parallel_size = self.vllm_config.parallel_config.data_parallel_size + parallel_config = self.vllm_config.parallel_config + old_data_parallel_size = parallel_config.data_parallel_size if old_data_parallel_size == new_data_parallel_size: logger.info( "Data parallel size is already %s, skipping scale", @@ -1282,9 +1315,36 @@ class AsyncLLM(EngineClient): if envs.VLLM_ELASTIC_EP_DRAIN_REQUESTS: await self._drain_requests_for_elastic_ep(drain_timeout) - await self.engine_core.commit_elastic_ep() - self.vllm_config.parallel_config.data_parallel_size = new_data_parallel_size - set_scaling_elastic_ep(False) + rank_will_retire = ( + parallel_config.data_parallel_external_lb + and new_data_parallel_size < old_data_parallel_size + and parallel_config.data_parallel_rank >= new_data_parallel_size + ) + if rank_will_retire: + set_elastic_ep_rank_retired(True) + + commit_succeeded = False + try: + await self.engine_core.commit_elastic_ep() + parallel_config.data_parallel_size = new_data_parallel_size + commit_succeeded = True + if rank_will_retire: + logger.info( + "[Elastic EP] Data-parallel rank %s is retired and awaiting " + "external process shutdown", + parallel_config.data_parallel_rank, + ) + finally: + if rank_will_retire and not commit_succeeded: + set_elastic_ep_rank_retired(False) + # A failed commit may leave engines partially reconfigured or paused. + if commit_succeeded: + set_scaling_elastic_ep(False) + + async def get_external_elastic_ep_status( + self, + ) -> dict[str, str | int | None] | None: + return await self.engine_core.get_external_elastic_ep_status() async def handle_fault( self, fault_tolerance_request: FaultToleranceRequest @@ -1308,6 +1368,10 @@ class AsyncLLM(EngineClient): def errored(self) -> bool: return self.engine_core.resources.engine_dead or not self.is_running + @property + def should_keep_api_server_alive(self) -> bool: + return get_elastic_ep_rank_retired() + @property def dead_error(self) -> BaseException: return EngineDeadError() diff --git a/vllm/v1/engine/coordinator.py b/vllm/v1/engine/coordinator.py index cb120ba5d9..28a4c7af6d 100644 --- a/vllm/v1/engine/coordinator.py +++ b/vllm/v1/engine/coordinator.py @@ -126,6 +126,7 @@ class DPCoordinator: self.stats_publish_address = front_publish_address self.coord_in_address = back_publish_address self.coord_out_address = back_output_address + self._coord_store: object | None = None self._finalizer = weakref.finalize(self, shutdown, [self.proc]) def get_stats_publish_address(self) -> str: diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index d8c976ff69..c9843dce96 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -1346,6 +1346,16 @@ class EngineCoreProc(EngineCore): ready_msg["parallel_config_hash"] = ( vllm_config.parallel_config.compute_hash() ) + if vllm_config.parallel_config.enable_elastic_ep: + ready_msg["coord_store_port"] = ( + vllm_config.parallel_config._coord_store_port + ) + ready_msg["elastic_ep_max_dp_size"] = ( + vllm_config.parallel_config.elastic_ep_max_dp_size + ) + ready_msg["num_redundant_experts"] = ( + vllm_config.parallel_config.eplb_config.num_redundant_experts + ) handshake_socket.send(msgspec.msgpack.encode(ready_msg)) @@ -1376,14 +1386,37 @@ class EngineCoreProc(EngineCore): f"minutes" ) init_bytes = handshake_socket.recv() - init_message: EngineHandshakeMetadata = msgspec.msgpack.decode( - init_bytes, type=EngineHandshakeMetadata - ) + num_redundant_experts: int | None = None + if ( + envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH + and parallel_config is not None + and parallel_config.data_parallel_external_lb + ): + from vllm.distributed.elastic_ep.external_elastic_ep import ( + ExternalElasticEPScaleUpHandshakeMetadata, + ) + + scale_up_message = msgspec.msgpack.decode( + init_bytes, type=ExternalElasticEPScaleUpHandshakeMetadata + ) + init_message = scale_up_message.engine_metadata + num_redundant_experts = scale_up_message.num_redundant_experts + parallel_config.elastic_ep_max_dp_size = ( + scale_up_message.elastic_ep_max_dp_size + ) + else: + init_message = msgspec.msgpack.decode( + init_bytes, type=EngineHandshakeMetadata + ) logger.debug("Received init message: %s", init_message) if parallel_config is not None: for key, value in init_message.parallel_config.items(): setattr(parallel_config, key, value) + if num_redundant_experts is not None: + parallel_config.eplb_config.num_redundant_experts = ( + num_redundant_experts + ) return init_message.addresses @@ -1826,6 +1859,9 @@ class EngineCoreProc(EngineCore): if self.vllm_config.weight_transfer_config is not None else False ), + coord_store_port=parallel_config._coord_store_port, + coordinator_input_address=self.addresses.coordinator_input, + coordinator_output_address=self.addresses.coordinator_output, ) def process_input_sockets( @@ -2342,6 +2378,13 @@ class DPEngineCoreProc(EngineCoreProc): if state.is_complete(): if state.worker_type == "removing": raise SystemExit + if ( + state.worker_type == "new" + and self.vllm_config.parallel_config.data_parallel_external_lb + ): + # An independently launched rank has no external utility + # caller to join the post-commit resume barrier. + self.resume_scheduler() self.process_input_queue_block = True self.eep_scaling_state = None elif not state.commit_requested and state.is_ready_for_switch(): @@ -2512,12 +2555,10 @@ class DPEngineCoreProc(EngineCoreProc): self.output_queue.put_nowait((0, outputs)) else: encoder = MsgpackEncoder() - with ( - zmq.Context() as ctx, - make_zmq_socket( - ctx, self.addresses.outputs[0], zmq.PUSH, linger=4000 - ) as socket, - ): + ctx = zmq.Context.instance() + with make_zmq_socket( + ctx, self.addresses.outputs[0], zmq.PUSH, linger=4000 + ) as socket: socket.send_multipart(encoder.encode(outputs)) def _eep_scale_up_before_kv_init(self): diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index dccc8cba53..cfe8976668 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -23,6 +23,9 @@ import zmq.asyncio from vllm import envs from vllm.config import VllmConfig from vllm.config.kv_events import KVEventsConfig +from vllm.distributed.elastic_ep.external_elastic_ep import ( + ExternalElasticEPScaleCoordinator, +) from vllm.envs import VLLM_ENGINE_READY_TIMEOUT_S from vllm.logger import init_logger from vllm.lora.request import LoRARequest @@ -271,6 +274,14 @@ class EngineCoreClient(ABC): async def prepare_elastic_ep(self, new_data_parallel_size: int) -> None: raise NotImplementedError + async def abort_external_elastic_ep(self, error: str) -> None: + raise NotImplementedError + + async def get_external_elastic_ep_status( + self, + ) -> dict[str, str | int | None] | None: + return None + async def get_output_async(self) -> EngineCoreOutputs: raise NotImplementedError @@ -626,6 +637,8 @@ class MPClient(EngineCoreClient): enable_input_socket_handover = parallel_config.enable_elastic_ep self.stats_update_address: str | None = None + self.coordinator_input_address: str | None = None + self.coordinator_output_address: str | None = None tensor_queue: Queue | None = None if client_addresses: # Engines are managed externally to this client. @@ -702,6 +715,11 @@ class MPClient(EngineCoreClient): coordinator.get_stats_publish_address() ) + self.input_address = self.input_socket.getsockopt_string(zmq.LAST_ENDPOINT) + output_socket = self.resources.output_socket + assert output_socket is not None + self.output_address = output_socket.getsockopt_string(zmq.LAST_ENDPOINT) + # Serialization setup with tensor queues for multimodal tensor IPC. tensor_ipc_sender: TensorIpcSender | None = None model_config = getattr(vllm_config, "model_config", None) @@ -885,6 +903,25 @@ class MPClient(EngineCoreClient): response.kv_events_config ) + parallel_config = vllm_config.parallel_config + if response.coord_store_port: + if parallel_config._coord_store_port: + assert response.coord_store_port == parallel_config._coord_store_port + else: + parallel_config._coord_store_port = response.coord_store_port + + for attribute in ( + "coordinator_input_address", + "coordinator_output_address", + ): + address = getattr(response, attribute) + if address is not None: + existing_address = getattr(self, attribute) + if existing_address is None: + setattr(self, attribute, address) + else: + assert address == existing_address + def get_kv_event_sources(self) -> dict[int, KVEventsConfig]: return dict(self._kv_event_sources) @@ -1431,6 +1468,12 @@ class DPAsyncMPClient(AsyncMPClient): renderer: BaseRenderer | None = None, ): self.current_wave = 0 + parallel_config = vllm_config.parallel_config + self.external_eep_coordinator: ExternalElasticEPScaleCoordinator | None = None + if parallel_config.enable_elastic_ep and ( + parallel_config.data_parallel_external_lb + ): + self.external_eep_coordinator = ExternalElasticEPScaleCoordinator(self) super().__init__( vllm_config, @@ -1514,23 +1557,19 @@ class DPAsyncMPClient(AsyncMPClient): dp_rank = parallel_config.data_parallel_rank assert dp_rank == 0 assert dp_size == new_engine_count - assert not ( - parallel_config.data_parallel_hybrid_lb - or parallel_config.data_parallel_external_lb - ) - num_ranks = dp_size - self.engine_ranks_managed = list( - range(dp_rank, dp_rank + num_ranks) - ) - if len(self.lb_engines) < new_engine_count: - self.lb_engines = self.lb_engines + [ - [0, 0, 0.0] - for _ in range( - new_engine_count - len(self.lb_engines) - ) - ] - else: - self.lb_engines = self.lb_engines[:new_engine_count] + assert not parallel_config.data_parallel_hybrid_lb + # External LB clients continue to manage only local ranks. + if not parallel_config.data_parallel_external_lb: + self.engine_ranks_managed = list(range(dp_size)) + if len(self.lb_engines) < new_engine_count: + self.lb_engines += [ + [0, 0, 0.0] + for _ in range( + new_engine_count - len(self.lb_engines) + ) + ] + else: + self.lb_engines = self.lb_engines[:new_engine_count] # Send scale up notification to coordinator scale_msg = msgspec.msgpack.encode( ("SCALE_ELASTIC_EP", new_engine_count) @@ -1599,6 +1638,93 @@ class DPAsyncMPClient(AsyncMPClient): def get_core_engine_for_request(self, request: EngineCoreRequest): return self.core_engine + def _eep_commit_pause_mode(self) -> PauseMode: + from vllm.distributed.elastic_ep.elastic_execute import ( + can_reuse_fused_moe_kernel, + ) + + # MRV2 re-warms through the request pool, so running requests must + # finish first while new ones stay queued in the scheduler. + parallel_config = self.vllm_config.parallel_config + if self.vllm_config.use_v2_model_runner and not can_reuse_fused_moe_kernel( + parallel_config + ): + return "wait" + return "keep" + + def _setup_elastic_ep_reconfig_bootstrap(self) -> tuple[str, int]: + from vllm.distributed.utils import create_tcp_store + from vllm.utils.network_utils import get_open_ports_list + + parallel_config = self.vllm_config.parallel_config + parallel_config._data_parallel_master_port_list = get_open_ports_list(5) + parallel_config.data_parallel_master_port = ( + parallel_config._data_parallel_master_port_list.pop() + ) + + ip = parallel_config.data_parallel_master_ip + store = create_tcp_store( + ip, + 0, + is_master=True, + world_size=-1, + wait_for_workers=False, + ) + parallel_config._coord_store_port = store.port + self._coord_store = store + return ip, store.port + + @staticmethod + async def eep_process_engine_core_notification( + self: "DPAsyncMPClient", notification_data: tuple[str, int] + ) -> None: + if self.external_eep_coordinator is not None: + await self.external_eep_coordinator.process_engine_core_notification( + notification_data + ) + + def _get_external_eep_coordinator(self) -> ExternalElasticEPScaleCoordinator: + parallel_config = self.vllm_config.parallel_config + if not parallel_config.enable_elastic_ep: + raise NotImplementedError( + "Elastic EP scaling requires enable_elastic_ep=True." + ) + if not parallel_config.data_parallel_external_lb: + raise NotImplementedError( + "DPAsyncMPClient only supports Elastic EP scaling in external " + "load-balancer mode." + ) + coordinator = self.external_eep_coordinator + assert coordinator is not None + return coordinator + + async def prepare_elastic_ep(self, new_data_parallel_size: int) -> None: + cur_data_parallel_size = self.vllm_config.parallel_config.data_parallel_size + + assert new_data_parallel_size != cur_data_parallel_size, ( + f"new_data_parallel_size {new_data_parallel_size} must be " + f"different from cur_data_parallel_size {cur_data_parallel_size}" + ) + await self._get_external_eep_coordinator().prepare( + cur_data_parallel_size, new_data_parallel_size + ) + + async def commit_elastic_ep(self) -> None: + await self._get_external_eep_coordinator().commit() + + async def abort_external_elastic_ep(self, error: str) -> None: + self._get_external_eep_coordinator().abort(error) + + async def get_external_elastic_ep_status( + self, + ) -> dict[str, str | int | None] | None: + coordinator = self.external_eep_coordinator + return ( + None + if coordinator is None + else await asyncio.to_thread(coordinator.get_status) + ) + class DPLBAsyncMPClient(DPAsyncMPClient): """Asyncio-compatible client for multi-proc, multi-engine (data parallel) @@ -1618,6 +1744,8 @@ class DPLBAsyncMPClient(DPAsyncMPClient): # To route aborts to the correct engine. self.reqs_in_flight: dict[str, EngineIdentity] = {} + # Open streaming-input sessions, whose aborts go to all engines. + self.streaming_reqs: set[str] = set() # Exact per-engine count of this client's unfinished requests. self.engine_inflight: Counter[EngineIdentity] = Counter() @@ -1640,6 +1768,12 @@ class DPLBAsyncMPClient(DPAsyncMPClient): ) // client_count def get_core_engine_for_request(self, request: EngineCoreRequest) -> EngineIdentity: + if request.resumable: + self.streaming_reqs.add(request.request_id) + else: + # Also removes a session on its final (end-of-stream) request. + self.streaming_reqs.discard(request.request_id) + if (engine := self.reqs_in_flight.get(request.request_id)) is not None: # This is for streaming-input session affinity. return engine @@ -1711,8 +1845,9 @@ class DPLBAsyncMPClient(DPAsyncMPClient): @staticmethod async def eep_process_engine_core_notification( - self: "DPLBAsyncMPClient", notification_data: tuple[str, int] - ): + self: "DPAsyncMPClient", notification_data: tuple[str, int] + ) -> None: + assert isinstance(self, DPLBAsyncMPClient) cache = self.eep_scaling_cache notification_type_str, dp_rank = notification_data try: @@ -1759,15 +1894,24 @@ class DPLBAsyncMPClient(DPAsyncMPClient): if not request_ids or self.resources.engine_dead: return - if len(request_ids) == 1: + if ( + len(request_ids) == 1 + and (req_id := request_ids[0]) not in self.streaming_reqs + and (engine := self.reqs_in_flight.get(req_id)) + ): # Fast-path common case. - if engine := self.reqs_in_flight.get(request_ids[0]): - await self._abort_requests(request_ids, engine) + await self._abort_requests(request_ids, engine) return by_engine = defaultdict[EngineIdentity, list[str]](list) for req_id in request_ids: - if engine := self.reqs_in_flight.get(req_id): + if req_id in self.streaming_reqs: + # A chunk in flight when the session finished may have + # re-created it on any engine. + self.streaming_reqs.discard(req_id) + for engine in self.core_engines: + by_engine[engine].append(req_id) + elif engine := self.reqs_in_flight.get(req_id): by_engine[engine].append(req_id) for engine, req_ids in by_engine.items(): await self._abort_requests(req_ids, engine) @@ -1864,27 +2008,6 @@ class DPLBAsyncMPClient(DPAsyncMPClient): new_engine_identities.discard(identity) self._apply_ready_response(payload) - def _setup_elastic_ep_reconfig_bootstrap(self) -> None: - from vllm.distributed.utils import create_tcp_store - from vllm.utils.network_utils import get_open_ports_list - - parallel_config = self.vllm_config.parallel_config - parallel_config._data_parallel_master_port_list = get_open_ports_list(5) - parallel_config.data_parallel_master_port = ( - parallel_config._data_parallel_master_port_list.pop() - ) - - ip = parallel_config.data_parallel_master_ip - store = create_tcp_store( - ip, - 0, - is_master=True, - world_size=-1, - wait_for_workers=False, - ) - parallel_config._coord_store_port = store.port - self._coord_store = store - def _make_reconfig_request( self, new_data_parallel_size: int, @@ -1939,20 +2062,6 @@ class DPLBAsyncMPClient(DPAsyncMPClient): await asyncio.to_thread(self._coord_store.wait, ready_keys) logger.info("[Elastic EP] Successfully started new engines") - def _eep_commit_pause_mode(self) -> PauseMode: - from vllm.distributed.elastic_ep.elastic_execute import ( - can_reuse_fused_moe_kernel, - ) - - # MRV2 re-warms through the request pool, so running requests must - # finish first while new ones stay queued in the scheduler. - parallel_config = self.vllm_config.parallel_config - if self.vllm_config.use_v2_model_runner and not can_reuse_fused_moe_kernel( - parallel_config - ): - return "wait" - return "keep" - async def _commit_scale_up_elastic_ep(self, new_data_parallel_size: int) -> None: new_core_engines = [ rank.to_bytes(2, "little") diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index eb499f6cc7..6008dd0d12 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -113,6 +113,53 @@ class RequestOutputCollector: self._input_stream_task = None +def _is_chunk_boundary( + output: RequestOutput | PoolingRequestOutput | Exception | None, +) -> bool: + """Whether the output marks the end of a streaming-input chunk: a + finish_reason on a request that has not finished.""" + return ( + isinstance(output, RequestOutput) + and not output.finished + and any(completion.finished() for completion in output.outputs) + ) + + +class StreamingInputOutputCollector(RequestOutputCollector): + """Collector for streaming-input requests. Doesn't merge across input + chunk boundaries (a finish_reason on a non-finished output).""" + + def __init__(self, output_kind: RequestOutputKind, request_id: str): + super().__init__(output_kind, request_id) + self._backlog: deque[RequestOutput] = deque() + + def put(self, output: RequestOutput | PoolingRequestOutput | Exception) -> None: + if isinstance(output, Exception): + self._backlog.clear() + elif self._backlog or _is_chunk_boundary(self.output): + assert isinstance(output, RequestOutput) + if self._backlog and not _is_chunk_boundary(self._backlog[-1]): + self._backlog[-1].add(output, aggregate=self.aggregate) + else: + self._backlog.append(output) + return + super().put(output) + + async def get(self) -> RequestOutput | PoolingRequestOutput: + output = await super().get() + if self._backlog: + self.output = self._backlog.popleft() + self.ready.set() + return output + + def get_nowait(self) -> RequestOutput | PoolingRequestOutput | None: + output = super().get_nowait() + if self._backlog: + self.output = self._backlog.popleft() + self.ready.set() + return output + + @dataclass class OutputProcessorOutput: request_outputs: list[RequestOutput | PoolingRequestOutput] @@ -784,6 +831,10 @@ class OutputProcessor: # if required. req_state.logprobs_processor.update_from_output(engine_core_output) + streaming_chunk_boundary = req_state.streaming_input and ( + finish_reason not in (FinishReason.ABORT, FinishReason.ERROR) + ) + # 4) Create and handle RequestOutput objects. if request_output := req_state.make_request_output( new_token_ids, @@ -794,7 +845,7 @@ class OutputProcessor: ec_transfer_params, error=request_error, ): - if req_state.streaming_input: + if streaming_chunk_boundary: request_output.finished = False if req_state.queue is not None: @@ -806,7 +857,7 @@ class OutputProcessor: # Free completed requests. if finish_reason is not None: - if req_state.streaming_input: + if streaming_chunk_boundary: if req_state.input_chunk_queue: update = req_state.input_chunk_queue.popleft() req_state.apply_streaming_update(update) @@ -814,9 +865,10 @@ class OutputProcessor: req_state.input_chunk_queue = None else: self._finish_request(req_state) - if not engine_core_output.finished: - # If req not finished in EngineCore, but Detokenizer - # detected stop string, abort needed in EngineCore. + if req_state.streaming_input or not engine_core_output.finished: + # EngineCore may still have the request: a stop string + # was detected here, or an in-flight streaming-input + # chunk may re-create the session. reqs_to_abort.append(req_id) # Track per-request stats diff --git a/vllm/v1/engine/utils.py b/vllm/v1/engine/utils.py index 4f1be3bacd..96df5f14eb 100644 --- a/vllm/v1/engine/utils.py +++ b/vllm/v1/engine/utils.py @@ -290,6 +290,7 @@ class CoreEngineProcManager: self._finalizer = weakref.finalize(self, shutdown, self.processes) self.manager_stopped = threading.Event() self.failed_proc_name: str | None = None + self._coord_store: object | None = None # All ranks share this config object: capture the user-provided # --device-ids list before the per-rank shard overwrites it. Mutating @@ -1139,7 +1140,9 @@ def get_engine_zmq_bind_addresses( offline_mode or local_engines_only or (local_engine_count == dp_size) ) # NOTE(yongji): handling scaling from intra-node to inter-node - if parallel_config.enable_elastic_ep: + if parallel_config.enable_elastic_ep and ( + not local_engines_only or parallel_config.data_parallel_rank == 0 + ): client_local_only = False def _addr() -> str: @@ -1228,6 +1231,32 @@ def launch_core_engines( vllm_config.needs_dp_coordinator and not offline_mode and dp_rank == 0 ) + coord_store = None + + create_coord_store = ( + parallel_config.enable_elastic_ep + and parallel_config.data_parallel_backend != "ray" + and dp_rank == 0 + and not parallel_config._coord_store_port + ) + + if create_coord_store: + from vllm.distributed.utils import create_tcp_store + + coord_store = create_tcp_store( + host, + 0, + is_master=True, + world_size=-1, + wait_for_workers=False, + ) + parallel_config._coord_store_port = coord_store.port + logger.info( + "Created Elastic EP coordination TCPStore on %s:%d", + host, + coord_store.port, + ) + if run_coordinator: coordinator = DPCoordinator( parallel_config, @@ -1243,12 +1272,13 @@ def launch_core_engines( ) logger.info("Started DP Coordinator process (PID: %d)", coordinator.proc.pid) + if coord_store is not None: + coordinator._coord_store = coord_store else: coordinator = None # Hold a coordination TCPStore (alive for this frame) so engines pick DP # master ports at bind time; pre-allocated ports can be taken before use. - coord_store = None if ( dp_size > 1 and not offline_mode @@ -1344,6 +1374,13 @@ def launch_core_engines( else: local_engine_manager = None + if ( + coord_store is not None + and coordinator is None + and (local_engine_manager is not None) + ): + local_engine_manager._coord_store = coord_store + launch = CoreEngineLaunch( local_engine_manager, coordinator, addresses, tensor_queue ) @@ -1474,21 +1511,33 @@ def wait_for_engine_startup( if status == "HELLO" and engine.state == CoreEngineState.NEW: # Send init message with DP config info. + parallel_config_payload: dict[str, int | str | list[int]] = {} + if coordinated_dp and not ( + parallel_config.local_engines_only + and parallel_config.data_parallel_rank > 0 + ): + parallel_config_payload = { + k: getattr(parallel_config, k) + for k in ( + "data_parallel_master_ip", + "data_parallel_master_port", + "_data_parallel_master_port_list", + "_coord_store_port", + "data_parallel_size", + ) + } + if ( + parallel_config.enable_elastic_ep + and not parallel_config._coord_store_port + ): + raise RuntimeError( + "Elastic EP startup requires a coordination TCPStore " + "port before engine handshake." + ) init_message = msgspec.msgpack.encode( EngineHandshakeMetadata( addresses=launch.addresses, - parallel_config={ - k: getattr(parallel_config, k) - for k in ( - "data_parallel_master_ip", - "data_parallel_master_port", - "_data_parallel_master_port_list", - "_coord_store_port", - "data_parallel_size", - ) - } - if coordinated_dp - else {}, + parallel_config=parallel_config_payload, ) ) handshake_socket.send_multipart((eng_identity, init_message), copy=False) @@ -1496,6 +1545,18 @@ def wait_for_engine_startup( start_pending[0 if local else 1] += 1 engine.state = CoreEngineState.CONNECTED elif status == "READY" and engine.state == CoreEngineState.CONNECTED: + coord_store_port = msg.get("coord_store_port") + if coord_store_port is not None and local: + parallel_config._coord_store_port = coord_store_port + num_redundant_experts = msg.get("num_redundant_experts") + if num_redundant_experts is not None and local: + parallel_config.eplb_config.num_redundant_experts = ( + num_redundant_experts + ) + elastic_ep_max_dp_size = msg.get("elastic_ep_max_dp_size") + if elastic_ep_max_dp_size is not None and local: + parallel_config.elastic_ep_max_dp_size = elastic_ep_max_dp_size + # Validate config hash consistency across DP workers for MoE models. if coordinated_dp: worker_config_hash = msg.get("parallel_config_hash") diff --git a/vllm/v1/executor/ray_executor_v2.py b/vllm/v1/executor/ray_executor_v2.py index 8d254020ff..63c3ff0fbe 100644 --- a/vllm/v1/executor/ray_executor_v2.py +++ b/vllm/v1/executor/ray_executor_v2.py @@ -292,7 +292,7 @@ class RayExecutorV2(MultiprocExecutor): # Step 1: Initialize Ray cluster and retrieve placement group if ray is None: raise ImportError("Using Ray backend requires installation of ray.") - initialize_ray_cluster(self.parallel_config, require_gpu_on_driver=False) + initialize_ray_cluster(self.parallel_config) placement_group = self.parallel_config.placement_group tp_size, pp_size, pcp_size = self._get_parallel_sizes() diff --git a/vllm/v1/executor/ray_utils.py b/vllm/v1/executor/ray_utils.py index 48dba63936..44fa7a835a 100644 --- a/vllm/v1/executor/ray_utils.py +++ b/vllm/v1/executor/ray_utils.py @@ -8,7 +8,6 @@ from collections import defaultdict from vllm.config import ParallelConfig from vllm.logger import init_logger from vllm.platforms import current_platform -from vllm.utils.network_utils import get_ip logger = init_logger(__name__) PG_WAIT_TIMEOUT = 1800 @@ -30,14 +29,6 @@ try: from ray.util import placement_group_table from ray.util.placement_group import PlacementGroup - try: - from ray._private.state import available_resources_per_node - except ImportError: - # Ray 2.9.x doesn't expose `available_resources_per_node` - from ray._private.state import state as _state - - available_resources_per_node = _state._available_resources_per_node - ray_import_err = None except ImportError as e: @@ -65,15 +56,8 @@ def _verify_bundles( placement_group: "PlacementGroup", parallel_config: ParallelConfig, device_str: str, - require_gpu_on_driver: bool = True, ): - """Verify a given placement group has bundles located in the right place. - - There are 2 rules. - - Warn if all tensor parallel workers cannot fit in a single node. - - Fail if driver node is not included in a placement group - (only when require_gpu_on_driver is True). - """ + """Warn if all tensor parallel workers cannot fit in a single node.""" assert ray.is_initialized(), ( "Ray is not initialized although distributed-executor-backend is ray." ) @@ -87,17 +71,6 @@ def _verify_bundles( for bundle_idx, node_id in bundle_to_node_ids.items(): node_id_to_bundle[node_id].append(bundles[bundle_idx]) - driver_node_id = ray.get_runtime_context().get_node_id() - - if require_gpu_on_driver and driver_node_id not in node_id_to_bundle: - raise RuntimeError( - f"driver node id {driver_node_id} is not included in a placement " - f"group {placement_group.id}. Node id -> bundles " - f"{node_id_to_bundle}. " - "You don't have enough GPUs available in a current node. Check " - "`ray status` and `ray list nodes` to see if you have available " - "GPUs in a node `{driver_node_id}` before starting an vLLM engine." - ) for node_id, bundles in node_id_to_bundle.items(): if len(bundles) < parallel_config.tensor_parallel_size: @@ -293,7 +266,6 @@ def _wait_until_pg_ready(current_placement_group: "PlacementGroup"): def initialize_ray_cluster( parallel_config: ParallelConfig, ray_address: str | None = None, - require_gpu_on_driver: bool = True, ): """Initialize the distributed cluster with Ray. @@ -305,10 +277,6 @@ def initialize_ray_cluster( parallel_config: The configurations for parallel execution. ray_address: The address of the Ray cluster. If None, uses the default Ray cluster address. - require_gpu_on_driver: If True (default), require at least one GPU - on the current (driver) node and pin the first PG bundle to it. - Set to False for executors like RayExecutorV2 where all GPU work - is delegated to remote Ray actors. """ assert_ray_available() @@ -404,25 +372,6 @@ def initialize_ray_cluster( {device_str: 1.0} for _ in range(parallel_config.world_size) ] - # vLLM engine is also a worker to execute model with an accelerator, - # so it requires to have the device in a current node. Check if - # the current node has at least one device. - current_ip = get_ip() - current_node_id = ray.get_runtime_context().get_node_id() - current_node_resource = available_resources_per_node()[current_node_id] - if require_gpu_on_driver: - if current_node_resource.get(device_str, 0) < 1: - raise ValueError( - f"Current node has no {device_str} available. " - f"{current_node_resource=}. vLLM engine cannot start " - f"without {device_str}. Make sure you have at least 1 " - f"{device_str} available in a node " - f"{current_node_id=} {current_ip=}." - ) - # This way, at least bundle is required to be created in a - # current node. - placement_group_specs[0][f"node:{current_ip}"] = 0.001 - # By default, Ray packs resources as much as possible. current_placement_group = ray.util.placement_group( placement_group_specs, strategy="PACK" @@ -430,8 +379,6 @@ def initialize_ray_cluster( _wait_until_pg_ready(current_placement_group) assert current_placement_group is not None - _verify_bundles( - current_placement_group, parallel_config, device_str, require_gpu_on_driver - ) + _verify_bundles(current_placement_group, parallel_config, device_str) # Set the placement group in the parallel config parallel_config.placement_group = current_placement_group diff --git a/vllm/v1/hisparse/binding.py b/vllm/v1/hisparse/binding.py index 1420a45396..6a02dd571c 100644 --- a/vllm/v1/hisparse/binding.py +++ b/vllm/v1/hisparse/binding.py @@ -17,6 +17,8 @@ from vllm.v1.hisparse.layout import ( from vllm.v1.hisparse.runtime import ( HiSparseCacheHandle, HiSparseHostPool, + drop_hisparse_host_pool, + get_hisparse_host_pool, initialize_hisparse_runtime_buffers, release_pinned_state, ) @@ -126,7 +128,7 @@ def init_hisparse_kv_cache( block_tables: "BlockTables", ) -> dict[str, torch.Tensor]: """Allocate and bind HiSparse caches within the caller's allocation context.""" - host_pool = HiSparseHostPool(vllm_config, kv_cache_config) + host_pool = get_hisparse_host_pool(vllm_config, kv_cache_config) try: kv_caches = allocate_hisparse_kv_caches( kv_cache_config, @@ -150,6 +152,9 @@ def init_hisparse_kv_cache( except Exception: if host_pool.shared_region is not None: host_pool.shared_region.cleanup() + elif host_pool.registered is not None: + release_pinned_state([], [host_pool.registered]) + drop_hisparse_host_pool(vllm_config) raise @@ -163,37 +168,14 @@ def _get_hisparse_cache( def release_hisparse_profiling_cache(forward_context: dict[str, Any]) -> None: - cache_handles = [ - cache - for layer in forward_context.values() - if (cache := getattr(layer, "hisparse_cache", None)) is not None - ] + """Detach the profiling caches but keep the host pool for the real KV cache.""" runtimes = { id(cache.runtime): cache.runtime - for cache in cache_handles - if hasattr(cache.runtime, "_host_cache") + for layer in forward_context.values() + if (cache := getattr(layer, "hisparse_cache", None)) is not None + and hasattr(cache.runtime, "_host_cache") } - if not runtimes: - return - - shared_regions = { - id(runtime.shared_host_region): runtime.shared_host_region - for runtime in runtimes.values() - if runtime.shared_host_region is not None - } - assert len(shared_regions) <= 1 - shared_region = next(iter(shared_regions.values()), None) - registered_pools = ( - [] - if shared_region is not None - else list( - { - runtime.registered_host_pool.data_ptr(): runtime.registered_host_pool - for runtime in runtimes.values() - }.values() - ) - ) - release_pinned_state(list(runtimes.values()), registered_pools, shared_region) + release_pinned_state(list(runtimes.values()), []) def bind_hisparse_kv_caches( diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py index 3cc45c680e..87916ba4ad 100644 --- a/vllm/v1/hisparse/runtime.py +++ b/vllm/v1/hisparse/runtime.py @@ -223,7 +223,9 @@ class HiSparseHostPool: self.shared_region: SharedOffloadRegion | None = None def allocate(self, size: int) -> torch.Tensor: - assert self.backing is None, "HiSparse host pool is already allocated." + if self.backing is not None: + assert self.backing.numel() == size, "HiSparse host pool size changed." + return self.backing config = self.kv_cache_config if config.hisparse_shared_host_pool: tensor_sizes = [ @@ -253,6 +255,23 @@ class HiSparseHostPool: return self.backing +_host_pools: dict[str, HiSparseHostPool] = {} + + +def get_hisparse_host_pool( + vllm_config: VllmConfig, kv_cache_config: KVCacheConfig +) -> HiSparseHostPool: + """Return the engine's host pool, which outlives the profiling KV cache.""" + instance_id = vllm_config.instance_id + if instance_id not in _host_pools: + _host_pools[instance_id] = HiSparseHostPool(vllm_config, kv_cache_config) + return _host_pools[instance_id] + + +def drop_hisparse_host_pool(vllm_config: VllmConfig) -> None: + _host_pools.pop(vllm_config.instance_id, None) + + def use_shared_hisparse_host_pool(vllm_config: VllmConfig) -> bool: """Whether replicated MLA host KV can share one local mmap.""" parallel = vllm_config.parallel_config diff --git a/vllm/v1/kv_offload/cpu/shared_offload_region.py b/vllm/v1/kv_offload/cpu/shared_offload_region.py index 34faf28b94..3457a34ff0 100644 --- a/vllm/v1/kv_offload/cpu/shared_offload_region.py +++ b/vllm/v1/kv_offload/cpu/shared_offload_region.py @@ -1,11 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import contextlib +import ctypes import errno +import functools import mmap import os import time from collections.abc import Callable, Iterator +from ctypes import c_int, c_long, c_size_t, c_void_p import numpy as np import torch @@ -21,6 +24,10 @@ logger = init_logger(__name__) # MADV_POPULATE_WRITE was added in Linux 5.14 (value 23). _MADV_POPULATE_WRITE = getattr(mmap, "MADV_POPULATE_WRITE", 23) +_MAP_FIXED = 0x10 +# Page-cache inserts are serialized per shm file, capping pre-faulting of one +# file near 10 GiB/s; smaller segments are not worth a file. +_MIN_SEGMENT_BYTES = 1 << 30 def _wait_for_file_size(fd: int, expected_size: int, timeout: float = 30.0) -> None: @@ -40,6 +47,14 @@ def _wait_for_file_size(fd: int, expected_size: int, timeout: float = 30.0) -> N time.sleep(0.005) +@functools.cache +def _libc_mmap() -> Callable[..., int | None]: + mmap_fn = ctypes.CDLL(None, use_errno=True).mmap + mmap_fn.argtypes = [c_void_p, c_size_t, c_int, c_int, c_int, c_long] + mmap_fn.restype = c_void_p + return mmap_fn + + def _madvise_populate_write(mmap_obj: mmap.mmap, offset: int, length: int) -> None: mmap_obj.madvise(_MADV_POPULATE_WRITE, offset, length) @@ -84,6 +99,7 @@ class SharedOffloadRegion: ``populate_shard=(index, count)`` pre-faults only this worker's contiguous share, for layouts where every worker maps the same bytes. The barrier also waits for all shards to finish pre-faulting before returning. + ``num_segments`` splits backing files beneath the contiguous mapping. """ BLOCK_SIZE_ALIGNMENT: int = mmap.PAGESIZE @@ -99,6 +115,7 @@ class SharedOffloadRegion: *, creator_memory_check: Callable[[int], None] | None = None, populate_shard: tuple[int, int] | None = None, + num_segments: int = 1, ) -> None: self.page_size = mmap.PAGESIZE assert kv_bytes_per_chunk % self.page_size == 0 @@ -109,6 +126,17 @@ class SharedOffloadRegion: self.cpu_page_size = cpu_page_size self.mmap_path = f"/dev/shm/vllm_offload_{engine_id}.mmap" + num_segments = min(num_segments, self.total_size_bytes // _MIN_SEGMENT_BYTES) + rows = max(1, cdiv(num_chunks, max(1, num_segments))) + # (path, offset, length) per file. + self._segments = [ + ( + f"{self.mmap_path}.{i}" if i else self.mmap_path, + row * self._row_stride, + min(rows, num_chunks - row) * self._row_stride, + ) + for i, row in enumerate(range(0, max(num_chunks, 1), rows)) + ] self._creator = False # set True only if this worker creates the file self.rank = rank if rank is not None: @@ -126,7 +154,7 @@ class SharedOffloadRegion: # Joiner path — another worker won O_EXCL. Reopen and wait # for the file to reach expected size. self.fd = os.open(self.mmap_path, os.O_RDWR) - _wait_for_file_size(self.fd, self.total_size_bytes) + _wait_for_file_size(self.fd, self._segments[0][2]) logger.info("Opened existing mmap file %s", self.mmap_path) else: # Creator path. We won O_EXCL, so we own the file: any @@ -140,24 +168,32 @@ class SharedOffloadRegion: self.total_size_bytes, allocation_name="CPU KV offload shared region in /dev/shm", ) - os.ftruncate(self.fd, self.total_size_bytes) + for path, _, length in self._segments[1:]: + fd = os.open(path, os.O_CREAT | os.O_EXCL | os.O_RDWR, 0o600) + os.ftruncate(fd, length) + os.close(fd) + # Joiners wait for the first file, so size it last. + os.ftruncate(self.fd, self._segments[0][2]) logger.info( "Created mmap file %s (%.2f GB)", self.mmap_path, self.total_size_bytes / 1e9, ) + single = len(self._segments) == 1 self.mmap_obj: mmap.mmap | None = mmap.mmap( - self.fd, + self.fd if single else -1, self.total_size_bytes, - flags=mmap.MAP_SHARED, + flags=mmap.MAP_SHARED if single else mmap.MAP_PRIVATE, prot=mmap.PROT_READ | mmap.PROT_WRITE, ) + if not single: + self._map_segments() if barrier is not None: barriers_remaining -= 1 barrier() if self._creator: - os.unlink(self.mmap_path) + self._unlink() self._creator = False logger.info("Unlinked mmap file %s", self.mmap_path) if populate_shard is not None: @@ -173,8 +209,7 @@ class SharedOffloadRegion: barrier() except Exception: if self._creator: - with contextlib.suppress(FileNotFoundError): - os.unlink(self.mmap_path) + self._unlink() self._creator = False if hasattr(self, "mmap_obj") and self.mmap_obj is not None: # The prefault traceback may retain an exported buffer. @@ -225,11 +260,42 @@ class SharedOffloadRegion: "MADV_POPULATE_WRITE entire region: %.3f s", time.perf_counter() - _t0 ) + def _map_segments(self) -> None: + """Replace the anonymous mapping with the segment files, back to back.""" + assert self.mmap_obj is not None and self.fd is not None + base = ctypes.addressof(ctypes.c_char.from_buffer(self.mmap_obj)) + for path, offset, length in self._segments: + fd = self.fd if offset == 0 else os.open(path, os.O_RDWR) + address = _libc_mmap()( + base + offset, + length, + mmap.PROT_READ | mmap.PROT_WRITE, + mmap.MAP_SHARED | _MAP_FIXED, + fd, + 0, + ) + if fd != self.fd: + os.close(fd) + if address != base + offset: + raise OSError(ctypes.get_errno(), f"Failed to map {path}") + + def _unlink(self) -> None: + for path, _, _ in self._segments: + with contextlib.suppress(FileNotFoundError): + os.unlink(path) + def iter_worker_slot_ranges(self) -> Iterator[tuple[int, int]]: """Yield this worker's slots as page-aligned (offset, length) ranges.""" assert self.rank is not None worker_offset = self.rank * self.cpu_page_size - for row in range(self.num_chunks): + first = 0 + if len(self._segments) > 1: + # Start each worker on a different file. + first = ( + self.rank * self.num_chunks // (self._row_stride // self.cpu_page_size) + ) + for i in range(self.num_chunks): + row = (first + i) % self.num_chunks start = row * self._row_stride + worker_offset aligned = start - start % self.page_size end = round_up(start + self.cpu_page_size, self.page_size) @@ -380,7 +446,7 @@ class SharedOffloadRegion: self.fd = None if self._creator and getattr(self, "mmap_path", None): try: - os.unlink(self.mmap_path) + self._unlink() logger.info("Removed mmap file %s", self.mmap_path) except Exception: logger.warning( diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index 4f6275d028..807f1b6e8f 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -22,6 +22,10 @@ from vllm.v1.kv_offload.cpu.gpu_worker import CPUOffloadingWorker from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion +# Split the shared region into this many shm files so workers pre-fault them +# in parallel (page-cache inserts are serialized per file). +_REGION_SEGMENTS = 32 + def _all_workers_barrier() -> None: """Block until every worker rank has reached this point (gloo cpu group). @@ -188,6 +192,7 @@ class CPUOffloadingSpec(OffloadingSpec): kv_bytes_per_chunk=self.kv_bytes_per_chunk, cpu_page_size=self.cpu_page_size_per_worker, barrier=_all_workers_barrier, + num_segments=_REGION_SEGMENTS, populate_shard=( (worker_index, world_size) if self.replicated_layout else None ), diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index daa94ebc3a..d38497b2ff 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -34,11 +34,34 @@ def register_top_k_top_p_warmups() -> None: _topp_split_mask.register_warmup() -def _skip_aiter_sampler_on_gfx1250() -> bool: - # Lazy ROCm-only import; keeps arch detection out of import time on CUDA/CPU. - from vllm.platforms.rocm import on_gfx1250 +def _aiter_temp_gumbel_sample( + logits: torch.Tensor, generators: dict[int, torch.Generator] +) -> torch.Tensor: + """Fused temperature Gumbel-max sampling via aiter. - return on_gfx1250() + Logits are already temperature-scaled by the parent Sampler, so temps=1. + Returns int64 token ids of shape (num_tokens,). + """ + from aiter import mixed_sample_outer_exponential + + n, vocab = logits.shape + out = torch.empty(n, dtype=torch.int32, device=logits.device) + if not generators: + # Common case: shared (1, vocab) exponential noise (cheap RNG). + exp = ( + torch.empty((1, vocab), dtype=torch.float32, device=logits.device) + .exponential_() + .expand(n, vocab) + ) + else: + exp = torch.empty((n, vocab), dtype=torch.float32, device=logits.device) + if len(generators) != n: + exp.exponential_() + for i, g in generators.items(): + exp[i].exponential_(generator=g) + temps = torch.ones(n, dtype=torch.float32, device=logits.device) + mixed_sample_outer_exponential(out, logits, exp, temps, eps=1e-10) + return out.to(torch.int64) def _flashinfer_jit_unsupported_reason(capability: DeviceCapability) -> str | None: @@ -246,8 +269,12 @@ class TopKTopPSampler(nn.Module): elif ( logprobs_mode not in PROCESSED_LOGPROBS_MODES and rocm_aiter_ops.is_enabled() - and not _skip_aiter_sampler_on_gfx1250() # TODO (JPVILLAM): Enable ): + # Lazy ROCm-only import; keeps arch detection out of import time + # on CUDA/CPU. + from vllm.platforms.rocm import on_gfx1250 + + self.on_gfx1250 = on_gfx1250() self.aiter_ops = None self._aiter_ops_import_failed = False logger.info_once( @@ -375,7 +402,16 @@ class TopKTopPSampler(nn.Module): k: torch.Tensor | None, p: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor | None]: - """Optimized ROCm/aiter path (same structure as forward_cuda).""" + """Optimized ROCm/aiter path.""" + if self.on_gfx1250: + if ( + k is None + and p is None + and not self.use_fp64_gumbel + and self.logprobs_mode not in PROCESSED_LOGPROBS_MODES + ): + return _aiter_temp_gumbel_sample(logits, generators), None + return self.forward_native(logits, generators, k, p) if (k is None and p is None) or generators: if generators: logger.warning_once( diff --git a/vllm/v1/sample/sampler.py b/vllm/v1/sample/sampler.py index d19f65f6d8..a9cb2f1a23 100644 --- a/vllm/v1/sample/sampler.py +++ b/vllm/v1/sample/sampler.py @@ -5,6 +5,7 @@ import torch import torch.nn as nn +from vllm._aiter_ops import rocm_aiter_ops from vllm.config.model import LogprobsMode from vllm.platforms import current_platform from vllm.utils.gpu_sync_debug import gpu_sync_allowed @@ -19,6 +20,61 @@ from vllm.v1.sample.ops.topk_topp_sampler import TopKTopPSampler _SAMPLING_EPS = 1e-5 +def mixed_sample_eligible(sm: SamplingMetadata, predict_bonus_token: bool) -> bool: + """Whether the fused aiter sample-from-raw-logits kernel can be used.""" + if not rocm_aiter_ops.is_enabled(): + return False + + from vllm.platforms.rocm import on_gfx1250 + + if not on_gfx1250(): + return False + if predict_bonus_token: + return False + if not sm.all_random or sm.temperature is None: + return False + if sm.top_k is not None or sm.top_p is not None: + return False + if sm.max_num_logprobs is not None or sm.logprob_token_ids: + return False + if ( + not sm.no_penalties + or sm.allowed_token_ids_mask is not None + or sm.bad_words_token_ids + ): + return False + lp = sm.logitsprocs + if lp.argmax_invariant or lp.non_argmax_invariant: + return False + h = sm.thinking_budget_state_holder + return h is None or not h.has_tracked_requests() + + +def _aiter_raw_fused_sample(logits: torch.Tensor, sm: SamplingMetadata) -> torch.Tensor: + """Sample token ids from raw (unscaled) logits via a fused aiter kernel. + + Returns int32 ids of shape (num_tokens, 1). + """ + from aiter import mixed_sample_outer_exponential + + n, vocab = logits.shape + out = torch.empty(n, dtype=torch.int32, device=logits.device) + gens = sm.generators + if not gens: + exp = ( + torch.empty((1, vocab), dtype=torch.float32, device=logits.device) + .exponential_() + .expand(n, vocab) + ) + else: + exp = torch.empty((n, vocab), dtype=torch.float32, device=logits.device) + exp.exponential_() + for i, g in gens.items(): + exp[i].exponential_(generator=g) + mixed_sample_outer_exponential(out, logits, exp, sm.temperature, eps=1e-10) + return out.unsqueeze(-1) + + class Sampler(nn.Module): """A layer that samples the next tokens from the model's outputs with the following steps in order: @@ -77,6 +133,11 @@ class Sampler(nn.Module): predict_bonus_token: bool = False, logprobs_mode_override: LogprobsMode | None = None, ) -> SamplerOutput: + if mixed_sample_eligible(sampling_metadata, predict_bonus_token): + return SamplerOutput( + sampled_token_ids=_aiter_raw_fused_sample(logits, sampling_metadata), + logprobs_tensors=None, + ) logprobs_mode = logprobs_mode_override or self.logprobs_mode # NOTE(woosuk): Use the original logits (before any penalties or # temperature scaling) for the top-k logprobs. diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index 3f461ed47f..7de9dc2216 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -5,6 +5,8 @@ from collections.abc import Iterable, Sequence from concurrent.futures import Future, ThreadPoolExecutor from typing import TYPE_CHECKING, NamedTuple +import numpy as np + from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.parser.engine.adapters import ParserEngineReasoningAdapter @@ -20,7 +22,6 @@ from vllm.v1.structured_output.backend_xgrammar import XgrammarBackend from vllm.v1.structured_output.utils import strip_speculative_padding if TYPE_CHECKING: - import numpy as np import numpy.typing as npt import torch @@ -211,10 +212,16 @@ class StructuredOutputManager: raise def _fill_bitmasks( - self, batch: Iterable[tuple[StructuredOutputGrammar, int, bool]] + self, + batch: Iterable[tuple[StructuredOutputGrammar, int, bool, tuple[int, ...]]], ) -> None: + """Fill one bitmask row per entry. + + A row the grammar does not constrain allows every token except + `masked_ids`, the implicit reasoning ends masked while reasoning. + """ assert self._grammar_bitmask is not None - for grammar, index, apply_bitmask in batch: + for grammar, index, apply_bitmask, masked_ids in batch: if apply_bitmask and not grammar.is_terminated(): grammar.fill_bitmask(self._grammar_bitmask, index) else: @@ -222,12 +229,41 @@ class StructuredOutputManager: # reset the relevant part of the bitmask for consequent # requests here. self._grammar_bitmask[index].fill_(self._full_mask) + if masked_ids: + row = self._grammar_bitmask[index].numpy().view(np.uint32) + for token_id in masked_ids: + row[token_id >> 5] &= ~np.uint32(1 << (token_id & 31)) def _async_submit_fill_bitmask( - self, batch: list[tuple[StructuredOutputGrammar, int, bool]] + self, batch: list[tuple[StructuredOutputGrammar, int, bool, tuple[int, ...]]] ) -> Future: return self.executor_for_fillmask.submit(self._fill_bitmasks, batch) + def _masked_reasoning_ends(self, request: "Request") -> tuple[int, ...]: + """Implicit reasoning-end tokens the grammar rejects as first content. + + Such a token (e.g. a tool-call start) ends reasoning and is also the + first content token, so sampling one where the grammar cannot start + with it would make the output invalid. These stay masked while the + request reasons. The grammar does not advance before reasoning ends, + so the set is fixed per request. + """ + structured_req = request.structured_output_request + assert structured_req is not None + if structured_req.masked_reasoning_end_ids is None: + masked: tuple[int, ...] = () + reasoner = self._get_reasoner(request) + if isinstance(reasoner, ParserEngineReasoningAdapter): + grammar = structured_req.grammar + assert isinstance(grammar, StructuredOutputGrammar) + masked = tuple( + token_id + for token_id in sorted(reasoner.implicit_reasoning_end_token_ids) + if not grammar.validate_tokens([token_id]) + ) + structured_req.masked_reasoning_end_ids = masked + return structured_req.masked_reasoning_end_ids + def _get_constraint_bounds( self, request: "Request", @@ -377,7 +413,10 @@ class StructuredOutputManager: apply_bitmask = ( self._get_constraint_bounds(request, ()).constraint_start == 0 ) - batch.append((grammar, cumulative_index, apply_bitmask)) + masked_ids = ( + () if apply_bitmask else self._masked_reasoning_ends(request) + ) + batch.append((grammar, cumulative_index, apply_bitmask, masked_ids)) if len(batch) == self.fill_bitmask_parallel_batch_size: promises.append(self._async_submit_fill_bitmask(batch)) batch = [] @@ -405,6 +444,10 @@ class StructuredOutputManager: grammar_start, constraint_start = self._get_constraint_bounds( request, strip_speculative_padding(req_tokens) ) + # Rows before constraint_start are still reasoning. + masked_ids = ( + self._masked_reasoning_ends(request) if constraint_start else () + ) state_advancements = 0 seen_padding = False # Row filled from the last valid grammar state before a draft @@ -417,7 +460,11 @@ class StructuredOutputManager: cumulative_index += 1 continue apply_bitmask = not seen_padding and i >= constraint_start - self._fill_bitmasks(((grammar, cumulative_index, apply_bitmask),)) + reasoning = not seen_padding and i < constraint_start + row_masked = masked_ids if reasoning else () + self._fill_bitmasks( + ((grammar, cumulative_index, apply_bitmask, row_masked),) + ) if token == -1: seen_padding = True elif not seen_padding and i >= grammar_start: @@ -444,7 +491,11 @@ class StructuredOutputManager: bonus_apply = not seen_padding and constraint_start <= len( req_tokens ) - self._fill_bitmasks(((grammar, cumulative_index, bonus_apply),)) + bonus_reasoning = not seen_padding and not bonus_apply + bonus_masked = masked_ids if bonus_reasoning else () + self._fill_bitmasks( + ((grammar, cumulative_index, bonus_apply, bonus_masked),) + ) cumulative_index += 1 if state_advancements > 0: diff --git a/vllm/v1/structured_output/request.py b/vllm/v1/structured_output/request.py index cd044f5466..2091cc12f9 100644 --- a/vllm/v1/structured_output/request.py +++ b/vllm/v1/structured_output/request.py @@ -28,6 +28,9 @@ class StructuredOutputRequest: # Cached per request; do not share reasoning parsers across requests because # their behavior can depend on reasoning_parser_kwargs. reasoner: "ReasoningParser | None" = None + # Implicit reasoning-end tokens the grammar rejects as the first content + # token; masked while the request reasons. Computed on first use. + masked_reasoning_end_ids: tuple[int, ...] | None = None @staticmethod def from_sampling_params( diff --git a/vllm/v1/watermarking/gumbel.py b/vllm/v1/watermarking/gumbel.py index 2f953dd389..31e7bdd616 100644 --- a/vllm/v1/watermarking/gumbel.py +++ b/vllm/v1/watermarking/gumbel.py @@ -14,6 +14,7 @@ from vllm.v1.watermarking.detector import ( ) from vllm.v1.watermarking.prfs import PhiloxPRF, WatermarkPRF, create_prf from vllm.v1.watermarking.watermarker import ( + DraftContextState, RandomSampler, SupportsSpeculativeDecoding, Watermarker, @@ -101,6 +102,42 @@ class GumbelWatermarker(Watermarker): ) return WatermarkSample(token_ids, logits) + def _try_sample_block( + self, + logits: torch.Tensor, + num_steps: int, + random_sampler: RandomSampler, + state: DraftContextState, + ) -> torch.Tensor | None: + if type(self.prf) is not PhiloxPRF or logits.device.type != "cuda": + return None + + from vllm.v1.worker.gpu.sample.watermark import draft_philox_gumbel_sample + + assert random_sampler.logits_cache is not None + assert random_sampler.logits_cache_col is not None + return draft_philox_gumbel_sample( + logits, + state.contexts, + self.prf.key, + num_steps=num_steps, + expanded_idx_mapping=random_sampler.expanded_idx_mapping, + temperatures=random_sampler.temperatures, + seeds=random_sampler.seeds, + positions=random_sampler.positions, + enabled=state.enabled, + logits_cache=random_sampler.logits_cache, + logits_cache_col=random_sampler.logits_cache_col, + use_fp64=random_sampler.use_fp64, + prior_contexts=state.prior_contexts, + all_token_ids=state.all_token_ids, + prompt_lens=state.prompt_lens, + total_lens=state.total_lens, + deduplicate=state.deduplicate_contexts != "none", + max_history=state.deduplicate_contexts_max_history, + include_prompt=state.deduplicate_contexts == "all", + ) + class DualKeyGumbelWatermarker(Watermarker, SupportsSpeculativeDecoding): def __init__( diff --git a/vllm/v1/watermarking/spec_decode.py b/vllm/v1/watermarking/spec_decode.py index 3cfe334a90..6c432db959 100644 --- a/vllm/v1/watermarking/spec_decode.py +++ b/vllm/v1/watermarking/spec_decode.py @@ -8,6 +8,7 @@ from vllm.v1.watermarking.factory import create_watermarker from vllm.v1.watermarking.gumbel import GumbelWatermarker from vllm.v1.watermarking.prfs import PhiloxPRF from vllm.v1.watermarking.watermarker import ( + DraftContextState, RandomSampler, SupportsSpeculativeDecoding, Watermarker, @@ -28,25 +29,25 @@ class DraftWatermarker: ) -> None: self.watermarker = watermarker self.num_speculative_steps = num_speculative_steps - self.deduplicate_contexts = deduplicate_contexts - self.deduplicate_contexts_max_history = deduplicate_contexts_max_history - self.contexts = torch.zeros( - max_num_reqs, - watermarker.context_width, - dtype=torch.int64, - device=device, + self.state = DraftContextState( + contexts=torch.zeros( + max_num_reqs, + watermarker.context_width, + dtype=torch.int64, + device=device, + ), + enabled=torch.zeros(max_num_reqs, dtype=torch.bool, device=device), + prior_contexts=torch.zeros( + max_num_reqs, + num_speculative_steps, + watermarker.context_width, + dtype=torch.int64, + device=device, + ), + prompt_lens=torch.zeros(max_num_reqs, dtype=torch.int32, device=device), + deduplicate_contexts=deduplicate_contexts, + deduplicate_contexts_max_history=deduplicate_contexts_max_history, ) - self.enabled = torch.zeros(max_num_reqs, dtype=torch.bool, device=device) - self.prior_contexts = torch.zeros( - max_num_reqs, - num_speculative_steps, - watermarker.context_width, - dtype=torch.int64, - device=device, - ) - self.prompt_lens = torch.zeros(max_num_reqs, dtype=torch.int32, device=device) - self.all_token_ids: torch.Tensor | None = None - self.total_lens: torch.Tensor | None = None def prepare( self, @@ -57,11 +58,12 @@ class DraftWatermarker: total_lens: torch.Tensor, ) -> None: num_reqs = contexts.shape[0] - self.contexts[:num_reqs].copy_(contexts) - self.enabled[:num_reqs].copy_(enabled) - self.all_token_ids = all_token_ids - self.prompt_lens.copy_(prompt_lens) - self.total_lens = total_lens + state = self.state + state.contexts[:num_reqs].copy_(contexts) + state.enabled[:num_reqs].copy_(enabled) + state.all_token_ids = all_token_ids + state.prompt_lens.copy_(prompt_lens) + state.total_lens = total_lens def _sampling_state( self, @@ -75,22 +77,23 @@ class DraftWatermarker: steps = torch.as_tensor(draft_step, device=logits.device, dtype=torch.int64) if steps.ndim == 0: steps = steps.expand(num_rows) - enabled = self.enabled[:num_rows] & (request_temperatures != 0) - if self.deduplicate_contexts == "none": + state = self.state + enabled = state.enabled[:num_rows] & (request_temperatures != 0) + if state.deduplicate_contexts == "none": return steps, enabled - assert self.all_token_ids is not None - assert self.total_lens is not None + assert state.all_token_ids is not None + assert state.total_lens is not None enabled = draft_watermarking_mask( - self.all_token_ids, + state.all_token_ids, idx_mapping, - self.prompt_lens, - self.total_lens, - self.prior_contexts, + state.prompt_lens, + state.total_lens, + state.prior_contexts, contexts, steps, enabled, - self.deduplicate_contexts_max_history, - include_prompt=self.deduplicate_contexts == "all", + state.deduplicate_contexts_max_history, + include_prompt=state.deduplicate_contexts == "all", ) return steps, enabled @@ -107,14 +110,26 @@ class DraftWatermarker: logits_cache_col: torch.Tensor | None = None, use_fp64: bool = False, logits_cache_source: torch.Tensor | None = None, + num_steps: int = 1, ) -> torch.Tensor: + if num_steps > 1: + return self.sample_block( + logits, + idx_mapping, + temperature, + seed, + pos, + logits_cache=logits_cache, + logits_cache_col=logits_cache_col, + use_fp64=use_fp64, + ) assert apply_temperature assert is_drafting assert logits_cache is not None assert logits_cache_col is not None num_rows = logits.shape[0] request_temperatures = temperature[idx_mapping] - contexts = self.contexts[:num_rows] + contexts = self.state.contexts[:num_rows] steps, enabled = self._sampling_state( logits, idx_mapping, request_temperatures, logits_cache_col, contexts ) @@ -145,6 +160,68 @@ class DraftWatermarker: contexts.copy_(torch.cat((contexts[:, 1:], sampled.unsqueeze(-1)), dim=-1)) return sampled + def sample_block( + self, + logits: torch.Tensor, + idx_mapping: torch.Tensor, + temperature: torch.Tensor, + seed: torch.Tensor, + pos: torch.Tensor, + logits_cache: torch.Tensor | None = None, + logits_cache_col: torch.Tensor | None = None, + use_fp64: bool = False, + ) -> torch.Tensor: + """Sample drafts flattened as (request, step), one step at a time. + The target verifies each draft with a context that includes the earlier + drafts of its request, so steps must be sampled in order even when + their logits were computed in parallel. + """ + assert logits_cache_col is not None + num_steps = self.num_speculative_steps + sampled_block = self.watermarker._try_sample_block( + logits, + num_steps, + RandomSampler( + expanded_idx_mapping=idx_mapping, + temperatures=temperature, + seeds=seed, + positions=pos, + use_fp64=use_fp64, + is_drafting=True, + logits_cache=logits_cache, + logits_cache_col=logits_cache_col, + ), + self.state, + ) + if sampled_block is not None: + return sampled_block + num_reqs = logits.shape[0] // num_steps + + def by_step(tensor: torch.Tensor) -> torch.Tensor: + return tensor.view(num_reqs, num_steps, *tensor.shape[1:]) + + logits = by_step(logits) + idx_mapping = by_step(idx_mapping) + pos = by_step(pos) + logits_cache_col = by_step(logits_cache_col) + sampled = torch.empty( + num_reqs, num_steps, dtype=torch.int64, device=logits.device + ) + for step in range(num_steps): + sampled[:, step] = self.sample( + logits[:, step], + idx_mapping[:, step], + temperature, + seed, + pos[:, step], + apply_temperature=True, + is_drafting=True, + logits_cache=logits_cache, + logits_cache_col=logits_cache_col[:, step], + use_fp64=use_fp64, + ) + return sampled.view(-1) + def create_speculative_target_watermarker(watermarker: Watermarker) -> Watermarker: if isinstance(watermarker, SupportsSpeculativeDecoding): diff --git a/vllm/v1/watermarking/watermarker.py b/vllm/v1/watermarking/watermarker.py index 15917d2520..13b806faac 100644 --- a/vllm/v1/watermarking/watermarker.py +++ b/vllm/v1/watermarking/watermarker.py @@ -7,6 +7,7 @@ from typing import Protocol, runtime_checkable import torch +from vllm.config.watermarking import WatermarkContextScope from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample @@ -44,6 +45,18 @@ class RandomSampler: ) +@dataclass +class DraftContextState: + contexts: torch.Tensor + enabled: torch.Tensor + prior_contexts: torch.Tensor + prompt_lens: torch.Tensor + deduplicate_contexts: WatermarkContextScope + deduplicate_contexts_max_history: int | None + all_token_ids: torch.Tensor | None = None + total_lens: torch.Tensor | None = None + + class Watermarker(ABC): @property @abstractmethod @@ -94,6 +107,20 @@ class Watermarker(ABC): ) -> WatermarkSample | None: return None + def _try_sample_block( + self, + logits: torch.Tensor, + num_steps: int, + random_sampler: RandomSampler, + state: DraftContextState, + ) -> torch.Tensor | None: + """Sample a (request, step) draft block in one fused pass. + + Returns None when there is no fused implementation, in which case the + caller samples the block one step at a time. + """ + return None + @runtime_checkable class SupportsSpeculativeDecoding(Protocol): diff --git a/vllm/v1/worker/cpu_worker.py b/vllm/v1/worker/cpu_worker.py index b103215d65..22b8618f3b 100644 --- a/vllm/v1/worker/cpu_worker.py +++ b/vllm/v1/worker/cpu_worker.py @@ -234,10 +234,24 @@ class CPUWorker(Worker): else: consumed_memory = psutil.Process(os.getpid()).memory_info().rss requested_memory_for_kv = int(self.requested_cpu_memory - consumed_memory) - if ( - requested_memory_for_kv <= 0 - or requested_memory_for_kv > available_memory - ): + if requested_memory_for_kv <= 0: + # The worker's own footprint already fills its reservation, so + # freeing memory elsewhere cannot help: only a larger + # reservation or an explicit KV cache size can. + in_use = math.ceil(consumed_memory / memory_status.total_memory * 100) + raise ValueError( + f"The worker already uses {format_gib(consumed_memory)} GiB " + "(model weights and runtime), more than the " + f"{format_gib(self.requested_cpu_memory)} GiB of CPU memory " + f"reserved for it on node {cpu_core.numa_node}, leaving none " + "for the KV cache. On the CPU backend, the " + "`--gpu-memory-utilization` flag controls the fraction of CPU " + "memory reserved (despite its name). To resolve: increase " + f"`--gpu-memory-utilization` to more than {in_use / 100:.2f}, " + "or set the KV cache size explicitly with " + "`--kv-cache-memory-bytes` or VLLM_CPU_KVCACHE_SPACE." + ) + if requested_memory_for_kv > available_memory: raise ValueError( f"Available memory on node {cpu_core.numa_node} " f"({format_gib(available_memory)}/" diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 25d0cbc153..664004ae1d 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -33,6 +33,7 @@ from vllm.distributed.parallel_state import ( ) from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.logger import init_logger +from vllm.model_executor.models.interfaces import requires_raw_input_tokens from vllm.model_executor.offloader.base import get_offloader from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors @@ -690,7 +691,8 @@ class ModelCudaGraphManager(CudaGraphManager): } if not self.is_first_pp_rank: # Update for non-first PP ranks. - model_inputs["input_ids"] = None + if not requires_raw_input_tokens(model): + model_inputs["input_ids"] = None model_inputs["inputs_embeds"] = None assert intermediate_tensors is not None model_inputs["intermediate_tensors"] = intermediate_tensors[:num_tokens] diff --git a/vllm/v1/worker/gpu/dp_utils.py b/vllm/v1/worker/gpu/dp_utils.py index 5c71ebcea0..2bcad34e05 100644 --- a/vllm/v1/worker/gpu/dp_utils.py +++ b/vllm/v1/worker/gpu/dp_utils.py @@ -5,7 +5,6 @@ from __future__ import annotations from dataclasses import dataclass, replace import torch -import torch.distributed as dist from vllm.config import ParallelConfig from vllm.config.compilation import CUDAGraphMode @@ -63,7 +62,6 @@ def sync_cudagraph_and_dp_padding( Returns (synced_batch_desc, sync). `sync` is None when no rank has work. """ assert dp_size > 1, "DP size must be greater than 1" - group = get_dp_group().cpu_group tensor = torch.zeros(6, dp_size, dtype=torch.int32, device="cpu") tensor[0][dp_rank] = num_tokens tensor[1][dp_rank] = desired_batch_desc.cg_mode.value @@ -74,7 +72,7 @@ def sync_cudagraph_and_dp_padding( if should_skip_dp_coordination(): tensor[:] = tensor[:, dp_rank, None].clone() else: - dist.all_reduce(tensor, group=group) + get_dp_group().cpu_all_reduce_(tensor) num_tokens_across_dp = tensor[0] cg_mode_across_dp = tensor[1] diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 5ba90b792e..1b75d73493 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -33,7 +33,7 @@ import vllm.envs as envs from vllm.compilation.counter import compilation_counter from vllm.compilation.cuda_graph import CUDAGraphStat from vllm.compilation.wrapper import compile_model_with_stock_torch -from vllm.config import VllmConfig +from vllm.config import VllmConfig, set_current_vllm_config from vllm.config.compilation import CompilationMode, CUDAGraphMode from vllm.distributed.aux_output_connector.worker import ( AuxOutputWorkerConnector, @@ -46,6 +46,9 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_ma from vllm.model_executor.layers.mamba.ops.ssu_dispatch import ( initialize_mamba_ssu_backend, ) +from vllm.model_executor.layers.quantization.online.lm_head import ( + refresh_quantized_lm_heads, +) from vllm.model_executor.model_loader import get_model_loader from vllm.model_executor.model_loader.reload.inplace import reload_weights from vllm.model_executor.models.interfaces import requires_raw_input_tokens @@ -569,6 +572,15 @@ class GPUModelRunner(LoRAModelRunnerMixin): return None return speculator.model + def reset_lora_state(self) -> None: + super().reset_lora_state() + # Called whenever the target's weights are replaced (reload, weight + # transfer), so draft weights derived from them must follow. + draft_model = self.get_draft_model() + if draft_model is not None: + with set_current_vllm_config(self.vllm_config): + refresh_quantized_lm_heads(draft_model) + def reload_weights( self, weights_iterator: Iterable[tuple[str, torch.Tensor]] | None = None, @@ -1970,7 +1982,8 @@ class GPUModelRunner(LoRAModelRunnerMixin): } if not self.is_first_pp_rank: # Update for non-first PP ranks. - model_inputs["input_ids"] = None + if not requires_raw_input_tokens(self.model): + model_inputs["input_ids"] = None model_inputs["inputs_embeds"] = None # Prepare the intermediate tensors. diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py index a34169b733..9ab0579a6f 100644 --- a/vllm/v1/worker/gpu/sample/gumbel.py +++ b/vllm/v1/worker/gpu/sample/gumbel.py @@ -70,6 +70,12 @@ def tl_rand32(seed, offset, includes_zero: tl.constexpr): return u +@triton.jit +def spec_step_seed(seed, step): + """Mix a speculative step into the high bits of a request seed.""" + return seed ^ (tl.cast(step, tl.int64) << 32) + + @triton.jit def _murmur3_rotl32(value, shift: tl.constexpr): return (value << shift) | (value >> (32 - shift)) @@ -257,6 +263,8 @@ def gumbel_block_argmax( seed = tl.load(seeds_ptr + req_state_idx, mask=is_valid_req, other=0) pos = tl.load(pos_ptr + token_idx) + if logits_cache_ptr is not None: + seed = spec_step_seed(seed, col) return gumbel_noised_argmax( logits, block, diff --git a/vllm/v1/worker/gpu/sample/watermark.py b/vllm/v1/worker/gpu/sample/watermark.py index 216d189f8d..00a137fd1b 100644 --- a/vllm/v1/worker/gpu/sample/watermark.py +++ b/vllm/v1/worker/gpu/sample/watermark.py @@ -25,7 +25,7 @@ from vllm.v1.watermarking.prfs.philox import ( from vllm.v1.watermarking.prfs.philox import ( _UINT32_MASK as _UINT32_MASK_VALUE, ) -from vllm.v1.worker.gpu.sample.gumbel import gumbel_noised_argmax +from vllm.v1.worker.gpu.sample.gumbel import gumbel_noised_argmax, spec_step_seed if HAS_TRITON: from triton.language import math as tl_math @@ -530,8 +530,10 @@ def _philox_gumbel_from_logits(logits, word): @triton.jit -def _gumbel_value(logits_ptr, output, mask): +def _gumbel_value(logits_ptr, output, mask, temperature): logits = tl.load(logits_ptr, mask=mask, other=float("-inf")).to(tl.float32) + if temperature is not None: + logits = tl_math.div_rn(logits, temperature) return _philox_gumbel_from_logits(logits, output) @@ -606,6 +608,7 @@ def _philox_gumbel_kernel( USE_FP64: tl.constexpr, IS_DRAFTING: tl.constexpr, PER_TOKEN_COL: tl.constexpr, + APPLY_TEMPERATURE: tl.constexpr, ): row = tl.program_id(0).to(tl.int64) block_index = tl.program_id(1) @@ -619,6 +622,13 @@ def _philox_gumbel_kernel( req_state_idx = tl.load(expanded_idx_mapping_ptr + row).to(tl.int64) valid_req = req_state_idx >= 0 + row_temperature = None + if APPLY_TEMPERATURE: + row_temperature = tl.load( + temp_ptr + req_state_idx, mask=valid_req, other=0.0 + ).to(tl.float32) + row_temperature = tl.where(row_temperature == 0, 1.0, row_temperature) + if logits_cache_ptr is not None: if PER_TOKEN_COL: col = tl.load(logits_cache_col_ptr + row) @@ -644,10 +654,14 @@ def _philox_gumbel_kernel( ) seed = tl.load(seeds_ptr + req_state_idx, mask=valid_req, other=0) pos = tl.load(pos_ptr + row) + if logits_cache_ptr is not None: + seed = spec_step_seed(seed, col) logits_row = logits_ptr + row * logits_stride logits = tl.load(logits_row + candidate, mask=mask, other=float("-inf")).to( tl.float32 ) + if row_temperature is not None: + logits = tl_math.div_rn(logits, row_temperature) value, index = gumbel_noised_argmax( logits, candidate, @@ -681,16 +695,16 @@ def _philox_gumbel_kernel( candidate_3 = candidate_0 + 3 logits_row = logits_ptr + row * logits_stride value_0 = _gumbel_value( - logits_row + candidate_0, output_0, candidate_0 < vocab_size + logits_row + candidate_0, output_0, candidate_0 < vocab_size, row_temperature ) value_1 = _gumbel_value( - logits_row + candidate_1, output_1, candidate_1 < vocab_size + logits_row + candidate_1, output_1, candidate_1 < vocab_size, row_temperature ) value_2 = _gumbel_value( - logits_row + candidate_2, output_2, candidate_2 < vocab_size + logits_row + candidate_2, output_2, candidate_2 < vocab_size, row_temperature ) value_3 = _gumbel_value( - logits_row + candidate_3, output_3, candidate_3 < vocab_size + logits_row + candidate_3, output_3, candidate_3 < vocab_size, row_temperature ) best_value = value_0 best_token = candidate_0 @@ -797,6 +811,238 @@ def philox_gumbel_sample( USE_FP64=use_fp64, IS_DRAFTING=is_drafting, PER_TOKEN_COL=logits_cache_col is not None and logits_cache_col.dim() > 0, + APPLY_TEMPERATURE=False, ) max_block_index = local_max.argmax(dim=-1, keepdim=True) return local_argmax.gather(dim=-1, index=max_block_index).view(-1) + + +@triton.jit(do_not_specialize=["step"]) +def _draft_step_kernel( + sampled_ptr, + sampled_stride, + step, + local_argmax_ptr, + local_max_ptr, + num_blocks, + contexts_ptr, + context_stride, + skip_mask_ptr, + next_req_indices_ptr, + next_steps_ptr, + temp_ptr, + enabled_ptr, + all_token_ids_ptr, + all_token_ids_stride, + prompt_lens_ptr, + total_lens_ptr, + prior_contexts_ptr, + prior_contexts_stride_0, + prior_contexts_stride_1, + CONTEXT_WIDTH: tl.constexpr, + NUM_SPECULATIVE_STEPS: tl.constexpr, + MAX_HISTORY: tl.constexpr, + INCLUDE_PROMPT: tl.constexpr, + DEDUPLICATE: tl.constexpr, + SAMPLE: tl.constexpr, + PREPARE_NEXT: tl.constexpr, + BLOCKS_POW2: tl.constexpr, +): + """Finish draft step `step` (if SAMPLE) and build the next step's skip mask.""" + row = tl.program_id(0).to(tl.int64) + contexts_row_ptr = contexts_ptr + row * context_stride + if SAMPLE: + # Same pick as local_max.argmax(dim=-1): the first NaN, else the first max. + blocks = tl.arange(0, BLOCKS_POW2) + values = tl.load( + local_max_ptr + row * num_blocks + blocks, + mask=blocks < num_blocks, + other=float("-inf"), + ) + is_nan = values != values + best = tl.max(tl.where(is_nan, float("-inf"), values), axis=0) + first_nan = tl.min(tl.where(is_nan, blocks, BLOCKS_POW2), axis=0) + first_max = tl.min( + tl.where((values == best) & (blocks < num_blocks), blocks, BLOCKS_POW2), + axis=0, + ) + block = tl.where(first_nan < BLOCKS_POW2, first_nan, first_max) + token = tl.load(local_argmax_ptr + row * num_blocks + block) + tl.store(sampled_ptr + row * sampled_stride + step, token) + + offsets = tl.arange(0, BLOCKS_POW2) + context = tl.load( + contexts_row_ptr + offsets + 1, + mask=offsets < CONTEXT_WIDTH - 1, + other=token, + ) + tl.debug_barrier() + tl.store(contexts_row_ptr + offsets, context, mask=offsets < CONTEXT_WIDTH) + tl.debug_barrier() + + if PREPARE_NEXT: + req_idx = tl.load(next_req_indices_ptr + row).to(tl.int64) + temperature = tl.load(temp_ptr + req_idx, mask=req_idx >= 0, other=0.0) + active = tl.load(enabled_ptr + row) & (temperature != 0) + if DEDUPLICATE: + # The dedup kernel body runs in place on the skip mask buffer. + tl.store(skip_mask_ptr + row, active) + tl.debug_barrier() + _repeated_context_mask_kernel( + skip_mask_ptr, + all_token_ids_ptr, + all_token_ids_stride, + next_req_indices_ptr, + prompt_lens_ptr, + total_lens_ptr, + next_steps_ptr, + 1, + contexts_ptr, + context_stride, + None, + prior_contexts_ptr, + prior_contexts_stride_0, + prior_contexts_stride_1, + next_steps_ptr, + 1, + skip_mask_ptr, + CONTEXT_WIDTH=CONTEXT_WIDTH, + NUM_SPECULATIVE_STEPS=NUM_SPECULATIVE_STEPS, + MAX_HISTORY=MAX_HISTORY, + INCLUDE_PROMPT=INCLUDE_PROMPT, + SKIP_PARTIAL_CONTEXT=False, + HAS_HISTORY_OFFSETS=True, + SPECULATIVE_CONTEXTS=False, + DRAFT_CONTEXTS=True, + BLOCK=512, + ) + tl.debug_barrier() + active = tl.load(skip_mask_ptr + row) + tl.store(skip_mask_ptr + row, active == 0) + + +def draft_philox_gumbel_sample( + logits: torch.Tensor, + contexts: torch.Tensor, + key: int, + *, + num_steps: int, + expanded_idx_mapping: torch.Tensor, + temperatures: torch.Tensor, + seeds: torch.Tensor, + positions: torch.Tensor, + enabled: torch.Tensor, + logits_cache: torch.Tensor, + logits_cache_col: torch.Tensor, + use_fp64: bool, + prior_contexts: torch.Tensor, + all_token_ids: torch.Tensor | None, + prompt_lens: torch.Tensor, + total_lens: torch.Tensor | None, + deduplicate: bool, + max_history: int | None, + include_prompt: bool, +) -> torch.Tensor: + """Sample `num_steps` draft steps flattened as (request, step), in step order. + + Bit-exact with sampling each step by `DraftWatermarker.sample`, but with two + kernel launches per step: the Philox block argmax (with the temperature + divide and the logits cache write), then a per-row kernel that picks the + token, shifts `contexts` and builds the next step's skip mask. Padding rows + (index -1) read temperature 0 and are never watermarked, so they may differ. + """ + num_rows, vocab_size = logits.shape + num_reqs = num_rows // num_steps + assert logits.stride(-1) == 1 and contexts.stride(-1) == 1 + assert logits.dtype == logits_cache.dtype, ( + "logits cache source and destination must have the same dtype" + ) + assert logits_cache.size(-1) >= vocab_size, ( + f"draft logits cache vocab dim ({logits_cache.size(-1)}) is narrower " + f"than the sampled logits ({vocab_size}). Cached logits would be truncated." + ) + assert not deduplicate or (all_token_ids is not None and total_lens is not None), ( + "context deduplication requires all_token_ids and total_lens" + ) + logits = logits.view(num_reqs, num_steps, vocab_size) + # Step-major copies, so each step reads contiguous per-request slices. + expanded_idx_mapping, positions, logits_cache_col = ( + tensor.view(num_reqs, num_steps).t().contiguous() + for tensor in (expanded_idx_mapping, positions, logits_cache_col) + ) + block_size = 1024 + num_blocks = triton.cdiv(vocab_size, block_size) + local_argmax = logits.new_empty(num_reqs, num_blocks, dtype=torch.int64) + local_max = logits.new_empty( + num_reqs, num_blocks, dtype=torch.float64 if use_fp64 else torch.float32 + ) + skip_mask = torch.empty(num_reqs, dtype=torch.bool, device=logits.device) + sampled = logits.new_empty(num_reqs, num_steps, dtype=torch.int64) + + def finish_step(step: int) -> None: + next_step = min(step + 1, num_steps - 1) + _draft_step_kernel[(num_reqs,)]( + sampled, + sampled.stride(0), + step, + local_argmax, + local_max, + num_blocks, + contexts, + contexts.stride(0), + skip_mask, + expanded_idx_mapping[next_step], + logits_cache_col[next_step], + temperatures, + enabled, + all_token_ids, + 0 if all_token_ids is None else all_token_ids.stride(0), + prompt_lens, + total_lens, + prior_contexts, + prior_contexts.stride(0), + prior_contexts.stride(1), + CONTEXT_WIDTH=contexts.shape[-1], + NUM_SPECULATIVE_STEPS=prior_contexts.shape[1], + MAX_HISTORY=0 if max_history is None else max_history, + INCLUDE_PROMPT=include_prompt, + DEDUPLICATE=deduplicate, + SAMPLE=step >= 0, + PREPARE_NEXT=step + 1 < num_steps, + BLOCKS_POW2=triton.next_power_of_2(max(num_blocks, contexts.shape[-1])), + ) + + finish_step(-1) + for step in range(num_steps): + _philox_gumbel_kernel[(num_reqs, num_blocks)]( + local_argmax, + local_argmax.stride(0), + local_max, + local_max.stride(0), + logits[:, step], + logits.stride(0), + contexts, + contexts.stride(0), + skip_mask, + expanded_idx_mapping[step], + seeds, + positions[step], + temperatures, + logits_cache, + logits_cache.stride(0), + logits_cache.stride(1), + logits_cache_col[step], + logits[:, step], + logits.stride(0), + key & _UINT32_MASK_VALUE, + key >> 32, + vocab_size, + CONTEXT_WIDTH=contexts.shape[-1], + BLOCK_SIZE=block_size, + USE_FP64=use_fp64, + IS_DRAFTING=True, + PER_TOKEN_COL=True, + APPLY_TEMPERATURE=True, + ) + finish_step(step) + return sampled.view(-1) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index a75e84100e..af519bb5db 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -366,6 +366,7 @@ class DFlashSpeculator(DraftModelSpeculator): self.seeds, self.sample_col[:num_sample], self.draft_logits, + num_steps=self.num_speculative_steps, ) self.draft_tokens[:num_reqs] = draft_tokens.view( num_reqs, self.num_speculative_steps diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py index bb14621129..110a2b34a7 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py @@ -8,7 +8,7 @@ import torch from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode from vllm.triton_utils import tl, triton -from vllm.v1.worker.gpu.sample.gumbel import gumbel_noised_argmax +from vllm.v1.worker.gpu.sample.gumbel import gumbel_noised_argmax, spec_step_seed from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator @@ -58,7 +58,7 @@ def _selector_walk_kernel( scores, candidates, mask & valid, - seed, + spec_step_seed(seed, step), sample_pos, temperature if SAMPLE_PROBABILISTIC else 0.0, IS_DRAFTING=True, diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py index 65289bb18e..63e999b75f 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py @@ -6,6 +6,9 @@ import torch.nn as nn from vllm.config import VllmConfig, replace from vllm.distributed.parallel_state import get_pp_group from vllm.lora.layers.base import BaseLayerWithLoRA +from vllm.model_executor.layers.quantization.online.lm_head import ( + quantized_lm_head_copy, +) from vllm.model_executor.model_loader import get_model from vllm.model_executor.model_loader.utils import get_draft_load_config from vllm.model_executor.models.utils import PPMissingLayer @@ -106,9 +109,14 @@ def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mod if target_lm_head is not None and _should_share( eagle_model, "has_own_lm_head", draft_lm_head, target_lm_head ): + draft_head = target_lm_head + if speculative_config.draft_lm_head_quantization is not None: + draft_head = quantized_lm_head_copy( + target_lm_head, speculative_config.draft_lm_head_quantization + ) if draft_lm_head is not None: del eagle_model.lm_head - eagle_model.lm_head = target_lm_head + eagle_model.lm_head = draft_head # MTP layers route logits through layer.shared_head.head, not # eagle_model.lm_head, so the per-layer copies need fixing up too. @@ -119,7 +127,11 @@ def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mod sh = getattr(layer, "shared_head", None) if sh is not None and hasattr(sh, "head"): del sh.head - sh.head = target_lm_head + sh.head = draft_head + elif speculative_config.draft_lm_head_quantization is not None: + raise ValueError( + "draft_lm_head_quantization needs a drafter that shares the target lm_head." + ) # MTP shares topk_indices_buffer with the target model. We update # every module in the draft that holds a buffer reference so that diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py index 4387b1d38a..f35468d62b 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py @@ -3,7 +3,11 @@ import torch from vllm.triton_utils import tl, tldevice, triton -from vllm.v1.worker.gpu.sample.gumbel import gumbel_block_argmax, tl_rand32 +from vllm.v1.worker.gpu.sample.gumbel import ( + gumbel_block_argmax, + spec_step_seed, + tl_rand32, +) from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_block_argmax @@ -569,7 +573,8 @@ def _rejection_kernel( verifying &= is_valid_draft if verifying: pos = tl.load(pos_ptr + logit_idx) - u = tl_rand32(seed, pos, includes_zero=False) + step_seed = spec_step_seed(seed, i) if USE_BLOCK_VERIFICATION else seed + u = tl_rand32(step_seed, pos, includes_zero=False) if is_greedy: # Greedy sampling. Accept IFF draft matches target argmax. # NOTE: Target argmax is stored directly so that resampling diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index d86a0b0351..c54df5ac17 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -3,6 +3,7 @@ from abc import ABC, abstractmethod from collections.abc import Mapping from dataclasses import replace +from functools import partial from typing import TYPE_CHECKING, Any import numpy as np @@ -446,6 +447,7 @@ class DraftModelSpeculator(BaseSpeculator): draft_step: torch.Tensor, draft_logits: torch.Tensor | None, spec_step_idx: int = 0, + num_steps: int = 1, ) -> torch.Tensor: if draft_logits is None and self.use_local_argmax_reduction: return self.get_draft_top_tokens(hidden_states, spec_step_idx) @@ -455,7 +457,7 @@ class DraftModelSpeculator(BaseSpeculator): sampler = ( gumbel_sample if self.draft_watermarker is None - else self.draft_watermarker.sample + else partial(self.draft_watermarker.sample, num_steps=num_steps) ) sampled = sampler( logits, diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index ccd50dd3e4..f3ffef5810 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -62,6 +62,10 @@ from vllm.distributed.weight_transfer import ( ) from vllm.logger import init_logger from vllm.lora.request import LoRARequest +from vllm.model_executor.layers.quantization.online.lm_head import ( + restore_quantized_lm_heads, + share_source_lm_heads, +) from vllm.model_executor.model_loader import get_model_loader from vllm.model_executor.model_loader.weight_checksum import ( compute_tensor_digests, @@ -1544,12 +1548,19 @@ class Worker(WorkerBase): "active. Call finish_weight_update first." ) + draft_model = self.get_draft_model() if is_draft else None try: if is_draft: self._set_draft_weight_update_target() + if draft_model is not None: + # Load through the unquantized target head, as without + # draft_lm_head_quantization; finish re-derives the copies. + share_source_lm_heads(draft_model) self.weight_transfer_engine.start_weight_update() except BaseException: self.weight_transfer_engine.reset_weight_update_target() + if draft_model is not None: + restore_quantized_lm_heads(draft_model, refresh=False) raise self._weight_update_active = True self._weight_update_is_draft = is_draft @@ -1589,6 +1600,10 @@ class Worker(WorkerBase): except BaseException: self._weight_update_active = False self.weight_transfer_engine.reset_weight_update_target() + if self._weight_update_is_draft and ( + (draft_model := self.get_draft_model()) is not None + ): + restore_quantized_lm_heads(draft_model, refresh=False) raise def finish_weight_update(self) -> None: @@ -1605,6 +1620,10 @@ class Worker(WorkerBase): self.weight_transfer_engine.finish_weight_update() self.weight_transfer_engine.reset_weight_update_target() self._weight_update_active = False + if self._weight_update_is_draft and ( + (draft_model := self.get_draft_model()) is not None + ): + restore_quantized_lm_heads(draft_model) # Weight transfer bypasses GPUModelRunner.reload_weights(). if not self._weight_update_is_draft: