From 067652c38fd1ede1d1ffab98378f88344ca56f9f Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Tue, 27 Oct 2026 22:43:10 +0800 Subject: [PATCH 1/4] [Model] DeepSeek-V4.1-Flash on 4x V100: 512K packed KV, faster prefill, API launch Runs DeepSeek-V4.1-Flash-1017 (GSQ-RCO IQ2_XXS GGUF + Engram IQ4_NL + DSpark5 draft) on 4x V100-32G (TP4, one x16 PCIe uplink, 62 GB host RAM) with expert offloading. All new paths are lossless and switchable. Prefill (G1-G7): - G1 activation memory: per-tile MoE scratch + FP32 accumulator, in-place expand_hc / large all-reduce, windowed hc_transition, pynccl all_gather, empty_cache before big passes; chunks up to 16K. - G2 cold-expert staging ring issued ahead across tiles and layers. - G3 routed down projection split by output rows (no duplicated TP boundary blocks, -2.25 GiB/GPU) with an IPC peer gather for decode. - G5 Engram readers, adaptive indexer tile, persistent activation grid. - G6 routed-only staging for short prompts. - G7 same-session prefix cache (scheduler guard + worker restore_prefix). Long context (G4): - Window layers keep 256-row private rings. - Packed KV (VLLM_V41_PACKED_KV): FP4 + E4M3 compressed latents in a shared host memfd plus a 32K GPU hot ring, FP4 index keys on GPU, CUDA decode index scores and segment-parallel exact top-k. Enables 512K. - RoPE rotates on the fly past its table; C2/Engram/guard limits and the DSpark draft length cap are lifted to the target length. Results vs phase 1: TTFT short 5.6 -> 0.8-1.1 s, 8K 27.7 -> 5.8 s, 15K 44.0 -> 11.1 s; 128K 110 s and 500K 651 s (needle passes); follow-up turns hit the prefix cache in ~1-1.7 s; fresh-text decode 57.8 -> 53.5 ms/round. Tokens at <=16K are identical with packed KV. Tooling: benchmarks/deepseek_v41/ gains the v41k SM70 extension source and build script, serve configs (phase 2: 16K, phase 3: 512K + prefix cache + tools + thinking), launch/validation tools, GPU micro-tests and HANDOFF.md. V4.1 switches are registered in vllm/envs.py and read through vllm.envs. Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 --- benchmarks/deepseek_v41/HANDOFF.md | 139 + benchmarks/deepseek_v41/README.md | 166 + benchmarks/deepseek_v41/acceptance.py | 375 ++ benchmarks/deepseek_v41/adapter_audit.py | 72 + benchmarks/deepseek_v41/capacity.py | 99 + benchmarks/deepseek_v41/checkpoint_audit.py | 275 ++ .../compare_official_primitives.py | 212 ++ .../deepseek_v41/draft_checkpoint_probe.py | 124 + benchmarks/deepseek_v41/finite_debug.py | 60 + benchmarks/deepseek_v41/main_smoke.py | 251 ++ benchmarks/deepseek_v41/oracle_primitives.py | 185 + benchmarks/deepseek_v41/performance.py | 177 + benchmarks/deepseek_v41/prepare_dspark.py | 74 + benchmarks/deepseek_v41/resources.py | 89 + .../deepseek_v41/results/54633-20261024.md | 172 + benchmarks/deepseek_v41/serve-54633.yaml | 57 + .../deepseek_v41/serve-ds41-phase1.yaml | 58 + .../deepseek_v41/serve-ds41-phase2.yaml | 59 + .../deepseek_v41/serve-ds41-phase3.yaml | 64 + benchmarks/deepseek_v41/shared_cache_probe.py | 235 ++ .../deepseek_v41/shared_checkpoint_probe.py | 259 ++ benchmarks/deepseek_v41/shared_host_probe.py | 267 ++ .../deepseek_v41/tests/bench_g3_down.py | 101 + .../deepseek_v41/tests/bench_index_decode.py | 47 + .../deepseek_v41/tests/test_fused_sparse.py | 63 + benchmarks/deepseek_v41/tests/test_g1mem.py | 96 + benchmarks/deepseek_v41/tests/test_g3_down.py | 187 + benchmarks/deepseek_v41/tests/test_hc_tok.py | 51 + .../deepseek_v41/tests/test_index_packed.py | 90 + benchmarks/deepseek_v41/tests/test_kv_pack.py | 33 + .../deepseek_v41/tests/test_packed_attn.py | 105 + .../deepseek_v41/tests/test_peer_gather.py | 87 + .../deepseek_v41/tests/test_rope_tail.py | 24 + .../deepseek_v41/tests/test_topk_par.py | 50 + .../deepseek_v41/tests/test_window_ring.py | 49 + .../deepseek_v41/tools/api_multiturn.py | 73 + benchmarks/deepseek_v41/tools/api_smoke.py | 103 + benchmarks/deepseek_v41/tools/api_soak.py | 213 ++ benchmarks/deepseek_v41/tools/api_think.py | 93 + benchmarks/deepseek_v41/tools/decode_loop.py | 45 + .../deepseek_v41/tools/fresh_prompts.py | 78 + benchmarks/deepseek_v41/tools/longctx_test.py | 109 + benchmarks/deepseek_v41/tools/memlog2.sh | 13 + .../deepseek_v41/tools/prefill_profile.py | 105 + .../deepseek_v41/tools/profile_torch.py | 109 + benchmarks/deepseek_v41/tools/real_decode.py | 188 + .../deepseek_v41/tools/run_decode_spy.sh | 20 + benchmarks/deepseek_v41/tools/run_e2e2.sh | 18 + benchmarks/deepseek_v41/tools/run_longctx.sh | 11 + .../deepseek_v41/tools/run_prefill_profile.sh | 11 + benchmarks/deepseek_v41/tools/run_real.sh | 12 + benchmarks/deepseek_v41/tools/run_serve2.sh | 13 + benchmarks/deepseek_v41/tools/run_serve3.sh | 13 + benchmarks/deepseek_v41/tools/run_session.sh | 11 + benchmarks/deepseek_v41/tools/session_test.py | 77 + benchmarks/deepseek_v41/tools/spy_decode.py | 48 + benchmarks/deepseek_v41/tools/spy_top.py | 34 + benchmarks/deepseek_v41/tools/tcp_forward.py | 41 + benchmarks/deepseek_v41/tools/top_decode.py | 56 + benchmarks/deepseek_v41/tools/v41_e2e.py | 335 ++ benchmarks/deepseek_v41/trace.py | 246 ++ benchmarks/deepseek_v41/v41k/.gitignore | 1 + benchmarks/deepseek_v41/v41k/build.py | 35 + benchmarks/deepseek_v41/v41k/v41k.cu | 3029 +++++++++++++++++ csrc/gguf_cpu_iq_tables.h | 733 ++++ csrc/gguf_expert_cache.cpp | 132 + csrc/gguf_shared_engram.h | 289 ++ csrc/gguf_shared_expert_cpu.h | 383 +++ csrc/gguf_shared_expert_dma.h | 186 + csrc/gguf_shared_expert_jobs.h | 235 ++ csrc/gguf_shared_expert_lane.h | 620 ++++ csrc/gguf_shared_expert_planner.h | 653 ++++ csrc/gguf_shared_expert_state.h | 194 ++ .../gguf_upstream/runtime_dequant_blas.cu | 204 +- .../gguf_upstream/torch_bindings.cpp | 13 +- csrc/sm70_turbomind/ops/awq_sm70_gemm.cu | 12 +- csrc/sm70_turbomind/ops/gguf_dp4a.cuh | 173 + csrc/sm70_turbomind/ops/gguf_dp4a_sm70.cu | 605 +++- .../ops/gguf_prefill_transcode.cuh | 86 + .../ops/gguf_prefill_transcode_sm70.cu | 3 + csrc/torch_bindings.cpp | 65 + docs/design/architecture/runtime_reference.md | 6 +- pyproject.toml | 2 + .../test_deepseek_v41_acceptance.py | 65 + .../test_deepseek_v41_checkpoint_audit.py | 150 + .../test_deepseek_v41_oracle_primitives.py | 113 + tests/benchmarks/test_deepseek_v41_trace.py | 91 + tests/config/test_gguf_shared_original.py | 57 + .../test_deepseek_v41_activations.py | 74 + .../quantization/test_deepseek_v41_dense.py | 60 + .../test_deepseek_v41_draft_mxfp4.py | 138 + .../quantization/test_deepseek_v41_tail.py | 143 + .../quantization/test_gguf_grouped_prefill.py | 28 +- .../quantization/test_gguf_q5_batched.py | 39 + .../kernels/quantization/test_gguf_q5_hmma.py | 97 + .../quantization/test_gguf_q5_vector.py | 66 + .../quantization/test_gguf_v41_dp4a.py | 264 ++ .../test_gguf_v41_expert_reuse.py | 112 + .../kernels/test_deepseek_v41_compression.py | 138 + tests/kernels/test_deepseek_v41_mhc.py | 202 ++ tests/models/gguf/shared_engram_test.cpp | 158 + .../models/gguf/shared_expert_state_test.cpp | 288 ++ .../models/gguf/test_deepseek_v41_adapter.py | 176 + tests/models/gguf/test_shared_acquire.py | 81 + tests/models/gguf/test_shared_engram_io.py | 220 ++ .../models/gguf/test_shared_engram_native.py | 57 + tests/models/gguf/test_shared_expert_copy.py | 174 + .../models/gguf/test_shared_expert_layout.py | 112 + .../models/gguf/test_shared_expert_loading.py | 286 ++ .../models/gguf/test_shared_expert_native.py | 54 + tests/models/gguf/test_shared_expert_state.py | 66 + tests/models/gguf/test_v41_device_codecs.py | 158 + .../models/test_deepseek_v41_attention_ops.py | 180 + tests/models/test_deepseek_v41_contract.py | 89 + tests/models/test_deepseek_v41_dspark.py | 105 + tests/models/test_deepseek_v41_engram.py | 189 + tests/models/test_deepseek_v41_fp8_linear.py | 171 + .../test_deepseek_v41_graph_metadata.py | 100 + tests/models/test_deepseek_v41_norm.py | 46 + tests/models/test_deepseek_v41_quant_rope.py | 156 + tests/models/test_deepseek_v41_query_graph.py | 147 + .../models/test_deepseek_v41_runner_state.py | 173 + tests/quantization/test_sm70_mxfp4_moe.py | 75 +- tests/tokenizers_/test_deepseek_v41.py | 141 + .../test_deepseekv41_tool_parser.py | 57 + .../test_deepseek_v41_config.py | 104 + tests/utils/test_shared_host.py | 156 + .../spec_decode/test_deepseek_v41_sampling.py | 155 + tools/config_boundaries.py | 26 +- vllm/config/kernel.py | 32 +- vllm/config/model.py | 7 +- vllm/config/speculative.py | 17 +- .../base_device_communicator.py | 17 +- .../device_communicators/cuda_communicator.py | 13 +- vllm/envs.py | 1130 ++++++ .../kernels/deepseek_v41_compression.py | 206 ++ .../kernels/mhc/deepseek_v41.py | 466 +++ .../kernels/mhc/deepseek_v41_fused.py | 974 ++++++ .../kernels/mhc/deepseek_v41_router.py | 63 + .../layers/deepseek_v41_engram.py | 363 ++ .../layers/deepseek_v41_engram_io.py | 392 +++ .../layers/fused_moe/sm70/fp4_workspace.py | 6 +- .../layers/quantization/gguf.py | 72 +- .../layers/quantization/gguf_dense_hmma.py | 12 +- .../quantization/gguf_device_transcode.py | 54 +- .../layers/quantization/gguf_expert_cache.py | 6 +- .../quantization/gguf_expert_storage.py | 54 +- .../layers/quantization/gguf_moe.py | 99 +- .../layers/quantization/gguf_native.py | 21 + .../layers/quantization/gguf_shared_copy.py | 297 ++ .../layers/quantization/gguf_shared_layout.py | 209 ++ .../quantization/gguf_shared_prefill.py | 1034 ++++++ .../quantization/gguf_shared_protocol.py | 581 ++++ .../quantization/gguf_shared_runtime.py | 926 +++++ .../quantization/gguf_shared_storage.py | 355 ++ .../layers/quantization/gguf_turbomind_moe.py | 6 +- .../layers/quantization/gguf_v41_moe.py | 424 +++ .../layers/quantization/mxfp4_sm70_moe.py | 82 +- .../model_loader/gguf_adapters/__init__.py | 4 + .../gguf_adapters/deepseek_v41.py | 329 ++ .../model_loader/gguf_loader.py | 6 +- vllm/model_executor/models/registry.py | 8 + vllm/models/deepseek_v41/__init__.py | 3 + vllm/models/deepseek_v41/activations.py | 23 + vllm/models/deepseek_v41/attention_ops.py | 1218 +++++++ vllm/models/deepseek_v41/cache.py | 136 + vllm/models/deepseek_v41/contract.py | 56 + vllm/models/deepseek_v41/dense_tm.py | 329 ++ vllm/models/deepseek_v41/dspark.py | 415 +++ vllm/models/deepseek_v41/fp8_linear.py | 310 ++ vllm/models/deepseek_v41/index_topk.py | 343 ++ vllm/models/deepseek_v41/kv_pack.py | 156 + vllm/models/deepseek_v41/kv_store.py | 223 ++ vllm/models/deepseek_v41/markov_fused.py | 144 + vllm/models/deepseek_v41/model.py | 1083 ++++++ vllm/models/deepseek_v41/norm.py | 45 + vllm/models/deepseek_v41/quant_rope.py | 192 ++ vllm/models/deepseek_v41/query_graph.py | 194 ++ vllm/renderers/deepseek_v41.py | 23 + vllm/renderers/registry.py | 1 + vllm/third_party/deepseek_v41/LICENSE | 21 + vllm/third_party/deepseek_v41/__init__.py | 2 + vllm/third_party/deepseek_v41/encoding.py | 1098 ++++++ vllm/tokenizers/deepseek_v41.py | 83 + vllm/tokenizers/registry.py | 1 + vllm/tool_parsers/__init__.py | 4 + vllm/tool_parsers/deepseekv32_tool_parser.py | 17 +- vllm/tool_parsers/deepseekv41_tool_parser.py | 29 + vllm/transformers_utils/config.py | 2 + vllm/transformers_utils/configs/__init__.py | 4 + .../configs/deepseek_v41.py | 79 + vllm/utils/shared_host.py | 169 + vllm/v1/core/kv_cache_manager.py | 27 + vllm/v1/sample/rejection_sampler.py | 29 +- vllm/v1/spec_decode/dspark.py | 24 + vllm/v1/spec_decode/utils.py | 8 +- vllm/v1/worker/gpu_model_runner.py | 109 +- 197 files changed, 34602 insertions(+), 186 deletions(-) create mode 100644 benchmarks/deepseek_v41/HANDOFF.md create mode 100644 benchmarks/deepseek_v41/README.md create mode 100644 benchmarks/deepseek_v41/acceptance.py create mode 100644 benchmarks/deepseek_v41/adapter_audit.py create mode 100644 benchmarks/deepseek_v41/capacity.py create mode 100644 benchmarks/deepseek_v41/checkpoint_audit.py create mode 100644 benchmarks/deepseek_v41/compare_official_primitives.py create mode 100644 benchmarks/deepseek_v41/draft_checkpoint_probe.py create mode 100644 benchmarks/deepseek_v41/finite_debug.py create mode 100644 benchmarks/deepseek_v41/main_smoke.py create mode 100644 benchmarks/deepseek_v41/oracle_primitives.py create mode 100644 benchmarks/deepseek_v41/performance.py create mode 100644 benchmarks/deepseek_v41/prepare_dspark.py create mode 100644 benchmarks/deepseek_v41/resources.py create mode 100644 benchmarks/deepseek_v41/results/54633-20261024.md create mode 100644 benchmarks/deepseek_v41/serve-54633.yaml create mode 100644 benchmarks/deepseek_v41/serve-ds41-phase1.yaml create mode 100644 benchmarks/deepseek_v41/serve-ds41-phase2.yaml create mode 100644 benchmarks/deepseek_v41/serve-ds41-phase3.yaml create mode 100644 benchmarks/deepseek_v41/shared_cache_probe.py create mode 100644 benchmarks/deepseek_v41/shared_checkpoint_probe.py create mode 100644 benchmarks/deepseek_v41/shared_host_probe.py create mode 100644 benchmarks/deepseek_v41/tests/bench_g3_down.py create mode 100644 benchmarks/deepseek_v41/tests/bench_index_decode.py create mode 100644 benchmarks/deepseek_v41/tests/test_fused_sparse.py create mode 100644 benchmarks/deepseek_v41/tests/test_g1mem.py create mode 100644 benchmarks/deepseek_v41/tests/test_g3_down.py create mode 100644 benchmarks/deepseek_v41/tests/test_hc_tok.py create mode 100644 benchmarks/deepseek_v41/tests/test_index_packed.py create mode 100644 benchmarks/deepseek_v41/tests/test_kv_pack.py create mode 100644 benchmarks/deepseek_v41/tests/test_packed_attn.py create mode 100644 benchmarks/deepseek_v41/tests/test_peer_gather.py create mode 100644 benchmarks/deepseek_v41/tests/test_rope_tail.py create mode 100644 benchmarks/deepseek_v41/tests/test_topk_par.py create mode 100644 benchmarks/deepseek_v41/tests/test_window_ring.py create mode 100644 benchmarks/deepseek_v41/tools/api_multiturn.py create mode 100644 benchmarks/deepseek_v41/tools/api_smoke.py create mode 100644 benchmarks/deepseek_v41/tools/api_soak.py create mode 100644 benchmarks/deepseek_v41/tools/api_think.py create mode 100644 benchmarks/deepseek_v41/tools/decode_loop.py create mode 100644 benchmarks/deepseek_v41/tools/fresh_prompts.py create mode 100644 benchmarks/deepseek_v41/tools/longctx_test.py create mode 100755 benchmarks/deepseek_v41/tools/memlog2.sh create mode 100644 benchmarks/deepseek_v41/tools/prefill_profile.py create mode 100644 benchmarks/deepseek_v41/tools/profile_torch.py create mode 100644 benchmarks/deepseek_v41/tools/real_decode.py create mode 100644 benchmarks/deepseek_v41/tools/run_decode_spy.sh create mode 100755 benchmarks/deepseek_v41/tools/run_e2e2.sh create mode 100644 benchmarks/deepseek_v41/tools/run_longctx.sh create mode 100755 benchmarks/deepseek_v41/tools/run_prefill_profile.sh create mode 100755 benchmarks/deepseek_v41/tools/run_real.sh create mode 100644 benchmarks/deepseek_v41/tools/run_serve2.sh create mode 100644 benchmarks/deepseek_v41/tools/run_serve3.sh create mode 100644 benchmarks/deepseek_v41/tools/run_session.sh create mode 100644 benchmarks/deepseek_v41/tools/session_test.py create mode 100644 benchmarks/deepseek_v41/tools/spy_decode.py create mode 100644 benchmarks/deepseek_v41/tools/spy_top.py create mode 100644 benchmarks/deepseek_v41/tools/tcp_forward.py create mode 100644 benchmarks/deepseek_v41/tools/top_decode.py create mode 100644 benchmarks/deepseek_v41/tools/v41_e2e.py create mode 100644 benchmarks/deepseek_v41/trace.py create mode 100644 benchmarks/deepseek_v41/v41k/.gitignore create mode 100644 benchmarks/deepseek_v41/v41k/build.py create mode 100644 benchmarks/deepseek_v41/v41k/v41k.cu create mode 100644 csrc/gguf_cpu_iq_tables.h create mode 100644 csrc/gguf_shared_engram.h create mode 100644 csrc/gguf_shared_expert_cpu.h create mode 100644 csrc/gguf_shared_expert_dma.h create mode 100644 csrc/gguf_shared_expert_jobs.h create mode 100644 csrc/gguf_shared_expert_lane.h create mode 100644 csrc/gguf_shared_expert_planner.h create mode 100644 csrc/gguf_shared_expert_state.h create mode 100644 tests/benchmarks/test_deepseek_v41_acceptance.py create mode 100644 tests/benchmarks/test_deepseek_v41_checkpoint_audit.py create mode 100644 tests/benchmarks/test_deepseek_v41_oracle_primitives.py create mode 100644 tests/benchmarks/test_deepseek_v41_trace.py create mode 100644 tests/config/test_gguf_shared_original.py create mode 100644 tests/kernels/quantization/test_deepseek_v41_activations.py create mode 100644 tests/kernels/quantization/test_deepseek_v41_dense.py create mode 100644 tests/kernels/quantization/test_deepseek_v41_draft_mxfp4.py create mode 100644 tests/kernels/quantization/test_deepseek_v41_tail.py create mode 100644 tests/kernels/quantization/test_gguf_q5_batched.py create mode 100644 tests/kernels/quantization/test_gguf_q5_hmma.py create mode 100644 tests/kernels/quantization/test_gguf_q5_vector.py create mode 100644 tests/kernels/quantization/test_gguf_v41_dp4a.py create mode 100644 tests/kernels/quantization/test_gguf_v41_expert_reuse.py create mode 100644 tests/kernels/test_deepseek_v41_compression.py create mode 100644 tests/kernels/test_deepseek_v41_mhc.py create mode 100644 tests/models/gguf/shared_engram_test.cpp create mode 100644 tests/models/gguf/shared_expert_state_test.cpp create mode 100644 tests/models/gguf/test_deepseek_v41_adapter.py create mode 100644 tests/models/gguf/test_shared_acquire.py create mode 100644 tests/models/gguf/test_shared_engram_io.py create mode 100644 tests/models/gguf/test_shared_engram_native.py create mode 100644 tests/models/gguf/test_shared_expert_copy.py create mode 100644 tests/models/gguf/test_shared_expert_layout.py create mode 100644 tests/models/gguf/test_shared_expert_loading.py create mode 100644 tests/models/gguf/test_shared_expert_native.py create mode 100644 tests/models/gguf/test_shared_expert_state.py create mode 100644 tests/models/gguf/test_v41_device_codecs.py create mode 100644 tests/models/test_deepseek_v41_attention_ops.py create mode 100644 tests/models/test_deepseek_v41_contract.py create mode 100644 tests/models/test_deepseek_v41_dspark.py create mode 100644 tests/models/test_deepseek_v41_engram.py create mode 100644 tests/models/test_deepseek_v41_fp8_linear.py create mode 100644 tests/models/test_deepseek_v41_graph_metadata.py create mode 100644 tests/models/test_deepseek_v41_norm.py create mode 100644 tests/models/test_deepseek_v41_quant_rope.py create mode 100644 tests/models/test_deepseek_v41_query_graph.py create mode 100644 tests/models/test_deepseek_v41_runner_state.py create mode 100644 tests/tokenizers_/test_deepseek_v41.py create mode 100644 tests/tool_parsers/test_deepseekv41_tool_parser.py create mode 100644 tests/transformers_utils/test_deepseek_v41_config.py create mode 100644 tests/utils/test_shared_host.py create mode 100644 tests/v1/spec_decode/test_deepseek_v41_sampling.py create mode 100644 vllm/model_executor/kernels/deepseek_v41_compression.py create mode 100644 vllm/model_executor/kernels/mhc/deepseek_v41.py create mode 100644 vllm/model_executor/kernels/mhc/deepseek_v41_fused.py create mode 100644 vllm/model_executor/kernels/mhc/deepseek_v41_router.py create mode 100644 vllm/model_executor/layers/deepseek_v41_engram.py create mode 100644 vllm/model_executor/layers/deepseek_v41_engram_io.py create mode 100644 vllm/model_executor/layers/quantization/gguf_shared_copy.py create mode 100644 vllm/model_executor/layers/quantization/gguf_shared_layout.py create mode 100644 vllm/model_executor/layers/quantization/gguf_shared_prefill.py create mode 100644 vllm/model_executor/layers/quantization/gguf_shared_protocol.py create mode 100644 vllm/model_executor/layers/quantization/gguf_shared_runtime.py create mode 100644 vllm/model_executor/layers/quantization/gguf_shared_storage.py create mode 100644 vllm/model_executor/layers/quantization/gguf_v41_moe.py create mode 100644 vllm/model_executor/model_loader/gguf_adapters/deepseek_v41.py create mode 100644 vllm/models/deepseek_v41/__init__.py create mode 100644 vllm/models/deepseek_v41/activations.py create mode 100644 vllm/models/deepseek_v41/attention_ops.py create mode 100644 vllm/models/deepseek_v41/cache.py create mode 100644 vllm/models/deepseek_v41/contract.py create mode 100644 vllm/models/deepseek_v41/dense_tm.py create mode 100644 vllm/models/deepseek_v41/dspark.py create mode 100644 vllm/models/deepseek_v41/fp8_linear.py create mode 100644 vllm/models/deepseek_v41/index_topk.py create mode 100644 vllm/models/deepseek_v41/kv_pack.py create mode 100644 vllm/models/deepseek_v41/kv_store.py create mode 100644 vllm/models/deepseek_v41/markov_fused.py create mode 100644 vllm/models/deepseek_v41/model.py create mode 100644 vllm/models/deepseek_v41/norm.py create mode 100644 vllm/models/deepseek_v41/quant_rope.py create mode 100644 vllm/models/deepseek_v41/query_graph.py create mode 100644 vllm/renderers/deepseek_v41.py create mode 100644 vllm/third_party/deepseek_v41/LICENSE create mode 100644 vllm/third_party/deepseek_v41/__init__.py create mode 100644 vllm/third_party/deepseek_v41/encoding.py create mode 100644 vllm/tokenizers/deepseek_v41.py create mode 100644 vllm/tool_parsers/deepseekv41_tool_parser.py create mode 100644 vllm/transformers_utils/configs/deepseek_v41.py create mode 100644 vllm/utils/shared_host.py diff --git a/benchmarks/deepseek_v41/HANDOFF.md b/benchmarks/deepseek_v41/HANDOFF.md new file mode 100644 index 0000000000..f4d7020956 --- /dev/null +++ b/benchmarks/deepseek_v41/HANDOFF.md @@ -0,0 +1,139 @@ +# DeepSeek-V4.1-Flash-1017 on 4× V100-32G: handoff (phases 1–3) + +**Hardware:** 54633, 4× V100-SXM2 32G with NVLink. All four GPUs sit behind one x16 Gen3 uplink (12.2 GB/s aggregate). Host RAM is 62 GB; the compute container's memcg is 54 GiB. + +**Model:** +- ISTA-DASLab `DeepSeek-V4.1-Flash-1017-GSQ-RCO-IQ2_XXS` (138 GB) +- Engram IQ4_NL (110.6 GB) +- DSpark5 draft from the base repository + +## Launch + +```bash +# once: build the v41k CUDA extension (inside the ds41 container) +.venv/bin/python benchmarks/deepseek_v41/v41k/build.py # -> benchmarks/deepseek_v41/v41k/build/v41k.so + +# API with 512K context, packed FP4 KV, session prefix cache, tools, thinking, DSpark5 +docker exec -d ds41-compute-20261024 bash benchmarks/deepseek_v41/tools/run_serve3.sh serve3 +# public exposure on 54633: forward 127.0.0.1:8000 (1cat-tunnel preset DFLASH2-API) -> container :18000 +python3 benchmarks/deepseek_v41/tools/tcp_forward.py 8000 18000 +``` + +### Configs + +`serve-ds41-phase3.yaml`: +- chunk 8192, max-model-len 524288 +- KV 256 MiB (window layers only) +- offload 0.33, prefill group 4096 +- prefix caching on, tool parser `deepseek_v41`, reasoning parser `deepseek_v4` +- default `{"thinking": true}` +- `api-key` + +Requests can pass `chat_template_kwargs: {"thinking": false}` or `reasoning_effort`. + +`serve-ds41-phase2.yaml` is the 16K, non-packed variant (`run_serve2.sh`). + +`run_serve3.sh` sets these env vars: +- `VLLM_V41_DOWN_ROWS=1` +- `VLLM_V41_SESSION_PREFIX=1` +- `VLLM_V41_PACKED_KV=1` +- `VLLM_V41_INDEX_TILE_MB=32` + +## Results + +| | Phase 1 (rel1) | Now | +|---|---:|---:| +| TTFT, short prompt (≤100 tok) | 5.6 s | 0.8–1.1 s | +| TTFT 2K / 8K / 15K | 5.7 / 27.7 / 44.0 s | 4.2 / 5.8 / 11.1 s | +| TTFT 32K / 128K / 500K (packed) | n/a (16K max) | 25 / 110 / 651 s | +| Follow-up turn TTFT (session prefix hit, 4K–500K) | full prefill | 0.9–1.7 s | +| Fresh-text decode, 24 prompts (`real_decode --corpus fresh`) | 57.8 ms/round, 50.8 tok/s | 53.5 ms (55.0 packed+512K), 54.5 tok/s | +| Warm decode round | 24.9 ms | about 25 ms | +| Steady VRAM allocation | 28.6 GiB | 26.9 GiB | +| KV per token | 54 KB | 12.4 KB paged (packed: about 1 KB GPU + host) | + +Other results: +- Quality: e2e 5/5 including the 16K context check. +- At ≤16K, packed KV gives tokens identical to the paged FP16 cache. +- Needle retrieval passes at 4K, 32K, 128K and 500K. +- API soak: 80/80 OK. + +## Changes (switch, default) + +### G1: prefill activation memory +- MoE down output goes to a per-tile scratch plus an FP32 token accumulator. +- `expand_hc` runs in place. +- `hc_transition` is windowed (`VLLM_V41_HC_WINDOW`). +- DSpark aux is stored as pre-scaled FP16. +- Large all-reduces run in place (`VLLM_INPLACE_AR_MIN`). +- `empty_cache` runs before prefill passes of ≥2048 tokens (`VLLM_V41_EMPTY_CACHE_TOKENS`). +- `all_gather` goes through pynccl (`VLLM_PYNCCL_ALLGATHER`). +- Chunks up to 16384. + +### G2: cold-expert staging ring +- `VLLM_V41_STAGE_RING=1`, `VLLM_V41_STAGE_SLOTS=128`. +- Stages cold experts only, issued ahead across layers. +- The host never blocks on a pinned row. + +### G3: down projection split by rows (`VLLM_V41_DOWN_ROWS`) +- Each rank owns 1280 output rows with full K, so there are no duplicated TP boundary blocks. +- Saves 2.25 GiB per GPU and about 10% of PCIe bytes. +- Decode: `gate_up7` → IPC `peer_gather` (about 8 µs) → `down_rows`. +- Prefill: P2P h-row banks plus a per-tile NCCL fence. + +### G4: KV +- **Window rings** (`VLLM_V41_WINDOW_RING`): a 256-row private ring per target layer ≥ 1. +- **Packed KV** (`VLLM_V41_PACKED_KV`; `kv_pack.py`, `kv_store.py`), lossless: + - Compressed latents are FP4 + E4M3 scales (288 B/row). They live once in a shared host memfd (rank 0 writes) plus a 32K-position GPU hot ring. + - Index keys are FP4 + power-of-two scales (68 B/row), stored on the GPU. + - Decode: selected-row buffer, CUDA `index_scores_decode`, and segment-parallel exact top-k. + - Prefill: staging plus an unpack → cuBLAS indexer. +- **RoPE:** positions past the table rotate on the fly. +- **Limits:** C2, Engram and the guard are lifted to 512K. The DSpark draft length cap is lifted (`speculative.py`). + +### G5: prefill compute +- Engram reader threads: `VLLM_V41_ENGRAM_READERS` = 64. +- Adaptive indexer tile: `VLLM_V41_INDEX_TILE_MB`. +- Persistent-grid activation: `VLLM_V41_ACT_PROGRAMS`. + +### G6: short prompts (`VLLM_V41_ROUTED_PREFILL_MAX` = 1536) +- Only routed experts are transcoded or staged, with one sync per layer. + +### G7: same-session prefix cache (`VLLM_V41_SESSION_PREFIX`) +- The scheduler accepts a hit only when the new request continues the last executed sequence, the hit position is even, and processed − 128 ≤ hit ≤ processed. +- The worker's `restore_prefix` rewinds the Engram history and the C2 positions. + +## Tests and tools + +- **`benchmarks/deepseek_v41/tests/`**: GPU micro-tests for each change (bit-exact or equivalent). + - `test_kv_pack.py`, `test_packed_attn.py`, `test_index_packed.py`, `test_topk_par.py` + - `test_window_ring.py`, `test_g3_down.py`, `test_peer_gather.py`, `test_g1mem.py`, `test_rope_tail.py` + - Run: `PYTHONPATH=. CUDA_VISIBLE_DEVICES=0 .venv/bin/python `. +- **`benchmarks/deepseek_v41/tools/`**: + - `v41_e2e.py` / `run_e2e2.sh`: full e2e. + - `real_decode.py` / `run_real.sh`: fresh-text decode. + - `prefill_profile.py` / `run_prefill_profile.sh`: prefill segment timers, enabled with `VLLM_V41_PREFILL_TIMING=1 VLLM_V41_ATTN_TIMING=1`. + - `session_test.py`, `longctx_test.py` (needle, TTFT, follow-up decode). + - `api_multiturn.py`, `api_think.py`, `api_soak.py`. + - `run_decode_spy.sh`: py-spy. + +## Dead ends (kept, default off) + +- Fused flash-style sparse attention (`VLLM_V41_FUSED_SPARSE`). +- Token-blocked HC (`VLLM_V41_HC_TOKEN_BLOCKED`). +- Exact per-tile GEMM sizes (`VLLM_V41_EXACT_GEMM`). +- TurboMind tuning pass (`VLLM_V41_TUNE_PREFILL_PASSES`). +- Prefill group 2048. + +Triton `tl.dot` on sm70 falls back to FMA, so heavy dots need CUDA or cuBLAS. + +## Open items + +1. **Packed KV at 500K:** minimum free VRAM drops to 25 MiB during 500K prefill. + - Cause: fp16 unpack of all index keys, plus staging. + - Fix: score keys in tiles; keep staging persistent and copy only the increment. +2. **Deep-context prefill:** non-MoE time grows from 4.4 s to 10+ s per 8K chunk, from the indexer O(T·N) work and staging. +3. **Thinking-mode prefix cache:** with thinking on, `drop_thinking` removes earlier reasoning from the history, so follow-up hits land at the previous prompt end. That is outside the ring reach, so a full prefill runs. + - Fix: snapshot the window rings at each prompt end (about 10 MB). +4. **Fresh-text decode:** 27 cold-expert misses per round dominate. Lossless levers are a better hit rate, prefetch accuracy and per-miss cost. +5. **Sync scheduling only (`v41_e2e`):** about 4 ms per round of host work is not overlapped. Serve uses async. diff --git a/benchmarks/deepseek_v41/README.md b/benchmarks/deepseek_v41/README.md new file mode 100644 index 0000000000..dcf923d332 --- /dev/null +++ b/benchmarks/deepseek_v41/README.md @@ -0,0 +1,166 @@ +# DeepSeek V4.1 Flash on four V100s + +> Phases 2–3 (prefill 4x, expert dedupe, 512K packed KV, prefix cache, thinking API): +> see [HANDOFF.md](HANDOFF.md) for launch commands, switches, results and open items. + +Baseline: `6a143197e6e5f6fa77bd483df0c5a09dcea6aa20`. +Branch: `codex/ds41-flash-sm70-20261024`. + +The first profile targets V100 32 GiB ×4, TP4, one request, 16,384 total tokens +and DSpark5. The current numerical profile completed 30 minutes of stability +qualification and OpenAI API checks, including a 16,255-token input. +See the [54633 qualification report](results/54633-20261024.md) for raw-run +performance, resource peaks, the startup registration incident and boundaries. + +## Implementation and numerical boundaries + +- Independent target/draft classes live in `vllm/models/deepseek_v41/`. + The runtime contract validates geometry and rejects unsupported topology, + prefix caching and asynchronous scheduling. Strict GGUF mapping covers all + 1008 backbone tensors, including 120 routed-expert tensors. +- Dense matrices retain original packed storage, with `prefill_min_m=1`, + FP16 Tensor Core GEMM and a 32 MiB bounded dequantization workspace. + Expert M1/M6 use DP4A. M2–5 use bounded FP16 GEMMs with device-side routes. + Prefill transcodes tiles of 128 experts into one reusable workspace. + There is no full FP16 expert library or CPU matrix fallback. +- Ratio-1/2 compression, four shared compressed KV sources, top-512 selection + and delayed mHC follow the reviewed V4.1 sources. The 16K bypass skips only + candidate-block filtering. FP16 caches hold values after official FP4/FP8 QDQ. +- Norms, routing, mHC/HC residuals, C2 projections/pooling and Engram gates retain + necessary FP32 arithmetic. Draft expert projections remain FP16; routed sums + and TP reduction enter HC in FP32. A real tool query produced finite card + partials summing to about 66,019, beyond FP16 range. +- Engram has one node-level reader, row deduplication, bounded shared raw + buffers, per-card dequantization and sharded projection. It preserves the + explicit Engram pad token 2 and never replicates the complete tables. +- DSpark5 uses three 128-expert/top-3 layers and pre-block HC means at target + layers 37/38/39, after Engram. Five noncausal query rows feed six target rows. + Embedding/head storage is shared. Native FP8/MXFP4 shards are linked unchanged. +- Engram history, C2 partial groups, window and compressed KV retain candidate + state. Runner acceptance selects the input prefix including its anchor. + EOS/output-budget clipping precedes state commit. +- The independent tokenizer uses the reviewed official MIT prompt encoder. + The renderer rejects unsupported media before media I/O. The parser supports + automatic and streaming tool calls using V4.1's spaced DSML tags. + Required/named constrained tool decoding is not admitted. + +There is no full FP32 reference model. Numerical checks decode bounded weight +blocks and respect each operator's actual FP16/DP4A boundary. + +## One original host expert pool + +`shared_original` stores unsliced original GGUF bytes in one size-sealed memfd +per TP group, mapped with `MAP_SHARED`. Processes independently register the +same object with CUDA and obtain process-local device pointers. Messages carry +64-bit offsets and validate object identity, size and peer uid. + +Rank 0 alone reads original payloads and owns the native coordination thread. +Four completion sequences and reader leases fence reuse. Evictions get distinct +host destinations; whole down blocks have one writeback owner. Completed bytes +are compacted by a bounded CPU byte copy. Unlock publication does not depend on +a CUDA stream. Owner death, timeouts and inconsistent plans fail the request. + +| Offload | Cold experts/layer | Shared cold bytes | GPU slots/card | Raw GPU bytes/card | +| ---: | ---: | ---: | ---: | ---: | +| 30% | 115 | 39,553,228,800 | 277 | 26,362,255,178 | +| 25% | 96 | 33,018,347,520 | 296 | 28,170,495,818 | +| 20% | 77 | 26,483,466,240 | 315 | 29,978,736,458 | + +Slots include eight separate GPU spares. There is one group-level 512 MiB +exchange reserve. The canonical prefill workspace is 566,231,040 bytes/card. +These exclude dense weights, KV, Engram and other buffers. The first version +replaces victim slots directly; overlap is not assumed. + +Set `expert_cache_host_storage="shared_original"`, +`expert_cache_offload_fraction=0.30`, `expert_storage="original"`, +`expert_arena=true`, `expert_cache_fraction=0`, `expert_cache_host_slack=0`. +Conflicting legacy budgets and predictive-prefill controls are rejected. + +The actual-size registration probe passed with 40,090,099,712 bytes. The +checkpoint probe verifies startup bytes and repeated M6 swaps on all 40 layers. +Matching device/inode and summed object PSS establish one physical pool. +Four RSS values overcount it. Storage probes do not establish full-model fit. + +## Checkpoint and environment + +Model files are read-only under `/home/ymzx/ds41-models` on 54633: + +| Component | Bytes | Decimal GB | +| --- | ---: | ---: | +| Backbone, four GGUF shards | 138,238,893,056 | 138.24 | +| Engram GGUF including header | 110,595,290,816 | 110.60 | +| DSpark, three HF shards | 7,933,129,808 | 7.93 | +| Total | 256,767,313,680 | 256.77 | + +Routed-expert payload: 132,073,390,080 bytes. The checkpoint audit reads headers +and checks the existing SHA256 log against its manifest; it does not rerun +complete file hashing. Sources are under `/home/ymzx/ds41-dev/src`, evidence +under `/home/ymzx/ds41-dev/state`. + +GPU qualification uses `ds41-compute-20261024`, exposing only GPUs 0–3 with a +54 GiB memory limit. `ds41-dev-20261024` is GPU-hidden for CPU tests. Run Python +inside these containers through their mapped `.venv/bin/python`. + +Native `_C` and `_moe_C` were rebuilt. The unchanged GGUF dense extension keeps +its seed binary; relevant source matches the baseline and bounded decoding was +checked independently. Seed package metadata may name `705aca212`; source HEAD, +native digests and per-file SHA256 manifests identify the running implementation. + +## Reproduce qualification + +From the source directory inside the appropriate container: + +```bash +ds41_state=/home/ymzx/ds41-dev/state +.venv/bin/python -m benchmarks.deepseek_v41.checkpoint_audit \ + --model-root /home/ymzx/ds41-models --output "$ds41_state/checkpoint_audit.json" +.venv/bin/python -m benchmarks.deepseek_v41.prepare_dspark \ + --model-root /home/ymzx/ds41-models --destination /home/ymzx/ds41-dev/dspark-v41 +.venv/bin/python -m benchmarks.deepseek_v41.shared_host_probe \ + --bytes 40090099712 --output "$ds41_state/shared-host-full.json" +.venv/bin/python -m benchmarks.deepseek_v41.shared_checkpoint_probe \ + --model-root /home/ymzx/ds41-models --layers 40 --offload 0.30 \ + --output "$ds41_state/shared-checkpoint-40.json" +.venv/bin/python -m benchmarks.deepseek_v41.acceptance \ + --model-root /home/ymzx/ds41-models --chunk 512 \ + --draft /home/ymzx/ds41-dev/dspark-v41 --context-checks --force-selectors \ + --soak-seconds 1800 --output "$ds41_state/acceptance.json" +``` + +Run performance comparisons sequentially in separate engine processes. Both use +KV 256 MiB, chunk 512 and 30% offload. Raw timing has two warmups and three +samples per length, without finite/stage hooks; a separate profile records +CUDA intervals, acceptance and byte counters. + +```bash +ds41_state=/home/ymzx/ds41-dev/state +.venv/bin/python -m benchmarks.deepseek_v41.performance \ + --model-root /home/ymzx/ds41-models --output "$ds41_state/performance-no-spec.json" +.venv/bin/python -m benchmarks.deepseek_v41.performance \ + --model-root /home/ymzx/ds41-models --draft /home/ymzx/ds41-dev/dspark-v41 \ + --output "$ds41_state/performance-dspark.json" +``` + +Checks cover forced acceptance 0–5, C2/window boundaries, rejection continuation, +EOS and the 16K tail. Sampling-law tests use 32,768 synthetic chains with real +Markov proposals/rejection sampling; they are separate from checkpoint quality. +Stage totals include nested work. File/carrier bytes are logical work, not +measured PCIe transactions. Do not sum nested times or replicated group counters. + +## Serving scope + +`serve-54633.yaml` records FP16, TP4, DSpark5, synchronous eager execution, +single-request/16K limits, prefill 512, MoE sub-batches 512 and prefetch 1. +It binds `127.0.0.1:18000`. Combined target/draft chunks 2048/1024 exceeded +memory; 512 is the approved fallback with unchanged offload. Preserve at least +8 GiB host MemAvailable. + +```bash +.venv/bin/python -m vllm.entrypoints.openai.api_server \ + --config benchmarks/deepseek_v41/serve-54633.yaml +``` + +Multimodal inference, concurrency above one, lengths above 16K, asynchronous +scheduling, graph execution and other GPU/TP topologies are outside this first +profile. Full-model capacity at 25%/20% needs separate qualification even when +raw expert-storage checks pass. diff --git a/benchmarks/deepseek_v41/acceptance.py b/benchmarks/deepseek_v41/acceptance.py new file mode 100644 index 0000000000..34a29ccdfa --- /dev/null +++ b/benchmarks/deepseek_v41/acceptance.py @@ -0,0 +1,375 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Single-request V4.1 qualification, matched greedy paths and an optional soak. + +The capped draft supplies an M1 target correctness reference while retaining +the same weights. It is not a no-speculative performance baseline; run without +--draft separately for that baseline. Forced proposals test state selectors, +not the trained draft's acceptance rate. +""" + +import argparse +import ast +import json +import time +from pathlib import Path + +from benchmarks.deepseek_v41 import trace +from benchmarks.deepseek_v41.main_smoke import ( + _install_finite_checks, + _memory_report, + create_llm, + trim_allocator, +) +from benchmarks.deepseek_v41.resources import ResourceMonitor + + +def validate_commits(records, expected_end=None): + """All recurrent sources must stop at the same accepted input boundary.""" + for rank in records: + for entry in rank["commits"]: + selected = entry["rows"] if entry["prefill"] else entry["accepted_inputs"] + end = entry["start"] + selected + if ( + not 0 < selected <= entry["rows"] + or entry["hash_history"][3] != end + or entry["hash_metadata"][2] != 0 + or entry["c2_positions"] != [end] * 3 + ): + raise AssertionError((rank["rank"], entry)) + if expected_end is not None: + if not rank["commits"]: + raise AssertionError("Request has no state commits") + last = rank["commits"][-1] + if last["hash_history"][3] != expected_end: + raise AssertionError((rank["rank"], expected_end, last)) + + +def request( + llm, + prompt, + maximum, + *, + seed=41, + temperature=0, + ignore_eos=False, + skip_special_tokens=True, +): + from vllm import SamplingParams + + engine = llm.llm_engine + begin = time.monotonic() + engine.add_request( + f"v41-{time.monotonic_ns()}", + {"prompt_token_ids": prompt}, + SamplingParams( + max_tokens=maximum, + temperature=temperature, + seed=seed, + ignore_eos=ignore_eos, + skip_special_tokens=skip_special_tokens, + ), + ) + arrivals, previous, final = [], 0, None + while engine.has_unfinished_requests(): + for output in engine.step(): + now = time.monotonic() - begin + count = len(output.outputs[0].token_ids) + if count > previous: + arrivals.append(dict(seconds=now, tokens=count - previous)) + previous = count + if output.finished: + final = output + elapsed = time.monotonic() - begin + if final is None or not arrivals: + raise RuntimeError("V4.1 request produced no completion") + completion = final.outputs[0] + first, last = arrivals[0]["seconds"], arrivals[-1]["seconds"] + return dict( + started_monotonic=begin, + prompt_tokens=len(prompt), + token_ids=list(completion.token_ids), + text=completion.text, + finish_reason=completion.finish_reason, + wall_seconds=elapsed, + ttft_seconds=first, + tpot_seconds=(last - first) / (previous - 1) if previous > 1 else None, + output_tokens_per_second=previous / elapsed, + arrivals=arrivals, + ) + + +def _unfence(text): + text = text.strip() + if text.startswith("```") and text.endswith("```"): + text = "\n".join(text.splitlines()[1:-1]).strip() + return text + + +def check_quality(name, text): + text = _unfence(text) + if name == "math": + assert text == "391", text + elif name == "chat": + assert len(text) > 5 and "月" in text, text + elif name == "json": + assert json.loads(text) == {"city": "北京", "temperature": 20}, text + elif name == "code": + functions = [n for n in ast.parse(text).body if isinstance(n, ast.FunctionDef)] + function = next(n for n in functions if n.name == "add") + returned = function.body[-1] + assert len(function.args.args) == 2 and isinstance(returned, ast.Return), text + value = returned.value + assert isinstance(value, ast.BinOp) and isinstance(value.op, ast.Add), text + assert isinstance(value.left, ast.Name) and isinstance(value.right, ast.Name) + assert {value.left.id, value.right.id} == {a.arg for a in function.args.args} + + +def run(args): + begin = time.monotonic() + resources = ResourceMonitor() + llm = None + report = dict( + status="running", + profile=( + "fp16-tensorcore-dp4a-dspark5-graph" + if args.decode_graph + else "fp16-dense-tail-fp32-draft-sum" + ), + dense_prefill_min_m=1, + chunk=args.chunk, + soak_requested_seconds=args.soak_seconds, + dspark=args.draft is not None, + offload_fraction=args.offload, + correctness_checks=True, + target_graph=args.decode_graph, + draft_graph=args.draft_graph, + custom_allreduce=args.custom_ar, + cases=[], + ) + + def save(): + report["resources"] = resources.report() + temporary = args.output.with_suffix(args.output.suffix + ".tmp") + temporary.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n") + temporary.replace(args.output) + + def execute(name, prompt, maximum, **kwargs): + result = request(llm, prompt, maximum, **kwargs) + result["name"] = name + result["trace_by_rank"] = llm.collective_rpc(trace.drain) + validate_commits( + result["trace_by_rank"], len(prompt) + len(result["token_ids"]) - 1 + ) + if args.decode_graph: + for rank in result["trace_by_rank"]: + small = [row for row in rank["target_forwards"] if row["rows"] <= 6] + assert all(row["mode"] == "FULL" for row in small), (name, rank) + if args.custom_ar: + assert all( + rank["custom_allreduce_enabled"] for rank in result["trace_by_rank"] + ) + result["memory_by_rank"] = llm.collective_rpc(_memory_report) + report["cases"].append(result) + save() + print( + json.dumps( + { + k: v + for k, v in result.items() + if k not in ("trace_by_rank", "memory_by_rank", "arrivals") + }, + ensure_ascii=False, + ), + flush=True, + ) + return result + + try: + llm = create_llm( + args.model_root, + args.chunk, + args.draft, + args.offload, + decode_graph=args.decode_graph, + draft_graph=args.draft_graph, + custom_allreduce=args.custom_ar, + ) + report["load_seconds"] = time.monotonic() - begin + llm.collective_rpc(trim_allocator) + llm.collective_rpc(_install_finite_checks) + llm.collective_rpc(trace.install) + tokenizer = llm.get_tokenizer() + if args.context_checks: + question = tokenizer.encode( + "\n请用中文简短回答：2+2等于几？\n答：", add_special_tokens=False + ) + source = tokenizer.encode( + "这是一段用于验证模型的测试上下文。" * 2048, add_special_tokens=False + ) + for context in (64, 2048, 16384): + maximum = 8 if context == 64 else 32 + length = context - maximum + prompt = source[: length - len(question)] + question + result = execute( + f"context-{context}", prompt, maximum, ignore_eos=context != 64 + ) + if context != 64 and len(result["token_ids"]) != maximum: + raise AssertionError("Context-bound probe ended too early") + prompts = { + "math": "只输出计算结果：17乘以23等于多少？", + "chat": "用一句话解释为什么月亮会有阴晴圆缺。", + "code": "只输出Python代码，写一个返回两个整数之和的add函数。", + "json": "只输出有效JSON，内容是城市北京、温度20，键为city和temperature。", + } + for name, question in prompts.items(): + prompt = tokenizer.apply_chat_template( + [{"role": "user", "content": question}], + tokenize=True, + add_generation_prompt=True, + enable_thinking=False, + ) + if args.draft is not None: + llm.collective_rpc(trace.set_draft_cap, args=(0,)) + baseline = execute(name + "-target-m1", prompt, 96) + llm.collective_rpc(trace.set_draft_cap, args=(None,)) + actual = execute(name + "-dspark5", prompt, 96) + if actual["token_ids"] != baseline["token_ids"]: + raise AssertionError(f"Greedy M1/M6 divergence in {name}") + else: + actual = execute(name + "-no-spec", prompt, 96) + check_quality(name, actual["text"]) + + from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ) + from vllm.tool_parsers.deepseekv41_tool_parser import DeepSeekV41ToolParser + + question = "查询北京的天气。必须调用get_weather工具，city参数设为北京。" + tools = [ + dict( + type="function", + function=dict( + name="get_weather", + description="查询指定城市的天气", + parameters=dict( + type="object", + properties=dict(city=dict(type="string")), + required=["city"], + ), + ), + ) + ] + prompt = tokenizer.apply_chat_template( + [{"role": "user", "content": question}], + tools=tools, + tokenize=True, + add_generation_prompt=True, + enable_thinking=False, + ) + tool_result = execute("tool-format", prompt, 128, skip_special_tokens=False) + tool_request = ChatCompletionRequest( + model="ds41", messages=[dict(role="user", content=question)], tools=tools + ) + parsed = DeepSeekV41ToolParser( + tokenizer, tool_request.tools + ).extract_tool_calls(tool_result["text"], tool_request) + assert parsed.tools_called and len(parsed.tool_calls) == 1, tool_result["text"] + call = parsed.tool_calls[0].function + assert call.name == "get_weather" and json.loads(call.arguments) == { + "city": "北京" + } + + if args.draft is not None and args.force_selectors: + filler = tokenizer.encode( + "甲乙丙丁春夏秋冬。" * 64, add_special_tokens=False + ) + question = tokenizer.encode( + "\n请依次列出从一到十的中文数字：一、", add_special_tokens=False + ) + for length in (127, 128): + prompt = filler[: length - len(question)] + question + llm.collective_rpc(trace.set_draft_cap, args=(0,)) + baseline = execute(f"selector-reference-{length}", prompt, 16) + if len(baseline["token_ids"]) < 7: + raise AssertionError("Selector reference ended before seven tokens") + llm.collective_rpc(trace.set_draft_cap, args=(None,)) + for rejected in range(6): + llm.collective_rpc( + trace.force_draft, args=(baseline["token_ids"], rejected) + ) + actual = execute( + f"selector-{length}-{rejected}", prompt, rejected + 4 + ) + if actual["token_ids"] != baseline["token_ids"][: rejected + 4]: + raise AssertionError( + "Forced proposal changed the greedy output" + ) + for rank in actual["trace_by_rank"]: + decode = [c for c in rank["commits"] if not c["prefill"]] + if not decode or decode[0]["accepted_inputs"] != rejected + 1: + raise AssertionError((rank["rank"], rejected, decode)) + if len(decode) < 2: + raise AssertionError( + "No continuation after forced rejection" + ) + llm.collective_rpc(trace.force_draft, args=(None, None)) + + soak_begin = time.monotonic() + deadline = soak_begin + args.soak_seconds + iteration = 0 + while time.monotonic() < deadline: + question = f"只输出结果：{iteration % 17 + 1}+2等于多少？" + prompt = tokenizer.apply_chat_template( + [{"role": "user", "content": question}], + tokenize=True, + add_generation_prompt=True, + enable_thinking=False, + ) + actual = execute( + f"soak-{iteration}", + prompt, + 32, + seed=41 + iteration, + temperature=0.8 if iteration % 4 == 0 else 0, + ) + assert actual["text"].strip() == str(iteration % 17 + 3), actual["text"] + iteration += 1 + if resources.report()["host_minimum_available_bytes"] < 8 * 1024**3: + raise MemoryError("Acceptance breached the 8 GiB host reserve") + report.update( + status="completed", + soak_requests=iteration, + soak_elapsed_seconds=time.monotonic() - soak_begin, + ) + save() + except BaseException as exception: + report.update(status="failed", error=repr(exception)) + save() + raise + finally: + if llm is not None: + llm.llm_engine.engine_core.shutdown() + resources.close() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-root", type=Path, required=True) + parser.add_argument("--draft", type=Path) + parser.add_argument("--chunk", type=int, choices=(512, 1024, 2048), default=1024) + parser.add_argument( + "--offload", type=float, choices=(0.20, 0.25, 0.30), default=0.30 + ) + parser.add_argument("--force-selectors", action="store_true") + parser.add_argument("--decode-graph", action="store_true") + parser.add_argument("--draft-graph", action="store_true") + parser.add_argument("--custom-ar", action="store_true") + parser.add_argument("--context-checks", action="store_true") + parser.add_argument("--soak-seconds", type=int, default=0) + parser.add_argument("--output", type=Path, required=True) + run(parser.parse_args()) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/deepseek_v41/adapter_audit.py b/benchmarks/deepseek_v41/adapter_audit.py new file mode 100644 index 0000000000..cf83116ce4 --- /dev/null +++ b/benchmarks/deepseek_v41/adapter_audit.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Check the V4.1 adapter against real GGUF headers without loading weights.""" + +import argparse +import hashlib +import inspect +import json +from collections import Counter +from pathlib import Path + +import gguf + +from vllm.model_executor.model_loader.gguf_adapters.deepseek_v41 import ( + DeepseekV41Adapter, +) +from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config + + +def audit(root: Path): + config = DeepseekV41Config.from_pretrained(root / "base", local_files_only=True) + adapter = DeepseekV41Adapter(config.get_text_config(), tp_size=4) + tensors, files = {}, [] + for path in sorted((root / "gguf").glob("*.gguf")): + reader = gguf.GGUFReader(path, mode="r") + architecture = reader.get_field("general.architecture") + if architecture and architecture.contents() == "deepseek41-engram": + continue + files.append({"path": str(path), "tensors": len(reader.tensors)}) + for tensor in reader.tensors: + if tensor.name in tensors: + raise ValueError(f"Duplicate backbone tensor: {tensor.name}") + tensors[tensor.name] = tensor + mapping = adapter.build_name_map(tensors) + manifest = [ + ( + raw, + list(map(int, tensors[raw].shape)), + int(tensors[raw].tensor_type), + mapping[raw], + ) + for raw in sorted(mapping) + ] + return { + "status": "passed", + "model_type": config.model_type, + "text_model_type": config.text_config.model_type, + "backbone_files": files, + "tensor_count": len(mapping), + "unique_target_count": len(set(mapping.values())), + "expert_tensor_count": sum("_exps.weight" in raw for raw in mapping), + "source_types": dict(Counter(t.tensor_type.name for t in tensors.values())), + "mapping_sha256": hashlib.sha256(json.dumps(manifest).encode()).hexdigest(), + "adapter_sha256": hashlib.sha256( + Path(inspect.getfile(DeepseekV41Adapter)).read_bytes() + ).hexdigest(), + "scope": "Header and mapping validation only; no model forward executed", + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-root", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + result = audit(args.model_root) + args.output.write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/deepseek_v41/capacity.py b/benchmarks/deepseek_v41/capacity.py new file mode 100644 index 0000000000..8dbfed03e0 --- /dev/null +++ b/benchmarks/deepseek_v41/capacity.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Project 25/20% raw-slot increments against an observed full 30% run. + +Reads checkpoint headers and an existing resource report; allocates no GPU or +model payload. Projections are labeled explicitly and are not runtime results +for the alternative fractions. Raw storage can be qualified separately with +shared_checkpoint_probe --offload. +""" + +import argparse +import json +from pathlib import Path + +import gguf + +from vllm.model_executor.layers.quantization.gguf_shared_layout import ( + SharedExpertBudget, + SharedExpertLayout, +) + + +def capacities(root, observation): + types = {} + for path in sorted((root / "gguf").glob("*-of-00004.gguf")): + reader = gguf.GGUFReader(path) + types.update({t.name: int(t.tensor_type) for t in reader.tensors}) + geometries = [] + for layer in range(40): + geometries.append( + tuple( + SharedExpertLayout.for_type( + 5120, 2304, 4, 0, role, types[f"blk.{layer}.ffn_{name}_exps.weight"] + ) + for role, name in (("w1", "gate"), ("w3", "up"), ("w2", "down")) + ) + ) + sizes = tuple( + sum(g.source_rows * g.source_row_bytes for g in layer) for layer in geometries + ) + records = [] + for fraction in (0.30, 0.25, 0.20): + budget = SharedExpertBudget(384, fraction, sizes) + raw = 0 + for layer in geometries: + for g in layer: + raw = (raw + 255) // 256 * 256 + raw += budget.gpu_slots * g.device_rows * g.device_row_bytes + raw += g.guard_bytes + records.append( + dict( + offload_fraction=fraction, + cold_experts=budget.cold_experts, + gpu_slots=budget.gpu_slots, + shared_cold_bytes=budget.cold_bytes, + exchange_bytes=budget.exchange_bytes, + raw_gpu_bytes_per_rank=raw, + ) + ) + if observation["offload_fraction"] != 0.30 or not observation["dspark"]: + raise ValueError("Projection requires a full DSpark 30% observation") + for record in records: + delta = record["raw_gpu_bytes_per_rank"] - records[0]["raw_gpu_bytes_per_rank"] + record["raw_gpu_increment_per_rank"] = delta + record["projected_driver_peaks"] = [ + dict( + uuid=gpu["uuid"], + projected_peak_bytes=gpu["peak_used_bytes"] + delta, + total_bytes=gpu["total_bytes"], + projected_deficit_bytes=max( + 0, gpu["peak_used_bytes"] + delta - gpu["total_bytes"] + ), + ) + for gpu in observation["resources"]["gpu"] + ] + return dict( + scope="Alternative-fraction peak projections; no full model execution", + assumption="Same non-expert live allocations, workspace and driver footprint", + baseline_status=observation["status"], + baseline_context_cases=[ + c["name"] for c in observation["cases"] if c["name"].startswith("context-") + ], + fractions=records, + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-root", type=Path, required=True) + parser.add_argument("--observation", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + result = capacities(args.model_root, json.loads(args.observation.read_text())) + args.output.write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/deepseek_v41/checkpoint_audit.py b/benchmarks/deepseek_v41/checkpoint_audit.py new file mode 100644 index 0000000000..31df5fd865 --- /dev/null +++ b/benchmarks/deepseek_v41/checkpoint_audit.py @@ -0,0 +1,275 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Read-only checkpoint and original-block TP storage audit for V4.1 Flash. + +No tensor payload is materialized. Bank sizes describe the current packed TP +representation, not measured runtime memory or a recommended cache fraction. +""" + +import argparse +import json +import math +import struct +from collections import Counter +from dataclasses import dataclass +from pathlib import Path + +import gguf +import regex as re + +EXPERT = re.compile(r"blk\.(\d+)\.ffn_(gate|up|down)_exps\.weight") + + +@dataclass(frozen=True) +class Tensor: + name: str + shape: tuple[int, ...] + kind: int + nbytes: int + + +def packed_bank_bytes(tensor: Tensor, tp_size: int) -> list[int]: + """Account for shared boundary blocks and the existing final-row guard.""" + match = EXPERT.fullmatch(tensor.name) + if match is None or len(tensor.shape) != 3 or tp_size <= 0: + raise ValueError("Expected a routed expert tensor and positive TP size") + k, rows, experts = tensor.shape + block, unit = gguf.GGML_QUANT_SIZES[gguf.GGMLQuantizationType(tensor.kind)] + if min(k, rows, experts) <= 0 or k % block: + raise ValueError("Expert tensor must contain complete source blocks") + if tensor.nbytes != math.prod(tensor.shape) // block * unit: + raise ValueError("Expert byte count does not match its quantization type") + result = [] + for rank in range(tp_size): + if match[2] == "down": + if k % tp_size: + raise ValueError("Expert input width is not divisible by TP") + local = k // tp_size + first = rank * local // block + last = ((rank + 1) * local + block - 1) // block + physical_k, local_rows = (last - first) * block, rows + else: + if rows % tp_size: + raise ValueError("Expert output width is not divisible by TP") + physical_k, local_rows = k, rows // tp_size + tail = ((-physical_k) % 512) // block * unit + result.append(experts * local_rows * (physical_k // block) * unit + tail) + return result + + +def verified_paths(manifest: str, log: str) -> set[str]: + expected = set() + for line in manifest.splitlines(): + digest, name = line.split(maxsplit=1) + name = name.lstrip(" *") + if re.fullmatch(r"[0-9a-fA-F]{64}", digest) is None: + raise ValueError("Invalid SHA256 manifest") + if Path(name).is_absolute() or ".." in Path(name).parts: + raise ValueError("Manifest path escapes checkpoint directory") + expected.add(name) + exits = re.findall(r"^EXIT=(\d+)$", log, re.MULTILINE) + actual = { + match[1] for match in re.finditer(r"^(.+): (?:OK|成功)$", log, re.MULTILINE) + } + if not expected or not exits or exits[-1] != "0" or actual != expected: + raise ValueError("Checkpoint SHA256 verification is incomplete or failed") + return expected + + +def safetensors_names(path: Path) -> dict[str, dict]: + with path.open("rb") as source: + prefix = source.read(8) + if len(prefix) != 8: + raise ValueError(f"Truncated safetensors header: {path}") + size = struct.unpack(" 16 * 1024**2 or 8 + size > path.stat().st_size: + raise ValueError(f"Invalid safetensors header size: {path}") + header = json.loads(source.read(size)) + tensors = {name: value for name, value in header.items() if name != "__metadata__"} + ranges = [] + for value in tensors.values(): + start, end = value["data_offsets"] + if not 0 <= start <= end <= path.stat().st_size - 8 - size: + raise ValueError(f"Safetensors payload is out of bounds: {path}") + ranges.append((start, end)) + ordered = sorted(ranges) + if any(left[1] > right[0] for left, right in zip(ordered, ordered[1:])): + raise ValueError(f"Overlapping safetensors payloads: {path}") + return tensors + + +def shard_header(reader: gguf.GGUFReader) -> dict: + def value(name): + field = reader.get_field(name) + return None if field is None else field.contents() + + architecture = value("general.architecture") + index, count = value("split.no"), value("split.count") + if index is None and count is None: + if architecture != "deepseek41-engram": + raise ValueError("Expected the standalone Engram GGUF") + else: + if index is None or count is None or not 0 <= index < count: + raise ValueError("Invalid GGUF split metadata") + if (index == 0 and architecture != "deepseek41") or ( + index > 0 and architecture not in (None, "deepseek41") + ): + raise ValueError("Unexpected backbone GGUF architecture") + return { + "architecture": architecture, + "split_index": index, + "split_count": count, + "split_tensor_count": value("split.tensors.count"), + } + + +def audit(root: Path, tp_size: int) -> dict: + verified = verified_paths( + (root / "SHA256SUMS").read_text(), (root / "verify.log").read_text() + ) + if len(verified) != 8: + raise ValueError("Expected four backbone, one Engram and three DSpark files") + config = json.loads((root / "base/config.json").read_text()) + if config["model_type"] != "deepseek_v41": + raise ValueError("Expected a DeepSeek V4.1 checkpoint") + text = config["text_config"] + layers = text["num_hidden_layers"] + tensors: dict[str, Tensor] = {} + file_records = [] + padding_ids = set() + for path in sorted((root / "gguf").glob("*.gguf")): + reader = gguf.GGUFReader(path, mode="r") + header = shard_header(reader) + padding = reader.get_field("tokenizer.ggml.padding_token_id") + if padding is not None: + padding_ids.add(int(padding.contents())) + intervals = [] + for item in reader.tensors: + if item.name in tensors: + raise ValueError(f"Duplicate tensor across shards: {item.name}") + if not reader.data_offset <= item.data_offset: + raise ValueError(f"Tensor overlaps GGUF header: {item.name}") + end = item.data_offset + item.n_bytes + if end > path.stat().st_size: + raise ValueError(f"Tensor is outside GGUF file: {item.name}") + intervals.append((item.data_offset, end)) + tensors[item.name] = Tensor( + item.name, + tuple(map(int, item.shape)), + int(item.tensor_type), + int(item.n_bytes), + ) + ordered = sorted(intervals) + if any(left[1] > right[0] for left, right in zip(ordered, ordered[1:])): + raise ValueError(f"Overlapping GGUF tensor payloads: {path}") + file_records.append({"name": path.name, "bytes": path.stat().st_size, **header}) + splits = [item for item in file_records if item["split_index"] is not None] + if ( + len(splits) != 4 + or {item["split_count"] for item in splits} != {4} + or {item["split_index"] for item in splits} != set(range(4)) + ): + raise ValueError("Expected the complete four-part backbone GGUF") + + expert_bytes = 0 + rank_bytes = [0] * tp_size + layer_records = [] + for layer in range(layers): + types = {} + for role in ("gate", "up", "down"): + name = f"blk.{layer}.ffn_{role}_exps.weight" + tensor = tensors[name] + if tensor.kind not in (16, 17, 29): + raise ValueError(f"Unexpected routed expert quantization: {name}") + expected = (text["hidden_size"], text["moe_intermediate_size"]) + if role == "down": + expected = expected[::-1] + if tensor.shape != (*expected, text["n_routed_experts"]): + raise ValueError(f"Expert geometry does not match HF config: {name}") + expert_bytes += tensor.nbytes + types[role] = gguf.GGMLQuantizationType(tensor.kind).name + rank_bytes = [ + a + b for a, b in zip(rank_bytes, packed_bank_bytes(tensor, tp_size)) + ] + layer_records.append({"layer": layer, "types": types}) + if sum(EXPERT.fullmatch(name) is not None for name in tensors) != layers * 3: + raise ValueError("Unexpected routed expert tensor count") + + engram_bytes = 0 + for layer, rows in zip( + text["engram_layer_ids"], text["engram_num_embeddings"], strict=True + ): + tensor = tensors[f"blk.{layer}.engram_embd.weight"] + if tensor.shape != (text["engram_head_dim"], rows) or tensor.kind != 20: + raise ValueError("Engram table does not match HF config or IQ4_NL") + engram_bytes += tensor.nbytes + + expected_mtp = { + name: file + for name, file in json.loads( + (root / "base/model.safetensors.index.json").read_text() + )["weight_map"].items() + if name.startswith("mtp.") + } + seen_mtp = {} + dtypes: Counter[str] = Counter() + for path in sorted((root / "base/dspark").glob("*.safetensors")): + for name, value in safetensors_names(path).items(): + if name in seen_mtp: + raise ValueError(f"Duplicate DSpark tensor: {name}") + seen_mtp[name] = path.name + dtypes[value["dtype"]] += 1 + if seen_mtp != expected_mtp: + raise ValueError("DSpark shards do not exactly cover the HF MTP index") + return { + "sha256_verified_files": sorted(verified), + "files": file_records, + "tensor_count": len(tensors), + "tp_size": tp_size, + "expert_source_bytes": expert_bytes, + "expert_packed_bank_bytes_per_rank": rank_bytes, + "expert_tp_boundary_overhead_bytes": sum(rank_bytes) - expert_bytes, + "cold_expert_storage_estimates": { + str(fraction): int(sum(rank_bytes) * (1 - fraction)) + for fraction in (0.65, 0.70, 0.75, 0.80) + }, + "cold_estimate_assumptions": ( + "Ideal fraction only; excludes spare slots, host slack, arena " + "alignment, load staging and prefill workspace" + ), + "expert_layers": layer_records, + "engram_table_bytes": engram_bytes, + "hf_pad_token_id": config["pad_token_id"], + "gguf_tokenizer_padding_ids": sorted(padding_ids), + "dspark_tensor_count": len(seen_mtp), + "dspark_dtypes": dict(dtypes), + "dspark_block_size": text["dspark_block_size"], + "dspark_layers": text["num_nextn_predict_layers"], + "memory_note": ( + "Packed layout estimate; excludes runtime buffers and measured allocation." + ), + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-root", type=Path, required=True) + parser.add_argument("--tp-size", type=int, default=4) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + report = audit(args.model_root, args.tp_size) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2) + "\n") + print( + json.dumps( + { + key: value + for key, value in report.items() + if key not in ("expert_layers", "files") + } + ) + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/deepseek_v41/compare_official_primitives.py b/benchmarks/deepseek_v41/compare_official_primitives.py new file mode 100644 index 0000000000..ea608ff66f --- /dev/null +++ b/benchmarks/deepseek_v41/compare_official_primitives.py @@ -0,0 +1,212 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compare oracle building blocks with a local official V4.1 reference. + +Run on CPU with --reference-dir pointing to the reviewed official inference +sources. This validates components only; it does not generate model logits. +""" + +import argparse +import ast +import hashlib +import json +import resource +import runpy +from pathlib import Path +from types import SimpleNamespace + +import gguf +import torch +import torch.nn.functional as F +from transformers import AutoTokenizer + +from benchmarks.deepseek_v41.oracle_primitives import ( + EngramHashState, + GGUFWeights, + engram_update, + rms_norm, + route, +) + + +def reference_method(path: Path, class_name: str, method_name: str): + """Load one reviewed method without importing vision or GPU-only modules.""" + tree = ast.parse(path.read_text()) + cls = next( + node + for node in tree.body + if isinstance(node, ast.ClassDef) and node.name == class_name + ) + method = next( + node + for node in cls.body + if isinstance(node, ast.FunctionDef) and node.name == method_name + ) + scope = {"torch": torch, "F": F, "linear": F.linear} + module = ast.Module(body=[method], type_ignores=[]) + exec(compile(module, str(path), "exec"), scope) + return scope[method_name] + + +def compare(root: Path, reference_dir: Path) -> dict: + torch.set_num_threads(4) + torch.manual_seed(41) + config = json.loads((root / "base/config.json").read_text())["text_config"] + args = SimpleNamespace( + **{key: value for key, value in config.items() if key.startswith("engram_")}, + max_batch_size=2, + max_seq_len=2056, + ) + args.engram_pad_id = args.engram_pad_token_id + official = runpy.run_path(str(reference_dir / "engram.py")) + layout = official["EngramLayout"].from_args(args) + tokenizer = AutoTokenizer.from_pretrained(root / "base", local_files_only=True) + expected = official["NgramHashState"](args, layout, tokenizer) + weights = GGUFWeights(sorted((root / "gguf").glob("*.gguf"))) + token_map = weights.integer("engram_token_map.weight") + primes = weights.integer("engram_primes.weight") + offsets = weights.integer("engram_offsets.weight") + multipliers = weights.integer("engram_multipliers.weight") + for name, actual in ( + ("token_map", token_map), + ("primes", primes.reshape_as(expected.primes)), + ("offsets", offsets), + ("multipliers", multipliers), + ): + torch.testing.assert_close(actual, getattr(expected, name), rtol=0, atol=0) + + def new_hash(): + return EngramHashState( + token_map, primes, offsets, multipliers, args.engram_pad_token_id + ) + + hash_cases = [] + for length in (1, 2, 3, 4, 64, 2048): + ids = torch.randint(len(token_map), (2, length + 3)) + mask = torch.ones_like(ids, dtype=torch.bool) + mask[0, ::7] = False + if length > 1: + mask[1, length - 1] = False + reference = expected(ids, 0, mask) + full = new_hash()(ids, 0, mask) + torch.testing.assert_close(full, reference, rtol=0, atol=0) + # Split at prefill/decode and within the 4-gram lookback window. + for split in sorted({1, min(3, length), length}): + state = new_hash() + chunks = [state(ids[:, :split], 0, mask[:, :split])] + for position in range(split, ids.shape[1]): + chunks.append( + state( + ids[:, position : position + 1], + position, + mask[:, position : position + 1], + ) + ) + torch.testing.assert_close( + torch.cat(chunks, dim=1), reference, rtol=0, atol=0 + ) + hash_cases.append({"prefill_tokens": length, "decode_tokens": 3}) + + model = reference_dir / "model.py" + norm = reference_method(model, "RMSNorm", "forward") + x = torch.randn(2, 7, 16) * 1e-12 + norm_weight = torch.randn(16) + torch.testing.assert_close( + rms_norm(x, norm_weight, 1e-20), + norm(SimpleNamespace(weight=norm_weight, eps=1e-20), x), + rtol=1e-6, + atol=1e-7, + ) + router = reference_method(model, "Gate", "forward") + router_state = SimpleNamespace( + weight=torch.randn(384, 16), + bias=torch.randn(384), + bias_vl=None, + topk=6, + score_func="sqrtsoftplus", + gate_temp=1.0, + norm_topk_prob=True, + route_scale=1.5, + ) + x = torch.randn(13, 16) + actual_weights, actual_ids = route(x, router_state.weight, router_state.bias) + expected_weights, expected_ids = router(router_state, x) + torch.testing.assert_close(actual_ids, expected_ids, rtol=0, atol=0) + torch.testing.assert_close(actual_weights, expected_weights, rtol=0, atol=0) + + update = reference_method(model, "Engram", "forward") + gate_errors = [] + for factor in (0.0, 1e-12, 1.0, 1e6): + hidden = torch.randn(2, 5, 4, 16) * factor + kv = torch.randn(2, 5, 80) * factor + q, k = torch.randn(4, 16), torch.randn(4, 16) + mask = torch.tensor([[True, False, True, False, True]]).expand(2, -1) + state = SimpleNamespace( + wkv=lambda value: value, + embed=lambda value, kv=kv: kv.unsqueeze(-2), + hc_mult=4, + dim=16, + q_weight=q, + k_weight=k, + eps=1e-20, + clamp_value=1e-6, + ) + reference = update(state, hidden, torch.zeros(2, 5, 24), mask) + actual = engram_update(hidden, kv, q, k, mask) + torch.testing.assert_close(actual, reference, rtol=1e-6, atol=1e-7) + assert torch.isfinite(actual).all() + assert torch.equal(actual[~mask], hidden[~mask]) + gate_errors.append(float((actual - reference).abs().max())) + + # Real IQ4_NL table rows plus chunked Q5_K projection for both Engram layers. + real_rows = [] + ids = torch.tensor([[0, 2]]) + hashes = new_hash()(ids, 0) + for index, layer in enumerate(args.engram_layer_ids): + name = f"blk.{layer}.engram_embd.weight" + row_ids = hashes[:, :, index] + embeddings = weights.rows(name, row_ids) + tensor = weights.tensors[name] + for position in range(2): + for head in range(24): + row = int(row_ids[0, position, head]) + raw = tensor.data[row : row + 1] + independent = torch.from_numpy( + gguf.quants.dequantize(raw, tensor.tensor_type).copy() + ).flatten() + torch.testing.assert_close( + embeddings[0, position, head], independent, rtol=0, atol=0 + ) + kv = weights.linear(f"blk.{layer}.engram_wkv.weight", embeddings.flatten(-2)) + assert kv.shape == (1, 2, 25600) and torch.isfinite(kv).all() + real_rows.append({"layer": layer, "lookup_rows": row_ids.numel()}) + return { + "status": "component_checks_passed", + "scope": "M1 building blocks only; no full-model logits or text validation", + "reference_sha256": { + name: hashlib.sha256((reference_dir / name).read_bytes()).hexdigest() + for name in ("model.py", "engram.py") + }, + "checkpoint_hash_auxiliaries_match_official": True, + "hash_cases": hash_cases, + "rms_norm_and_router_match_official": True, + "engram_gate_max_absolute_errors": gate_errors, + "real_engram_row_and_projection_checks": real_rows, + "peak_rss_kib": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss, + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-root", required=True, type=Path) + parser.add_argument("--reference-dir", required=True, type=Path) + parser.add_argument("--output", required=True, type=Path) + args = parser.parse_args() + with torch.inference_mode(): + result = compare(args.model_root, args.reference_dir) + args.output.write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/deepseek_v41/draft_checkpoint_probe.py b/benchmarks/deepseek_v41/draft_checkpoint_probe.py new file mode 100644 index 0000000000..bda12dd639 --- /dev/null +++ b/benchmarks/deepseek_v41/draft_checkpoint_probe.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Load the real DSpark checkpoint on TP4 without loading the target weights. + +This qualifies weight coverage and draft dense/MoE execution, not speculative +generation. Attention correctness has separate paged, non-causal component tests. +""" + +import argparse +import json +import socket +import time +from pathlib import Path + +import torch +import torch.multiprocessing as mp + + +def _worker(rank, address, root, draft, output): + from vllm.config import set_current_vllm_config + from vllm.distributed import ( + cleanup_dist_env_and_memory, + init_distributed_environment, + initialize_model_parallel, + ) + from vllm.engine.arg_utils import EngineArgs + from vllm.forward_context import set_forward_context + from vllm.model_executor.model_loader import get_model + + torch.accelerator.set_device_index(rank) + model_path = next((root / "gguf").glob("*-00001-of-00004.gguf")) + config = EngineArgs( + model=str(model_path), + hf_config_path=str(root / "base"), + tokenizer=str(root / "base"), + dtype="half", + quantization="gguf", + load_format="gguf", + tensor_parallel_size=4, + max_model_len=16384, + max_num_seqs=1, + max_num_batched_tokens=1024, + enforce_eager=True, + enable_prefix_caching=False, + disable_custom_all_reduce=True, + speculative_config={ + "method": "dspark", + "model": str(draft), + "quantization": "fp8", + "num_speculative_tokens": 5, + "enforce_eager": True, + "draft_load_config": {"load_format": "safetensors"}, + }, + kernel_config={"sm70_ring": {"enabled": False}}, + ).create_engine_config() + started = time.monotonic() + with set_current_vllm_config(config): + init_distributed_environment(4, rank, address, rank, backend="nccl") + initialize_model_parallel(4) + spec = config.speculative_config + model = get_model( + vllm_config=config, + model_config=spec.draft_model_config, + load_config=spec.draft_load_config, + prefix="eagle_head", + ) + loaded = time.monotonic() - started + with torch.inference_mode(), torch.device("cuda", rank): + torch.manual_seed(415) + auxiliary = torch.randn(64, 15360, dtype=torch.float16) * 0.01 + context = model.combine_hidden_states(auxiliary) + positions = torch.arange(64) + model.precompute_and_store_context_kv(context, positions) + embeddings = torch.randn(5, 5120, dtype=torch.float16) * 0.01 + with set_forward_context(None, config, num_tokens=5): + hidden = model( + torch.full((5,), 128799), torch.arange(64, 69), embeddings + ) + transition = model.markov_bias(model.markov_embed(torch.tensor([22]))) + confidence = model.confidence_logits( + hidden, model.markov_embed(torch.full((5,), 22)) + ) + for value in (context, hidden, transition, confidence): + if not torch.isfinite(value).all(): + raise ValueError("Nonfinite values in native V4.1 draft probe") + torch.accelerator.synchronize(rank) + report = dict( + rank=rank, + passed=True, + scope="Draft weights/dense/MoE only; speculative acceptance untested", + load_seconds=loaded, + allocated_bytes=torch.accelerator.memory_allocated(rank), + peak_bytes=torch.accelerator.max_memory_allocated(rank), + hidden_abs_max=float(hidden.abs().max()), + markov_abs_max=float(transition.abs().max()), + ) + (output / f"draft-checkpoint-rank{rank}.json").write_text( + json.dumps(report, indent=2) + "\n" + ) + torch.distributed.barrier() + del model + cleanup_dist_env_and_memory() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-root", type=Path, required=True) + parser.add_argument("--draft", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + args.output.mkdir(parents=True, exist_ok=True) + with socket.socket() as listener: + listener.bind(("127.0.0.1", 0)) + port = listener.getsockname()[1] + mp.spawn( + _worker, + args=(f"tcp://127.0.0.1:{port}", args.model_root, args.draft, args.output), + nprocs=4, + join=True, + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/deepseek_v41/finite_debug.py b/benchmarks/deepseek_v41/finite_debug.py new file mode 100644 index 0000000000..4f514f8cff --- /dev/null +++ b/benchmarks/deepseek_v41/finite_debug.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Locate the first nonfinite module boundary in a real FP16 target forward.""" + +import json + +import torch +import torch.distributed as dist + + +def install(worker): + from vllm.distributed import get_tp_group + + model = worker.model_runner.model + group = get_tp_group().device_group + context = {} + + def inputs(module, args, kwargs): + del module, args + ids, positions = kwargs["input_ids"], kwargs["positions"] + context.update( + rows=ids.numel(), + start=int(positions[0]), + end=int(positions[-1]), + tokens=ids.tolist() if ids.numel() <= 6 else ids[-8:].tolist(), + ) + + model.register_forward_pre_hook(inputs, with_kwargs=True) + + def check(name): + def hook(module, args, output): + del module, args + tensors = output if isinstance(output, tuple) else (output,) + for tensor in tensors: + if ( + not isinstance(tensor, torch.Tensor) + or not tensor.is_floating_point() + ): + continue + finite = torch.isfinite(tensor).all().to(torch.int32) + # Stop every rank at the same boundary, before a peer can + # enter the next collective or shared-memory wait. + dist.all_reduce(finite, op=dist.ReduceOp.MIN, group=group) + if not finite.item(): + detail = dict( + module=name, + shape=list(tensor.shape), + dtype=str(tensor.dtype), + context=context, + nan=int(torch.isnan(tensor).sum()), + inf=int(torch.isinf(tensor).sum()), + ) + print("V41_FIRST_NONFINITE " + json.dumps(detail), flush=True) + raise FloatingPointError(json.dumps(detail)) + + return hook + + for name, module in model.named_modules(): + if name: + module.register_forward_hook(check(name)) diff --git a/benchmarks/deepseek_v41/main_smoke.py b/benchmarks/deepseek_v41/main_smoke.py new file mode 100644 index 0000000000..38823ba167 --- /dev/null +++ b/benchmarks/deepseek_v41/main_smoke.py @@ -0,0 +1,251 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Real target/DSpark generation and finite-logit correctness probe.""" + +import argparse +import json +import os +import time +from pathlib import Path + + +def _install_finite_checks(worker): + import torch + + runner = worker.model_runner + models = [runner.model] + drafter = getattr(runner, "drafter", None) + if drafter is not None: + models.append(drafter.model) + for module in models: + original = module.compute_logits + + def checked(hidden, original=original): + logits = original(hidden) + if logits is not None: + finite = torch.isfinite(logits).all() + if logits.is_cuda: + # Correctness probes can wrap the head before its lazy + # graph capture. A Python bool would synchronize that + # capture; the device assertion is checked on drain. + torch._assert_async(finite, "Nonfinite V4.1 correctness logits") + elif not finite: + raise FloatingPointError( + "Nonfinite logits in V4.1 correctness probe" + ) + return logits + + module.compute_logits = checked + + +def trim_allocator(worker): + import torch + + torch.accelerator.synchronize(worker.device) + # Release dead profiling allocations before any lazy cubin loads. The + # raw experts, canonical workspace, KV and loaded parameters stay live. + torch.accelerator.empty_cache() + + +def _memory_report(worker): + import torch + + from benchmarks.deepseek_v41.shared_checkpoint_probe import _pss + from benchmarks.deepseek_v41.shared_host_probe import memory_available + + device = worker.device + target = worker.model_runner.model.model + cache = target.layers[0].ffn.experts.quant_method.expert_cache + counts = cache.runtime.statistics.cpu().tolist() + loads = writes = 0 + for spec, (_, misses) in zip(cache.specs, counts): + for geometry in spec.layouts(cache.rank): + loads += misses * geometry.device_rows * geometry.device_row_bytes + writes += ( + misses + * geometry.device_rows + * (geometry.write_end_block - geometry.first_block) + * geometry.block_bytes + ) + return dict( + device=str(device), + allocated_bytes=torch.accelerator.memory_allocated(device), + peak_allocated_bytes=torch.accelerator.max_memory_allocated(device), + reserved_bytes=torch.accelerator.memory_reserved(device), + host_available_bytes=memory_available(), + expert_pool_identity=vars(cache.data.identity), + expert_pool_pss_bytes=_pss(), + expert_decode_requests_and_misses_by_layer=counts, + expert_decode_load_bytes=loads, + expert_decode_writeback_bytes=writes, + prefill_cold_transcodes_by_layer=cache.tm.cold_transcodes.cpu().tolist(), + expert_workspace_bytes=cache.tm.storage_bytes(), + engram_requests_rows_bytes_reader_ns=target.engram_io.runtime.statistics.cpu().tolist(), + ) + + +def create_llm( + root, + chunk=1024, + draft=None, + offload=0.30, + *, + kv_cache_mib=None, + decode_graph=False, + draft_graph=False, + custom_allreduce=False, +): + os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + # Developer-only callable RPCs run in this private local worker group. + os.environ.setdefault("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") + from vllm import LLM + + paths = sorted((root / "gguf").glob("*-00001-of-00004.gguf")) + sidecar = root / "gguf/DeepSeek-V4.1-Flash-Engram-IQ4_NL.gguf" + if len(paths) != 1 or not sidecar.is_file(): + raise ValueError("Expected the audited backbone and explicit Engram sidecar") + if kv_cache_mib is None: + kv_cache_mib = 256 if draft is not None else 384 + return LLM( + model=str(paths[0]), + hf_config_path=str(root / "base"), + tokenizer=str(root / "base"), + dtype="half", + quantization="gguf", + load_format="gguf", + tensor_parallel_size=4, + max_model_len=int(os.environ.get("VLLM_V41_MAX_LEN", "16384")), + max_num_seqs=1, + max_num_batched_tokens=chunk, + model_loader_extra_config={"engram_path": str(sidecar)}, + gpu_memory_utilization=0.96, + kv_cache_memory_bytes=kv_cache_mib * 1024**2, + enforce_eager=not decode_graph, + compilation_config=( + { + "mode": 0, + "cudagraph_mode": "FULL", + "cudagraph_capture_sizes": [1, 2, 3, 4, 5, 6], + "max_cudagraph_capture_size": 6, + } + if decode_graph + else None + ), + enable_prefix_caching=os.environ.get("VLLM_V41_SESSION_PREFIX", "0") == "1", + enable_chunked_prefill=True, + async_scheduling=os.environ.get("VLLM_V41_ASYNC", "0") == "1", + disable_custom_all_reduce=not custom_allreduce, + speculative_config=( + { + "method": "dspark", + "model": str(draft), + "quantization": "fp8", + "num_speculative_tokens": 5, + "draft_load_config": {"load_format": "safetensors"}, + "enforce_eager": not draft_graph, + } + if draft is not None + else None + ), + kernel_config={ + "sm70_ring": {"enabled": False}, + "sm70_gguf": { + "expert_storage": "original", + "expert_arena": True, + "expert_cache_host_storage": "shared_original", + "expert_cache_offload_fraction": offload, + "expert_cache_fraction": 0, + "expert_cache_host_slack": 0, + "expert_load_prefetch_banks": 1, + "expert_cache_prefill_tokens": chunk, + "dense_storage": "original", + "embedding_storage": "original", + "prefill_min_m": 1, + "dequant_workspace_bytes": 32 * 1024**2, + }, + }, + ) + + +def run(root, contexts, chunk, output, draft=None, debug_finite=False): + from vllm import SamplingParams + + started = time.monotonic() + llm = create_llm(root, chunk, draft) + load_seconds = time.monotonic() - started + llm.collective_rpc(trim_allocator) + if debug_finite: + from benchmarks.deepseek_v41.finite_debug import install + + llm.collective_rpc(install) + llm.collective_rpc(_install_finite_checks) + tokenizer = llm.get_tokenizer() + question = tokenizer.encode( + "\n请用中文简短回答：2+2等于几？\n答：", add_special_tokens=False + ) + source = tokenizer.encode( + "这是一段用于验证模型的测试上下文。" * 2048, add_special_tokens=False + ) + reports = [] + for context in contexts: + new_tokens = 8 if context == 64 else 32 + count = context - new_tokens + prompt = source[: count - len(question)] + question + if len(prompt) != count: + raise ValueError("Probe prompt cannot fill requested context") + begin = time.monotonic() + results = llm.generate( + [{"prompt_token_ids": prompt}], + # Longer probes exercise the exact total-context boundary, + # including positions after an early EOS. Quality is checked + # separately with the ordinary EOS stopping policy. + SamplingParams( + temperature=0, max_tokens=new_tokens, ignore_eos=context != 64 + ), + ) + elapsed = time.monotonic() - begin + generated = results[0].outputs[0] + report = dict( + status="completed", + scope="Context-bound execution; not a steady-state performance result", + dspark=draft is not None, + context_limit=context, + prompt_tokens=len(prompt), + output_tokens=len(generated.token_ids), + token_ids=list(generated.token_ids), + text=generated.text, + finish_reason=generated.finish_reason, + load_seconds=load_seconds, + generation_wall_seconds=elapsed, + tokens_per_second_including_prefill=len(generated.token_ids) / elapsed, + memory_by_rank=llm.collective_rpc(_memory_report), + ) + reports.append(report) + output.write_text(json.dumps(reports, ensure_ascii=False, indent=2) + "\n") + print(json.dumps(report, ensure_ascii=False, indent=2), flush=True) + llm.llm_engine.engine_core.shutdown() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-root", type=Path, required=True) + parser.add_argument( + "--context", type=int, choices=(64, 2048, 16384), nargs="+", default=[64] + ) + parser.add_argument("--chunk", type=int, choices=(512, 1024, 2048), default=2048) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--draft", type=Path) + parser.add_argument("--debug-finite", action="store_true") + args = parser.parse_args() + run( + args.model_root, + args.context, + args.chunk, + args.output, + args.draft, + args.debug_finite, + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/deepseek_v41/oracle_primitives.py b/benchmarks/deepseek_v41/oracle_primitives.py new file mode 100644 index 0000000000..a45312e0e2 --- /dev/null +++ b/benchmarks/deepseek_v41/oracle_primitives.py @@ -0,0 +1,185 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU reference building blocks for the V4.1 GGUF oracle. + +These are validation helpers, not a registered model or a complete M1 oracle. +Weights remain memory mapped; only requested rows are dequantized to FP32. +""" + +from collections.abc import Iterable +from pathlib import Path + +import gguf +import numpy as np +import torch +import torch.nn.functional as F + + +class GGUFWeights: + def __init__(self, paths: Iterable[Path]): + self.readers = [gguf.GGUFReader(path, mode="r") for path in paths] + self.tensors = {} + for reader in self.readers: + for tensor in reader.tensors: + if tensor.name in self.tensors: + raise ValueError(f"Duplicate tensor: {tensor.name}") + self.tensors[tensor.name] = tensor + + def integer(self, name: str) -> torch.Tensor: + tensor = self.tensors[name] + if tensor.tensor_type not in ( + gguf.GGMLQuantizationType.I32, + gguf.GGMLQuantizationType.I64, + ): + raise ValueError(f"Expected integer tensor: {name}") + return torch.from_numpy(np.array(tensor.data, dtype=np.int64, copy=True)) + + def rows(self, name: str, row_ids: torch.Tensor) -> torch.Tensor: + """Read flattened outer rows, preserving duplicates and caller shape.""" + tensor = self.tensors[name] + ids = row_ids.to(device="cpu", dtype=torch.int64).numpy() + width = int(tensor.shape[0]) + row_count = int(np.prod(tensor.shape[1:])) + if np.any(ids < 0) or np.any(ids >= row_count): + raise ValueError(f"Row outside tensor: {name}") + unique, inverse = np.unique(ids, return_inverse=True) + packed = tensor.data.reshape(row_count, -1)[unique] + values = gguf.quants.dequantize(packed, tensor.tensor_type) + values = np.asarray(values, dtype=np.float32).reshape(len(unique), width) + return torch.from_numpy(values[inverse].reshape(*ids.shape, width)) + + @torch.inference_mode() + def linear( + self, name: str, x: torch.Tensor, rows_per_chunk: int = 256 + ) -> torch.Tensor: + tensor = self.tensors[name] + if len(tensor.shape) != 2 or x.shape[-1] != tensor.shape[0]: + raise ValueError(f"Linear shape mismatch: {name}") + if rows_per_chunk <= 0: + raise ValueError("rows_per_chunk must be positive") + x = x.to(device="cpu", dtype=torch.float32) + out_features = int(tensor.shape[1]) + result = torch.empty((*x.shape[:-1], out_features), dtype=torch.float32) + for start in range(0, out_features, rows_per_chunk): + end = min(start + rows_per_chunk, out_features) + weight = self.rows(name, torch.arange(start, end)) + result[..., start:end] = F.linear(x, weight) + return result + + +class EngramHashState: + """Hash from checkpoint auxiliaries with only max_ngram_size-1 cached ids. + + A masked token blocks lookback across its position. Calls must be contiguous; + start_pos=0 explicitly starts a new sequence, as in the official reference. + """ + + DEAD = -1 + + def __init__( + self, + token_map: torch.Tensor, + primes: torch.Tensor, + offsets: torch.Tensor, + multipliers: torch.Tensor, + pad_token_id: int, + ): + self.token_map = token_map.to(device="cpu", dtype=torch.int64) + self.primes = primes.to(device="cpu", dtype=torch.int64) + self.offsets = offsets.to(device="cpu", dtype=torch.int64) + self.multipliers = multipliers.to(device="cpu", dtype=torch.int64) + self.pad_id = int(self.token_map[pad_token_id]) + self.max_ngram_size = multipliers.shape[-1] + self.history = None + self.next_pos = 0 + + @torch.inference_mode() + def __call__( + self, + input_ids: torch.Tensor, + start_pos: int, + token_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + input_ids = input_ids.to(device="cpu", dtype=torch.int64) + if input_ids.ndim != 2 or input_ids.shape[1] == 0: + raise ValueError("Expected a non-empty [batch, sequence] token tensor") + if torch.any(input_ids < 0) or torch.any(input_ids >= len(self.token_map)): + raise ValueError("Token id outside compressed token map") + batch, length = input_ids.shape + history_length = self.max_ngram_size - 1 + if start_pos == 0: + self.history = torch.full((batch, history_length), self.DEAD) + self.next_pos = 0 + if start_pos != self.next_pos or self.history is None: + raise ValueError("Hash calls must have contiguous start_pos") + if self.history.shape[0] != batch: + raise ValueError("Batch size changed during a sequence") + compressed = self.token_map[input_ids] + if token_mask is not None: + if token_mask.shape != input_ids.shape: + raise ValueError("Token mask shape mismatch") + compressed = compressed.masked_fill(~token_mask.cpu().bool(), self.DEAD) + context = torch.cat((self.history, compressed), dim=1) + tokens = [] + blocked = torch.zeros_like(compressed, dtype=torch.bool) + for shift in range(self.max_ngram_size): + begin = history_length - shift + source = context[:, begin : begin + length] + blocked |= source == self.DEAD + tokens.append(torch.where(blocked, self.pad_id, source)) + products = torch.stack(tokens, dim=-1).unsqueeze(2) * self.multipliers + heads = self.primes.shape[-1] // history_length + primes = self.primes.unflatten(-1, (history_length, heads)) + rolling, hashes = products[..., 0], [] + for shift in range(1, self.max_ngram_size): + rolling = torch.bitwise_xor(rolling, products[..., shift]) + hashes.append(rolling.unsqueeze(-1) % primes[:, shift - 1]) + self.history = context[:, -history_length:].clone() + self.next_pos += length + return torch.cat(hashes, dim=-1) + self.offsets + + +def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + x = x.float() + return x * torch.rsqrt(x.square().mean(-1, keepdim=True) + eps) * weight.float() + + +def route( + x: torch.Tensor, + weight: torch.Tensor, + selection_bias: torch.Tensor, + topk: int = 6, + scale: float = 1.5, + temperature: float = 1.0, +) -> tuple[torch.Tensor, torch.Tensor]: + scores = F.softplus(F.linear(x.float(), weight.float()) / temperature).sqrt() + indices = (scores + selection_bias.float()).topk(topk, dim=-1).indices + weights = scores.gather(-1, indices) + if topk > 1: + weights = weights / (weights.sum(-1, keepdim=True) + 1e-20) + return weights * scale, indices + + +def engram_update( + x: torch.Tensor, + kv: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + token_mask: torch.Tensor | None = None, + eps: float = 1e-20, +) -> torch.Tensor: + """FP32 signed-square-root gate, normalized separately for each HC copy.""" + copies, dim = x.shape[-2:] + key, value = kv.float().split((copies * dim, dim), dim=-1) + key = key.unflatten(-1, (copies, dim)) + h = x.float() + rstd = torch.rsqrt(h.square().mean(-1) + eps) + rstd *= torch.rsqrt(key.square().mean(-1) + eps) + weight = q_weight.float() * k_weight.float() + dot = (h * weight * key).sum(-1) + dot = dot * rstd * dim**-0.5 + signed_root = torch.copysign(dot.abs().clamp_min(1e-6).sqrt(), dot) + gate = torch.sigmoid(signed_root) + if token_mask is not None: + gate = gate.masked_fill(~token_mask.unsqueeze(-1), 0) + return h + gate.unsqueeze(-1) * value.unsqueeze(-2) diff --git a/benchmarks/deepseek_v41/performance.py b/benchmarks/deepseek_v41/performance.py new file mode 100644 index 0000000000..dd09ad9cdf --- /dev/null +++ b/benchmarks/deepseek_v41/performance.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Matched single-request timing, with separate opt-in stage instrumentation. + +Run twice, with and without --draft, in separate engine processes. Both runs +use the same 256 MiB KV allocation, chunk 512, TP4 and 30% shared offload. +Raw timings have no finite-logit or stage hooks. Profiled timings are labeled +separately; nested expert/Engram stream intervals are not additive to target. +""" + +import argparse +import json +import time +from pathlib import Path + +from benchmarks.deepseek_v41 import trace +from benchmarks.deepseek_v41.acceptance import request, validate_commits +from benchmarks.deepseek_v41.main_smoke import ( + _memory_report, + create_llm, + trim_allocator, +) +from benchmarks.deepseek_v41.resources import ResourceMonitor + + +def make_prompt(tokenizer, length): + marker = "__V41_BENCH_FILL__" + text = tokenizer.apply_chat_template( + [ + dict( + role="user", + content=marker + "\n忽略前面的背景资料。只输出从0到1000的整数JSON数组，" + "不要省略任何数字，不要解释。", + ) + ], + tokenize=False, + enable_thinking=False, + ) + before, after = [ + tokenizer.encode(s, add_special_tokens=False) for s in text.split(marker) + ] + fill = tokenizer.encode( + "背景资料：春夏秋冬，日月星辰。" * 4096, add_special_tokens=False + ) + size = length - len(before) - len(after) + if not 0 <= size <= len(fill): + raise ValueError("Prompt length cannot hold the benchmark instruction") + return before + fill[:size] + after + + +def run(args): + resources = ResourceMonitor() + llm = None + report = dict( + status="running", + dspark=args.draft is not None, + scope="single request; raw timings separate from instrumented profiles", + kv_cache_bytes=256 * 1024**2, + chunk=512, + offload_fraction=0.30, + target_graph=args.decode_graph, + draft_graph=args.draft_graph, + custom_allreduce=args.custom_ar, + cases=[], + ) + + def save(): + report["resources"] = resources.report() + temporary = args.output.with_suffix(args.output.suffix + ".tmp") + temporary.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n") + temporary.replace(args.output) + + try: + start = time.monotonic() + llm = create_llm( + args.model_root, + 512, + args.draft, + kv_cache_mib=256, + decode_graph=args.decode_graph, + draft_graph=args.draft_graph, + custom_allreduce=args.custom_ar, + ) + report["load_seconds"] = time.monotonic() - start + llm.collective_rpc(trim_allocator) + tokenizer = llm.get_tokenizer() + prompts = {n: make_prompt(tokenizer, n) for n in args.lengths} + for length, prompt in prompts.items(): + # Warm each shape and resident set before recording raw latency. + for _ in range(2): + request(llm, prompt, args.output_tokens) + before = llm.collective_rpc(_memory_report) + runs = [ + request(llm, prompt, args.output_tokens) for _ in range(args.repeats) + ] + if any(len(r["token_ids"]) != args.output_tokens for r in runs): + raise AssertionError("Timing workload ended before its output budget") + if any(r["token_ids"] != runs[0]["token_ids"] for r in runs[1:]): + raise AssertionError("Repeated greedy completions differ") + report["cases"].append( + dict( + prompt_tokens=length, + instrumentation=False, + runs=runs, + memory_before=before, + memory_after=llm.collective_rpc(_memory_report), + ) + ) + save() + print( + json.dumps(dict(raw_prompt_tokens=length, completed=True)), flush=True + ) + + # The same prompts now qualify stage and acceptance reporting; these + # wall times are excluded from the raw performance comparison. + llm.collective_rpc(trace.install) + for length, prompt in prompts.items(): + result = request(llm, prompt, args.output_tokens) + result["trace_by_rank"] = llm.collective_rpc(trace.drain) + if args.decode_graph: + for rank in result["trace_by_rank"]: + small = [row for row in rank["target_forwards"] if row["rows"] <= 6] + assert all(row["mode"] == "FULL" for row in small), (length, rank) + if args.custom_ar: + assert all( + rank["custom_allreduce_enabled"] for rank in result["trace_by_rank"] + ) + validate_commits( + result["trace_by_rank"], len(prompt) + len(result["token_ids"]) - 1 + ) + report["cases"].append( + dict( + prompt_tokens=length, + instrumentation=True, + result=result, + memory_after=llm.collective_rpc(_memory_report), + ) + ) + save() + if resources.report()["host_minimum_available_bytes"] < 8 * 1024**3: + raise MemoryError("Performance qualification breached the host reserve") + report["status"] = "completed" + save() + except BaseException as error: + report.update(status="failed", error=repr(error)) + save() + raise + finally: + if llm is not None: + llm.llm_engine.engine_core.shutdown() + resources.close() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-root", type=Path, required=True) + parser.add_argument("--draft", type=Path) + parser.add_argument("--lengths", type=int, nargs="+", default=[128, 2048]) + parser.add_argument("--output-tokens", type=int, default=128) + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--decode-graph", action="store_true") + parser.add_argument("--draft-graph", action="store_true") + parser.add_argument("--custom-ar", action="store_true") + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if ( + args.repeats < 1 + or args.output_tokens < 2 + or min(args.lengths) < 1 + or max(args.lengths) + args.output_tokens > 16384 + ): + parser.error("Invalid request count or context budget") + run(args) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/deepseek_v41/prepare_dspark.py b/benchmarks/deepseek_v41/prepare_dspark.py new file mode 100644 index 0000000000..495cc5917d --- /dev/null +++ b/benchmarks/deepseek_v41/prepare_dspark.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Write a small DSpark-only config/index and symlink original weight shards.""" + +import argparse +import json +from pathlib import Path + +from benchmarks.deepseek_v41.checkpoint_audit import safetensors_names + + +def prepare(root: Path, destination: Path): + original = json.loads((root / "base/config.json").read_text()) + config = dict(original["text_config"]) + if original["model_type"] != "deepseek_v41": + raise ValueError("Expected the audited DeepSeek V4.1 checkpoint") + config.update( + architectures=["DSparkDeepseekV41ForCausalLM"], + dtype="float16", + num_hidden_layers=3, + n_routed_experts=config["dspark_n_routed_experts"], + num_experts_per_tok=config["dspark_num_experts_per_tok"], + num_nextn_predict_layers=0, + compress_ratios=[0, 0, 0], + kv_source_layer_ids=[], + index_source_layer_ids=[], + engram_layer_ids=[], + max_position_embeddings=16384, + rope_scaling=None, + rope_parameters={"rope_type": "default", "rope_theta": 10000.0}, + quantization_config=original["quantization_config"], + bos_token_id=original["bos_token_id"], + eos_token_id=original["eos_token_id"], + pad_token_id=original["pad_token_id"], + ) + index = json.loads((root / "base/model.safetensors.index.json").read_text()) + expected = {k: v for k, v in index["weight_map"].items() if k.startswith("mtp.")} + actual, total = {}, 0 + shards = sorted((root / "base/dspark").glob("*.safetensors")) + for path in shards: + for name, record in safetensors_names(path).items(): + if name in actual: + raise ValueError(f"Duplicate draft tensor {name}") + actual[name] = path.name + start, end = record["data_offsets"] + total += end - start + if actual != expected or len(actual) != 2401: + raise ValueError("DSpark shard coverage does not match the original index") + destination.mkdir(parents=True, exist_ok=True) + for source in shards: + link = destination / source.name + if link.is_symlink() and link.resolve() == source.resolve(): + continue + if link.exists() or link.is_symlink(): + raise FileExistsError(link) + link.symlink_to(source.resolve()) + (destination / "config.json").write_text(json.dumps(config, indent=2) + "\n") + (destination / "model.safetensors.index.json").write_text( + json.dumps(dict(metadata=dict(total_size=total), weight_map=actual), indent=2) + + "\n" + ) + return dict(tensors=len(actual), payload_bytes=total, weight_conversion=False) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-root", type=Path, required=True) + parser.add_argument("--destination", type=Path, required=True) + args = parser.parse_args() + print(json.dumps(prepare(args.model_root, args.destination))) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/deepseek_v41/resources.py b/benchmarks/deepseek_v41/resources.py new file mode 100644 index 0000000000..d9ac7891c5 --- /dev/null +++ b/benchmarks/deepseek_v41/resources.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Sample node/driver usage separately from per-worker allocator peaks.""" + +import threading +import time +from collections import deque + +from benchmarks.deepseek_v41.shared_host_probe import cgroup_memory, memory_available + + +class ResourceMonitor: + def __init__(self): + from vllm.third_party import pynvml + + self.nvml = pynvml + pynvml.nvmlInit() + self.handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in range(4)] + self.gpus = [ + dict( + uuid=pynvml.nvmlDeviceGetUUID(handle), + total_bytes=pynvml.nvmlDeviceGetMemoryInfo(handle).total, + peak_used_bytes=0, + minimum_free_bytes=pynvml.nvmlDeviceGetMemoryInfo(handle).total, + ) + for handle in self.handles + ] + self.minimum_available = memory_available() + self.peak_cgroup = 0 + self.samples: deque[dict] = deque(maxlen=20000) + self.samples_seen = 0 + self.stopped, self.lock = threading.Event(), threading.Lock() + self.thread = threading.Thread(target=self._run, daemon=True) + self.thread.start() + + def _run(self): + while not self.stopped.is_set(): + info = [self.nvml.nvmlDeviceGetMemoryInfo(h) for h in self.handles] + usage = [] + for handle in self.handles: + try: + utilization = self.nvml.nvmlDeviceGetUtilizationRates(handle) + usage.append( + dict( + gpu_utilization_pct=utilization.gpu, + memory_controller_utilization_pct=utilization.memory, + power_watts=self.nvml.nvmlDeviceGetPowerUsage(handle) + / 1000, + sm_clock_mhz=self.nvml.nvmlDeviceGetClockInfo( + handle, self.nvml.NVML_CLOCK_SM + ), + memory_clock_mhz=self.nvml.nvmlDeviceGetClockInfo( + handle, self.nvml.NVML_CLOCK_MEM + ), + throttle_reasons=self.nvml.nvmlDeviceGetCurrentClocksThrottleReasons( + handle + ), + ) + ) + except self.nvml.NVMLError as error: + usage.append(dict(unavailable=str(error))) + available, current = memory_available(), cgroup_memory()["current"] + with self.lock: + self.samples.append(dict(monotonic=time.monotonic(), gpu=usage)) + self.samples_seen += 1 + for entry, value in zip(self.gpus, info): + entry["peak_used_bytes"] = max(entry["peak_used_bytes"], value.used) + entry["minimum_free_bytes"] = min( + entry["minimum_free_bytes"], value.free + ) + self.minimum_available = min(self.minimum_available, available) + self.peak_cgroup = max(self.peak_cgroup, current) + self.stopped.wait(0.25) + + def report(self): + with self.lock: + return dict( + sample_interval_seconds=0.25, + gpu=[dict(entry) for entry in self.gpus], + host_minimum_available_bytes=self.minimum_available, + cgroup_peak_observed_bytes=self.peak_cgroup, + gpu_usage_samples=list(self.samples), + gpu_usage_samples_dropped=self.samples_seen - len(self.samples), + ) + + def close(self): + self.stopped.set() + self.thread.join(2) + self.nvml.nvmlShutdown() diff --git a/benchmarks/deepseek_v41/results/54633-20261024.md b/benchmarks/deepseek_v41/results/54633-20261024.md new file mode 100644 index 0000000000..3a68f50222 --- /dev/null +++ b/benchmarks/deepseek_v41/results/54633-20261024.md @@ -0,0 +1,172 @@ +# DeepSeek V4.1 Flash — 54633 首版验收 + +代码基线：`6a143197e6e5f6fa77bd483df0c5a09dcea6aa20`。 +分支：`codex/ds41-flash-sm70-20261024`。 +本地/远端改动源码 SHA256 一致：`2f3aae4456672c6081f7f6d631758b078ce47c3d2785667dd00ebb67369ac922`。 +稳定性与性能两轮的模型/native 源码一致；性能 profile 另加 prefill 计时。 + +最终范围：V100 32GiB ×4、TP4、单请求、16,384 总上下文、DSpark5。 +dense 使用有界 FP16 GEMM（prefill_min_m=1、32MiB 工作区）； +专家 M1/M6 保留 DP4A，prefill 使用 Tensor Core。 +M2–5 尾批次采用有界 GPU FP16 GEMM，路由不回到 CPU。 +归一化、路由、mHC/HC 残差、C2 池化及 Engram 门控保留必要的 FP32。 +草稿专家各投影仍为 FP16，路由求和和 TP 归约以 FP32 进入 HC，避免有限分片相加溢出。 +没有完整 FP32 参考模型或完整 FP16 专家副本。 + +服务容器：`ds41-serve-20261024`，主机回环 `127.0.0.1:18000`。 +只暴露 GPU 0–3；P400 显示卡和其他部署保持独立。 + +## 正确性与稳定性 + +- 最后一次归约、mHC、MXFP4、Engram、状态和采样回归：149 passed。 +- 后续加载、混合量化和状态回归：98 passed；原生协调器故障场景：10 passed。 +- 64、2K、16K 上下文：目标/草稿 logits 有限，四 rank 的状态提交一致。 +- 数学、聊天、代码、JSON：M1 与 DSpark5 贪心 token 序列一致。 +- 127/128 边界的接受 0–5 个草稿，共 12 组；拒绝后继续生成与参考一致。 +- 每条请求结束状态与已输出前缀对应；实际 Engram 历史内容另用检查点 token map 核对。 +- 稳定性：1804.88s，210 条请求；每四条包含一条 temperature=0.8 请求。 +- OpenAI 接口：聊天、代码、数学、JSON、工具调用、流式工具调用和工具回填通过。 +- 服务端真实长请求：16255 输入 token，实际回答 `17 × 23 = 391`。 +- 长请求最初因只接受裸字符串 391 的展示格式断言而失败；保留原始响应，重新核对全部响应后，仅放行正确数值或完整正确算式。 +- 多模态和超过 16K 的输入被拒绝。 + +## 上下文边界调用 + +以下为带有限值及状态检查、强制 EOS 后继续执行的边界探针，不代表聊天质量或稳态性能。 + +| 总上下文上限 | 输入 | 输出 | 总耗时 s | TTFT s | 平均 TPOT ms | +| ---: | ---: | ---: | ---: | ---: | ---: | +| 64 | 56 | 2 | 8.599 | 7.772 | 827.322 | +| 2048 | 2016 | 32 | 33.598 | 29.356 | 136.834 | +| 16384 | 16352 | 32 | 243.309 | 239.068 | 136.826 | + +## 独立引擎性能对比 + +两次运行的 KV 都为 256MiB、chunk 512、30% 卸载；无推测运行不加载草稿。 +DSpark 为草稿预留槽位后，实际单步输入上限为 508；无推测为 512。 +TTFT 包含这部分实际调度差异。 +负载是背景填充后生成 0–1000 整数 JSON 数组，固定输出 128 token。 +速度和接受率适用于该固定负载，会随任务变化。 +每种长度预热两次，再测三次；以下耗时不带有限值/阶段 hook。 +吞吐包含 prefill；TPOT 按前端实际到达的 token 数计，包含推测输出的成批到达。 + +| 配置 | 输入/输出 | TTFT P50/P95 s | TPOT P50/P95 ms | 吞吐 P50 tok/s | +| --- | ---: | ---: | ---: | ---: | +| 无推测 | 128/128 | 7.740/7.799 | 231.397/232.559 | 3.454 | +| 无推测 | 2048/128 | 28.759/28.761 | 230.182/230.873 | 2.208 | +| DSpark5 | 128/128 | 7.750/7.772 | 46.575/46.754 | 9.374 | +| DSpark5 | 2048/128 | 36.873/36.887 | 46.373/46.514 | 2.994 | + +本轮引擎初始化：无推测 186.80s，DSpark5 229.06s；不计入 TTFT，初始化时间受文件缓存状态影响。 +这两组性能负载的贪心 token 序列也一致。P95 为三次样本的线性分位估计，样本量有限。 +无推测原始采样完成后，仪表安装曾因缺少 drafter 属性而失败；修复后在同配置的新引擎补采阶段 profile，并逐 token 核对一致。原始采样记录保留，资源表取两轮观测极值。 + +## 阶段与接受率 + +以下来自另外一次带 CUDA event 的调用；每阶段取最慢 rank。 +target 包含专家传输、Engram 等待及主机发射间隙，嵌套项不能再次相加。 +draft 的投影、context、head、Markov 采样分别列出；head 嵌套在 Markov 采样中。 + +CPU 调用时间记录同一 rank 的函数调用耗时；它也包含嵌套调用。 + +| 模式/输入 | 阶段 | 最慢 rank | 流间隔 ms | 同 rank CPU 调用 ms | +| --- | --- | ---: | ---: | ---: | +| 无推测/128 | commit | 0 | 73.571 | 70.429 | +| 无推测/128 | engram_wait_and_dequant | 0 | 65.745 | 57.899 | +| 无推测/128 | expert_transfer_and_group_wait | 3 | 6685.002 | 1814.691 | +| 无推测/128 | prefill_expert_transfer_and_transcode | 1 | 7611.717 | 1276.173 | +| 无推测/128 | target | 0 | 36750.694 | 35775.320 | +| 无推测/128 | target_head | 3 | 273.603 | 337.323 | +| 无推测/128 | target_sampling | 0 | 20.945 | 20.215 | +| 无推测/2048 | commit | 1 | 73.183 | 71.142 | +| 无推测/2048 | engram_wait_and_dequant | 0 | 66.925 | 57.779 | +| 无推测/2048 | expert_transfer_and_group_wait | 2 | 6535.304 | 1846.733 | +| 无推测/2048 | prefill_expert_transfer_and_transcode | 3 | 26965.159 | 623.073 | +| 无推测/2048 | target | 0 | 57577.930 | 54553.685 | +| 无推测/2048 | target_head | 2 | 274.742 | 186.255 | +| 无推测/2048 | target_sampling | 3 | 20.052 | 20.749 | +| DSpark5/128 | commit | 1 | 12.663 | 13.046 | +| DSpark5/128 | draft | 2 | 434.620 | 579.440 | +| DSpark5/128 | draft_context | 1 | 35.398 | 40.179 | +| DSpark5/128 | draft_head | 2 | 51.157 | 84.477 | +| DSpark5/128 | draft_main_projection | 1 | 39.943 | 10.796 | +| DSpark5/128 | draft_markov_sampling | 2 | 75.871 | 247.795 | +| DSpark5/128 | engram_wait_and_dequant | 1 | 10.312 | 11.055 | +| DSpark5/128 | expert_transfer_and_group_wait | 2 | 1151.096 | 316.306 | +| DSpark5/128 | prefill_expert_transfer_and_transcode | 0 | 7550.248 | 1273.077 | +| DSpark5/128 | target | 1 | 12962.979 | 11959.558 | +| DSpark5/128 | target_head | 2 | 51.844 | 180.873 | +| DSpark5/128 | target_sampling | 1 | 10.067 | 10.967 | +| DSpark5/2048 | commit | 1 | 13.256 | 15.165 | +| DSpark5/2048 | draft | 3 | 480.075 | 2025.202 | +| DSpark5/2048 | draft_context | 1 | 51.236 | 110.159 | +| DSpark5/2048 | draft_head | 3 | 58.308 | 91.678 | +| DSpark5/2048 | draft_main_projection | 1 | 151.073 | 12.650 | +| DSpark5/2048 | draft_markov_sampling | 3 | 87.349 | 256.855 | +| DSpark5/2048 | engram_wait_and_dequant | 1 | 18.458 | 12.267 | +| DSpark5/2048 | expert_transfer_and_group_wait | 2 | 1161.438 | 313.281 | +| DSpark5/2048 | prefill_expert_transfer_and_transcode | 0 | 34774.976 | 4068.742 | +| DSpark5/2048 | target | 1 | 41837.813 | 37398.233 | +| DSpark5/2048 | target_head | 3 | 58.778 | 187.951 | +| DSpark5/2048 | target_sampling | 1 | 9.573 | 11.273 | + +DSpark5 输入 128：非终止、完整 M6 的 21 轮中，平均接受草稿 5.000 个，平均每次验证输出 6.000 个 token。 +位置 1–5 的累计接受率：100.00%, 100.00%, 100.00%, 100.00%, 100.00%。 + +DSpark5 输入 2048：非终止、完整 M6 的 21 轮中，平均接受草稿 4.857 个，平均每次验证输出 5.857 个 token。 +位置 1–5 的累计接受率：100.00%, 100.00%, 95.24%, 95.24%, 95.24%。 + +## 内存和共享存储 + +专家冷池保存未切 TP 的原始 GGUF 字节，仅一份 39,553,228,800B，另有组级 512MiB 交换区。 +每卡 raw GPU 槽 26,362,255,178B；单套 canonical 工作区 566,231,040B。 +0.25s 采样的主机最低 MemAvailable：11.493GiB；要求至少 8GiB。 + +| GPU | 驱动采样峰值 GiB | Torch 分配峰值 GiB | 最低采样余量 MiB | +| ---: | ---: | ---: | ---: | +| 0 | 31.843 | 30.086 | 160.6 | +| 1 | 31.843 | 30.086 | 160.6 | +| 2 | 31.843 | 30.086 | 160.6 | +| 3 | 31.843 | 30.086 | 160.6 | + +| 独立性能运行 | 最高 GPU 峰值 GiB | 最低主机余量 GiB | +| --- | ---: | ---: | +| 无推测 | 29.732 | 11.816 | +| DSpark5 | 31.826 | 11.616 | + +真实服务 API 观察窗口：552.4s，GPU 采样最高 31.958GiB，最低采样显存余量 42.6MiB，主机最低余量 11.055GiB。 + +四 rank 的专家对象 PSS 合计 40,090,099,712B，共享对象大小 40,090,099,712B；对象 device/inode 一致。不能相加四个进程 RSS 计算物理占用。 + +| 卸载 | 冷专家/层 | 共用冷数据 GB | 每卡 raw GB | 实际资格 | +| ---: | ---: | ---: | ---: | --- | +| 0.25 | 96 | 33.018 | 28.170 | 40 层原始存储及换入换出；未执行完整模型 | +| 0.20 | 77 | 26.483 | 29.979 | 40 层原始存储及换入换出；未执行完整模型 | + +以 30% 稳定性验收的驱动峰值为基准，保持其他分配不变，25% 预计每卡缺 1.640GB；这是推算，不是该比例完整模型的实测峰值。 +以 30% 稳定性验收的驱动峰值为基准，保持其他分配不变，20% 预计每卡缺 3.448GB；这是推算，不是该比例完整模型的实测峰值。 + +## 传输与 I/O + +完整验收运行累计 decode 未命中 3,677 次；四卡载入/回写字节合计 34.935/31.569GB。 + +每个性能 profile 从计数清零开始；缓存未命中和 Engram 计数是组级复制观测，仅取一份。GPU 拷贝字节按各 rank 累加。 +prefill 直接读取共享映射，按完整 GGUF 量化块计算冷转码逻辑读取量；这些是逻辑工作量，不是 PCIe 总线事务计数。 + +| 模式/输入 | decode 未命中 | GPU 载入/回写 GB | prefill 冷转码 次数/GB | Engram 文件读取 MB | 读取端请求时间 ms | +| --- | ---: | ---: | ---: | ---: | ---: | +| 无推测/128 | 0 | 0.000/0.000 | 4600/43.778 | 1.157 | 186.357 | +| 无推测/2048 | 0 | 0.000/0.000 | 18400/175.114 | 1.343 | 182.493 | +| DSpark5/128 | 0 | 0.000/0.000 | 4600/43.778 | 1.191 | 33.317 | +| DSpark5/2048 | 0 | 0.000/0.000 | 23000/218.892 | 1.378 | 52.402 | + +## 验证边界 + +数值路径、状态提交、真实请求和采样法则分别验证；32,768 链的采样分布测试是已知 PDF 的算法检查，不等同于完整模型的统计质量评测。 +不承诺多模态、四并发、超过 16K 或其他 GPU/TP 拓扑。25%/20% 目前仅通过存储资格检查。 +首次 API 容器启动在 TP2 的 Engram CUDA 主存注册时报 invalid argument，进程退出且没有 OOM；同配置重启后完成 API 验收。该次启动异常的根因尚未确认，原始日志保留在 service-startup-first.log。 +全量 pre-commit 曾受 Node 环境下载阻塞；针对改动的 Ruff、格式、typos、clang-format 和 Linux 目标 mypy 单独执行。 +基础版本 check_layering 的 moe_runner.py 计数不匹配与本分支无关，该文件未修改。 +Linux 目标 mypy：44 个源文件、35 个测试文件通过；注册表的三个既有报错已用原始基线复现，并在本次检查中单独隔离。 +原始 JSON 和逐项日志保存在 54633 的 `/home/ymzx/ds41-dev/state/`。 +代码尚未提交或创建 PR，供人工审查。 diff --git a/benchmarks/deepseek_v41/serve-54633.yaml b/benchmarks/deepseek_v41/serve-54633.yaml new file mode 100644 index 0000000000..a1e2bcabbc --- /dev/null +++ b/benchmarks/deepseek_v41/serve-54633.yaml @@ -0,0 +1,57 @@ +# Development launch profile. Run only after the acceptance report passes. +# The container must expose only V100 GPUs 0-3 and preserve 8 GiB host reserve. +model: /home/ymzx/ds41-models/gguf/DeepSeek-V4.1-Flash-1017-GSQ-RCO-IQ2_XXS-00001-of-00004.gguf +hf-config-path: /home/ymzx/ds41-models/base +tokenizer: /home/ymzx/ds41-models/base +tokenizer-mode: deepseek_v41 +served-model-name: deepseek-v4.1-flash +host: 127.0.0.1 +port: 18000 +dtype: half +quantization: gguf +load-format: gguf +tensor-parallel-size: 4 +max-model-len: 16384 +max-num-seqs: 1 +max-num-batched-tokens: 512 +kv-cache-memory-bytes: 268435456 +gpu-memory-utilization: 0.96 +no-enforce-eager: true +compilation-config: + mode: 0 + cudagraph_mode: FULL + cudagraph_capture_sizes: [1, 2, 3, 4, 5, 6] + max_cudagraph_capture_size: 6 +# This baseline's YAML loader omits false scalar flags; use explicit negations. +no-async-scheduling: true +no-enable-prefix-caching: true +enable-chunked-prefill: true +no-disable-custom-all-reduce: true +enable-auto-tool-choice: true +tool-call-parser: deepseek_v41 +model-loader-extra-config: + engram_path: /home/ymzx/ds41-models/gguf/DeepSeek-V4.1-Flash-Engram-IQ4_NL.gguf +speculative-config: + method: dspark + model: /home/ymzx/ds41-dev/dspark-v41 + quantization: fp8 + num_speculative_tokens: 5 + draft_load_config: + load_format: safetensors + enforce_eager: false +kernel-config: + sm70_ring: + enabled: false + sm70_gguf: + expert_storage: original + expert_arena: true + expert_cache_host_storage: shared_original + expert_cache_offload_fraction: 0.30 + expert_cache_fraction: 0 + expert_cache_host_slack: 0 + expert_load_prefetch_banks: 1 + expert_cache_prefill_tokens: 512 + dense_storage: original + embedding_storage: original + prefill_min_m: 1 + dequant_workspace_bytes: 33554432 diff --git a/benchmarks/deepseek_v41/serve-ds41-phase1.yaml b/benchmarks/deepseek_v41/serve-ds41-phase1.yaml new file mode 100644 index 0000000000..f13a14b057 --- /dev/null +++ b/benchmarks/deepseek_v41/serve-ds41-phase1.yaml @@ -0,0 +1,58 @@ +# DS-V4.1 Phase 1 API (2026-10-26): async, chunk 2048, KV 384 MiB, offload 0.34. +# Development launch profile. Run only after the acceptance report passes. +# The container must expose only V100 GPUs 0-3 and preserve 8 GiB host reserve. +model: /home/ymzx/ds41-models/gguf/DeepSeek-V4.1-Flash-1017-GSQ-RCO-IQ2_XXS-00001-of-00004.gguf +hf-config-path: /home/ymzx/ds41-models/base +tokenizer: /home/ymzx/ds41-models/base +tokenizer-mode: deepseek_v41 +served-model-name: deepseek-v4.1-flash +host: 0.0.0.0 +port: 18000 +dtype: half +quantization: gguf +load-format: gguf +tensor-parallel-size: 4 +max-model-len: 16384 +max-num-seqs: 1 +max-num-batched-tokens: 2048 +kv-cache-memory-bytes: 402653184 +gpu-memory-utilization: 0.96 +no-enforce-eager: true +compilation-config: + mode: 0 + cudagraph_mode: FULL + cudagraph_capture_sizes: [1, 2, 3, 4, 5, 6] + max_cudagraph_capture_size: 6 +# This baseline's YAML loader omits false scalar flags; use explicit negations. +async-scheduling: true +no-enable-prefix-caching: true +enable-chunked-prefill: true +no-disable-custom-all-reduce: true +enable-auto-tool-choice: true +tool-call-parser: deepseek_v41 +model-loader-extra-config: + engram_path: /home/ymzx/ds41-models/gguf/DeepSeek-V4.1-Flash-Engram-IQ4_NL.gguf +speculative-config: + method: dspark + model: /home/ymzx/ds41-dev/dspark-v41 + quantization: fp8 + num_speculative_tokens: 5 + draft_load_config: + load_format: safetensors + enforce_eager: false +kernel-config: + sm70_ring: + enabled: false + sm70_gguf: + expert_storage: original + expert_arena: true + expert_cache_host_storage: shared_original + expert_cache_offload_fraction: 0.34 + expert_cache_fraction: 0 + expert_cache_host_slack: 0 + expert_load_prefetch_banks: 1 + expert_cache_prefill_tokens: 2048 + dense_storage: original + embedding_storage: original + prefill_min_m: 1 + dequant_workspace_bytes: 33554432 diff --git a/benchmarks/deepseek_v41/serve-ds41-phase2.yaml b/benchmarks/deepseek_v41/serve-ds41-phase2.yaml new file mode 100644 index 0000000000..099f4abdab --- /dev/null +++ b/benchmarks/deepseek_v41/serve-ds41-phase2.yaml @@ -0,0 +1,59 @@ +# DS-V4.1 Phase 2 API (2026-10-27): async, chunk 8192, KV 288 MiB, offload 0.33, session prefix cache. +# Needs run_serve2.sh envs: VLLM_V41_DOWN_ROWS=1, VLLM_V41K_LIB=v41k_g3, VLLM_V41_SESSION_PREFIX=1. +# Development launch profile. Run only after the acceptance report passes. +# The container must expose only V100 GPUs 0-3 and preserve 8 GiB host reserve. +model: /home/ymzx/ds41-models/gguf/DeepSeek-V4.1-Flash-1017-GSQ-RCO-IQ2_XXS-00001-of-00004.gguf +hf-config-path: /home/ymzx/ds41-models/base +tokenizer: /home/ymzx/ds41-models/base +tokenizer-mode: deepseek_v41 +served-model-name: deepseek-v4.1-flash +host: 0.0.0.0 +port: 18000 +dtype: half +quantization: gguf +load-format: gguf +tensor-parallel-size: 4 +max-model-len: 16384 +max-num-seqs: 1 +max-num-batched-tokens: 8192 +kv-cache-memory-bytes: 301989888 +gpu-memory-utilization: 0.96 +no-enforce-eager: true +compilation-config: + mode: 0 + cudagraph_mode: FULL + cudagraph_capture_sizes: [1, 2, 3, 4, 5, 6] + max_cudagraph_capture_size: 6 +# This baseline's YAML loader omits false scalar flags; use explicit negations. +async-scheduling: true +enable-prefix-caching: true +enable-chunked-prefill: true +no-disable-custom-all-reduce: true +enable-auto-tool-choice: true +tool-call-parser: deepseek_v41 +model-loader-extra-config: + engram_path: /home/ymzx/ds41-models/gguf/DeepSeek-V4.1-Flash-Engram-IQ4_NL.gguf +speculative-config: + method: dspark + model: /home/ymzx/ds41-dev/dspark-v41 + quantization: fp8 + num_speculative_tokens: 5 + draft_load_config: + load_format: safetensors + enforce_eager: false +kernel-config: + sm70_ring: + enabled: false + sm70_gguf: + expert_storage: original + expert_arena: true + expert_cache_host_storage: shared_original + expert_cache_offload_fraction: 0.33 + expert_cache_fraction: 0 + expert_cache_host_slack: 0 + expert_load_prefetch_banks: 1 + expert_cache_prefill_tokens: 4096 + dense_storage: original + embedding_storage: original + prefill_min_m: 1 + dequant_workspace_bytes: 33554432 diff --git a/benchmarks/deepseek_v41/serve-ds41-phase3.yaml b/benchmarks/deepseek_v41/serve-ds41-phase3.yaml new file mode 100644 index 0000000000..e869491e28 --- /dev/null +++ b/benchmarks/deepseek_v41/serve-ds41-phase3.yaml @@ -0,0 +1,64 @@ +# DS-V4.1 Phase 3 API (2026-10-27): 512K context (packed FP4 KV, host latents), prefix cache, tools, DSpark 5. +# Needs run_serve3.sh envs (packed KV, G3 rows, session prefix). +# Development launch profile. Run only after the acceptance report passes. +# The container must expose only V100 GPUs 0-3 and preserve 8 GiB host reserve. +model: /home/ymzx/ds41-models/gguf/DeepSeek-V4.1-Flash-1017-GSQ-RCO-IQ2_XXS-00001-of-00004.gguf +hf-config-path: /home/ymzx/ds41-models/base +tokenizer: /home/ymzx/ds41-models/base +tokenizer-mode: deepseek_v41 +served-model-name: deepseek-v4.1-flash +host: 0.0.0.0 +port: 18000 +dtype: half +quantization: gguf +load-format: gguf +tensor-parallel-size: 4 +max-model-len: 524288 +max-num-seqs: 1 +max-num-batched-tokens: 8192 +kv-cache-memory-bytes: 268435456 +gpu-memory-utilization: 0.96 +no-enforce-eager: true +compilation-config: + mode: 0 + cudagraph_mode: FULL + cudagraph_capture_sizes: [1, 2, 3, 4, 5, 6] + max_cudagraph_capture_size: 6 +# This baseline's YAML loader omits false scalar flags; use explicit negations. +async-scheduling: true +enable-prefix-caching: true +enable-chunked-prefill: true +no-disable-custom-all-reduce: true +enable-auto-tool-choice: true +tool-call-parser: deepseek_v41 +# Thinking on by default (R1-style ...); a request can pass +# chat_template_kwargs {"thinking": false} or reasoning_effort. +reasoning-parser: deepseek_v4 +default-chat-template-kwargs: '{"thinking": true}' +model-loader-extra-config: + engram_path: /home/ymzx/ds41-models/gguf/DeepSeek-V4.1-Flash-Engram-IQ4_NL.gguf +speculative-config: + method: dspark + model: /home/ymzx/ds41-dev/dspark-v41 + quantization: fp8 + num_speculative_tokens: 5 + draft_load_config: + load_format: safetensors + enforce_eager: false +kernel-config: + sm70_ring: + enabled: false + sm70_gguf: + expert_storage: original + expert_arena: true + expert_cache_host_storage: shared_original + expert_cache_offload_fraction: 0.33 + expert_cache_fraction: 0 + expert_cache_host_slack: 0 + expert_load_prefetch_banks: 1 + expert_cache_prefill_tokens: 4096 + dense_storage: original + embedding_storage: original + prefill_min_m: 1 + dequant_workspace_bytes: 33554432 +api-key: sk-ds41-dNCMTcG1OUSl diff --git a/benchmarks/deepseek_v41/shared_cache_probe.py b/benchmarks/deepseek_v41/shared_cache_probe.py new file mode 100644 index 0000000000..7911b719b1 --- /dev/null +++ b/benchmarks/deepseek_v41/shared_cache_probe.py @@ -0,0 +1,235 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Synthetic TP4 M1/M6 expert-cache protocol, not a model speed benchmark.""" + +import argparse +import json +import multiprocessing as mp +import socket +import time +import traceback +from pathlib import Path + +import numpy as np +import torch + +from vllm.model_executor.layers.quantization.gguf_shared_runtime import ( + SHARDS, + SharedExpertLayerSpec, + SharedExpertOwner, + SharedExpertRank, +) +from vllm.utils.shared_host import SharedHostArena + + +def _run_rank(rank, data, control, specs, iterations, barrier): + banks, storage = {}, [] + for layer, spec in enumerate(specs): + for shard, layout in zip(SHARDS, spec.layouts(rank)): + size = (spec.resident + 8) * layout.device_rows * layout.device_row_bytes + backing = torch.zeros( + size + layout.guard_bytes, dtype=torch.uint8, device=rank + ) + bank = backing[:size].view( + spec.resident + 8, layout.device_rows, layout.device_row_bytes + ) + for expert in range(spec.resident): + bank[expert].view(torch.int16).fill_(expert + layer * 1024) + banks[layer, shard] = bank + storage.append((backing, layout)) + cache = SharedExpertRank(data, control, rank, specs, banks) + try: + cache.warmup() + graphs = {} + for layer in range(len(specs)): + for count in (6, 36): + ids = torch.zeros(count, dtype=torch.int32, device=rank) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = cache.resolve(layer, ids) + graphs[layer, count] = graph, ids, output + barrier.wait(180) + misses, milliseconds = [], [] + for i in range(iterations): + count = 6 if i % 2 == 0 else 36 + for layer, spec in enumerate(specs): + graph, ids, output = graphs[layer, count] + requested = ( + torch.arange(count, device=rank, dtype=torch.int32) + + spec.resident + + i * 36 + ) % spec.experts + ids.copy_(requested) + before, after = ( + torch.cuda.Event(enable_timing=True), + torch.cuda.Event(enable_timing=True), + ) + before.record() + graph.replay() + after.record() + after.synchronize() + milliseconds.append(before.elapsed_time(after)) + misses.append(int(cache.control_host[3])) + for shard in SHARDS: + actual = banks[layer, shard][output.long()].view(torch.int16) + if not bool( + (actual == requested[:, None, None] + layer * 1024).all() + ): + raise AssertionError( + (rank, i, layer, shard, "incorrect expert bytes") + ) + for backing, layout in storage: + if layout.guard_bytes and bool(backing[-layout.guard_bytes :].any()): + raise AssertionError("Bank guard was overwritten") + return { + "rank": rank, + "passed": True, + "misses": misses, + "protocol_milliseconds": milliseconds, + } + finally: + cache.close() + + +def _consumer(rank, connection, spec, iterations, barrier): + data = control = None + try: + torch.set_num_threads(1) + torch.accelerator.set_device_index(rank) + data = SharedHostArena.receive(connection) + control = SharedHostArena.receive(connection) + report = _run_rank(rank, data, control, spec, iterations, barrier) + connection.send(json.dumps(report).encode()) + barrier.wait(180) + except BaseException: + connection.send( + json.dumps({"rank": rank, "error": traceback.format_exc()}).encode() + ) + barrier.abort() + raise + finally: + if control is not None: + control.close() + if data is not None: + data.close() + connection.close() + + +def probe(iterations=8, layers=1): + if iterations < 2 or torch.accelerator.device_count() != 4 or not 1 <= layers <= 6: + raise ValueError("Use at least two iterations and exactly four GPUs") + kinds = ( + (29, 16, 17), + (29, 29, 29), + (16, 16, 16), + (17, 17, 17), + (29, 29, 16), + (17, 17, 29), + ) + specs = tuple( + SharedExpertLayerSpec(5120, 2304, 384, 269, kinds[i]) for i in range(layers) + ) + context = mp.get_context("spawn") + barrier = context.Barrier(5) + connections, processes = [], [] + started = time.monotonic() + with SharedExpertOwner(specs, timeout_ns=60_000_000_000) as owner: + for layer, spec in enumerate(specs): + for expert in range(spec.resident, spec.experts): + offset = owner.coordinator.cold_offset(layer, expert) + row = np.ndarray( + (spec.expert_bytes // 2,), + dtype="= 0: + continue + cold += 1 + offset = owner.coordinator.cold_offset(layer, expert) + row = np.ndarray( + (spec.expert_bytes // 2,), + dtype=" dict[str, int]: + totals = dict.fromkeys(("Size", "Rss", "Pss", "Shared_Dirty", "Private_Dirty"), 0) + selected = False + for line in Path(f"/proc/{pid}/smaps").read_text().splitlines(): + if line and line[0] in "0123456789abcdef" and "-" in line.split()[0]: + selected = f"memfd:{_NAME}" in line + elif selected and ":" in line: + key, value = line.split(":", 1) + if key in totals: + totals[key] += int(value.split()[0]) * 1024 + return totals + + +def memory_available() -> int: + line = next( + line + for line in Path("/proc/meminfo").read_text().splitlines() + if line.startswith("MemAvailable:") + ) + return int(line.split()[1]) * 1024 + + +def cgroup_memory() -> dict[str, int | str]: + """Record cgroup usage separately from this object's proportional size.""" + for line in Path("/proc/self/cgroup").read_text().splitlines(): + hierarchy, controllers, relative = line.split(":", 2) + if hierarchy == "0" and not controllers: + root = Path("/sys/fs/cgroup") / relative.lstrip("/") + if not (root / "memory.current").exists(): + # A cgroup namespace can present the process's group at root. + root = Path("/sys/fs/cgroup") + if (root / "memory.current").exists(): + return { + "path": str(root), + "current": int((root / "memory.current").read_text()), + "limit": (root / "memory.max").read_text().strip(), + } + if "memory" in controllers.split(","): + root = Path("/sys/fs/cgroup/memory") / relative.lstrip("/") + if not (root / "memory.usage_in_bytes").exists(): + root = Path("/sys/fs/cgroup/memory") + return { + "path": str(root), + "current": int((root / "memory.usage_in_bytes").read_text()), + "limit": (root / "memory.limit_in_bytes").read_text().strip(), + } + raise RuntimeError("Memory cgroup accounting is required for this probe") + + +def _consumer(rank, connection, barrier): + host = mapped = view = arena = native = marker = peer_data = None + registered = False + try: + torch.set_num_threads(1) + torch.accelerator.set_device_index(rank) + # The prebuilt helper registers an existing region and wraps its device + # alias; it must not use a tensor helper that silently makes a copy. + from vllm.model_executor.layers.quantization.gguf_expert_cache import _native + + native = _native() + arena = SharedHostArena.receive(connection) + view = memoryview(arena.buffer) + host = torch.frombuffer(view, dtype=torch.uint8) + native.host_register(host.data_ptr(), arena.size) + registered = True + mapped = native.mapped_view(host, rank) + sample = min(_SAMPLE, arena.size) + offsets = {0, arena.size - sample} + if arena.size > 4 * 1024**3 + sample: + offsets.add(4 * 1024**3 + mmap.PAGESIZE) + for offset in sorted(offsets): + checksum = int(mapped[offset : offset + sample].sum().item()) + if checksum != 37 * sample: + raise AssertionError((rank, offset, checksum)) + report = { + "rank": rank, + "pid": os.getpid(), + "identity": vars(arena.identity), + "cpu_pointer": host.data_ptr(), + "device_pointer": mapped.data_ptr(), + "sample_offsets": sorted(offsets), + "samples_match": True, + "shared_mapping": shared_pss(os.getpid()), + } + connection.send(json.dumps(report).encode()) + barrier.wait(180) + # A GPU kernel, rather than merely a DMA copy, writes shared host pages. + start = rank * mmap.PAGESIZE + mapped[start : start + mmap.PAGESIZE].fill_(rank + 1) + torch.accelerator.synchronize() + barrier.wait(180) + barrier.wait(180) + marker = mapped[4 * mmap.PAGESIZE : 5 * mmap.PAGESIZE] + if int(marker.sum().item()) != 99 * mmap.PAGESIZE: + raise AssertionError("CPU updates are not visible to the GPU") + peer = (rank + 1) % 4 + peer_data = mapped[peer * mmap.PAGESIZE : (peer + 1) * mmap.PAGESIZE] + if int(peer_data.sum().item()) != (peer + 1) * mmap.PAGESIZE: + raise AssertionError("Other GPU updates are not visible") + connection.send(json.dumps({"rank": rank, "bidirectional_ok": True}).encode()) + barrier.wait(180) + except BaseException: + connection.send( + json.dumps({"rank": rank, "error": traceback.format_exc()}).encode() + ) + barrier.abort() + raise + finally: + marker = peer_data = mapped = None + if registered: + torch.accelerator.synchronize() + native.host_unregister(host.data_ptr()) + mapped = host = None + if view is not None: + view.release() + if arena is not None: + arena.close() + connection.close() + + +def probe(size: int) -> dict: + if sys.platform != "linux" or not hasattr(os, "memfd_create"): + raise RuntimeError("The CUDA shared-host probe requires Linux") + if not torch.cuda.is_available() or torch.accelerator.device_count() != 4: + raise RuntimeError("Expose exactly the four test GPUs to the probe") + if size < 8 * mmap.PAGESIZE: + raise ValueError("Probe needs at least eight pages") + if memory_available() < size + 12 * 1024**3: + raise RuntimeError( + "Insufficient RAM for arena, worker overhead and 8 GiB reserve" + ) + context = mp.get_context("spawn") + barrier = context.Barrier(5) + processes, connections = [], [] + started = time.monotonic() + baseline_memory = memory_available() + baseline_cgroup = cgroup_memory() + with SharedHostArena.create(size, _NAME) as arena: + view = memoryview(arena.buffer) + host = torch.frombuffer(view, dtype=torch.uint8) + host.fill_(37) + host = None + view.release() + try: + for rank in range(4): + parent, child = socket.socketpair(socket.AF_UNIX, socket.SOCK_SEQPACKET) + parent.settimeout(180) + process = context.Process(target=_consumer, args=(rank, child, barrier)) + process.start() + child.close() + processes.append(process) + connections.append(parent) + arena.send(parent) + reports = [json.loads(conn.recv(65536)) for conn in connections] + if any("error" in report for report in reports): + raise RuntimeError(reports) + if any(report["identity"] != vars(arena.identity) for report in reports): + raise AssertionError("Workers mapped different memory objects") + # Read accounting only after every worker has mapped the object. + for report in reports: + report["shared_mapping"] = shared_pss(report["pid"]) + owner_pss = shared_pss(os.getpid()) + physical_pss = owner_pss["Pss"] + sum( + report["shared_mapping"]["Pss"] for report in reports + ) + # Attribute only this memfd, not CUDA contexts or repeated RSS. + if not 0.9 * arena.size <= physical_pss <= 1.1 * arena.size: + raise AssertionError( + ("Expected one physical arena", physical_pss, size) + ) + mapped_cgroup = cgroup_memory() + mapped_memory = memory_available() + if mapped_memory < 8 * 1024**3: + raise RuntimeError("Shared mapping left less than 8 GiB host reserve") + barrier.wait(180) + barrier.wait(180) + for rank in range(4): + data = arena.buffer[rank * mmap.PAGESIZE : (rank + 1) * mmap.PAGESIZE] + if data != bytes([rank + 1]) * mmap.PAGESIZE: + raise AssertionError("GPU writes are not visible to CPU") + arena.buffer[4 * mmap.PAGESIZE : 5 * mmap.PAGESIZE] = ( + b"\x63" * mmap.PAGESIZE + ) + barrier.wait(180) + updates = [json.loads(conn.recv(65536)) for conn in connections] + if not all(report.get("bidirectional_ok") for report in updates): + raise RuntimeError(updates) + barrier.wait(180) + for process in processes: + process.join(30) + if process.exitcode != 0: + raise RuntimeError(f"Worker exited with {process.exitcode}") + result = { + "status": "passed", + "arena": vars(arena.identity), + "shared_pss_sum_bytes": physical_pss, + "owner_mapping": owner_pss, + "workers": reports, + "bidirectional_checks": updates, + "host_memory_available_before": baseline_memory, + "host_memory_available_mapped": mapped_memory, + "cgroup_before": baseline_cgroup, + "cgroup_mapped": mapped_cgroup, + "cgroup_mapped_increment_bytes": ( + mapped_cgroup["current"] - baseline_cgroup["current"] + ), + } + finally: + barrier.abort() + for process in processes: + if process.is_alive(): + process.terminate() + process.join(5) + if process.is_alive(): + process.kill() + process.join(5) + for connection in connections: + connection.close() + if shared_pss(os.getpid())["Size"]: + raise AssertionError("Owner retained the arena after cleanup") + result["owner_mapping_released"] = True + result["host_memory_available_after_release"] = memory_available() + result["cgroup_after_release"] = cgroup_memory() + result["wall_seconds"] = time.monotonic() - started + return result + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--bytes", type=int, default=256 * 1024**2) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + torch.set_num_threads(2) + result = probe(args.bytes) + args.output.write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/deepseek_v41/tests/bench_g3_down.py b/benchmarks/deepseek_v41/tests/bench_g3_down.py new file mode 100644 index 0000000000..7cd371e33a --- /dev/null +++ b/benchmarks/deepseek_v41/tests/bench_g3_down.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""G3 kernel timing: moe7 (K-split) vs gate_up7 + down_rows (row-split), one rank, M=6.""" + +import torch + +torch.ops.load_library("/home/ymzx/ds41-dev/tools/v41k_g3/build/v41k.so") +dev = torch.device("cuda") +SIZES = {16: 66, 17: 74, 29: 56} +E, m = 128, 6 +routes = m * 6 +torch.manual_seed(0) + + +def timed(fn, n=300): + for _ in range(20): + fn() + a, b = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + a.record() + for _ in range(n): + fn() + b.record() + torch.accelerator.synchronize() + return a.elapsed_time(b) / n * 1000 + + +def rand_bank(kind, experts, rows, blocks): + # Random qs bytes; fp16 scales forced small so dots stay finite. + bank = torch.randint( + 0, 256, (experts, rows, blocks, SIZES[kind]), dtype=torch.uint8, device=dev + ) + if kind in (16, 17): + bank[..., 1] = 0x20 + else: + bank[..., 49] &= 0x0F + bank[..., 51] &= 0x0F + bank[..., 53] &= 0x0F + bank[..., 55] = 0x20 + return bank.view(experts, rows, blocks * SIZES[kind]).contiguous() + + +for gkind, dkind in ((29, 29), (29, 16), (16, 17)): + gate = rand_bank(gkind, E, 576, 20) + up = rand_bank(gkind, E, 576, 20) + dcol = rand_bank(dkind, E, 5120, 3) + drow = rand_bank(dkind, E, 1280, 9) + gbook, dbook = torch.ops.v41k.table(gkind, dev), torch.ops.v41k.table(dkind, dev) + x = (torch.randn(m, 5120, device=dev) * 0.5).half() + ids = ( + torch.stack([torch.randperm(E, device=dev)[:6] for _ in range(m)]) + .int() + .contiguous() + ) + prob = torch.full((m, 6), 0.25, device=dev) + xq = torch.empty((m, 160, 32), dtype=torch.int8, device=dev) + xd = torch.empty((m, 160), dtype=torch.float32, device=dev) + torch.ops.v41k.q8_aligned(xq, xd, x) + hq = torch.empty((routes, 18, 32), dtype=torch.int8, device=dev) + hd = torch.empty((routes, 18), dtype=torch.float32, device=dev) + out = torch.empty(m, 5120, dtype=torch.float16, device=dev) + plan = torch.zeros(512, dtype=torch.int32, device=dev) + rows, warps = (16, 8) if dkind == 29 else (8, 4) + t_moe7 = timed( + lambda: torch.ops.v41k.moe7( + out, + xq, + xd, + hq, + hd, + ids, + prob, + gate, + up, + dcol, + gkind, + dkind, + 64, + gbook, + dbook, + rows, + warps, + plan, + ) + ) + t_gu = timed( + lambda: torch.ops.v41k.gate_up7( + xq, xd, hq, hd, ids, prob, gate, up, gkind, gbook, plan + ) + ) + stride = (routes * 648 + 255) // 256 * 256 + gathered = torch.zeros(4 * stride, dtype=torch.uint8, device=dev) + res = [f"g{gkind}/d{dkind}: moe7 {t_moe7:.1f} us, gate_up7 {t_gu:.1f} us"] + for r_tile, w in ((4, 8), (2, 8), (2, 4), (8, 8)): + o = torch.zeros(m, 5120, dtype=torch.float16, device=dev) + t = timed( + lambda: torch.ops.v41k.down_rows( + o, gathered, stride, routes, drow, dkind, 0, dbook, r_tile, w, plan + ) + ) + res.append(f"down_rows[{r_tile}x{w}] {t:.1f}") + print(" | ".join(res)) diff --git a/benchmarks/deepseek_v41/tests/bench_index_decode.py b/benchmarks/deepseek_v41/tests/bench_index_decode.py new file mode 100644 index 0000000000..9a791dd269 --- /dev/null +++ b/benchmarks/deepseek_v41/tests/bench_index_decode.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os + +import torch + +os.environ.setdefault( + "VLLM_V41K_LIB", "/home/ymzx/ds41-dev/tools/v41k_g3/build/v41k.so" +) +import vllm.models.deepseek_v41.attention_ops as ops +from vllm.models.deepseek_v41.index_topk import index_topk +from vllm.models.deepseek_v41.kv_pack import pack_rows + +torch.manual_seed(0) +H, T = 8, 6 + + +def tm(fn, n=10): + fn() + torch.accelerator.synchronize() + a, b = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + a.record() + for _ in range(n): + fn() + b.record() + torch.accelerator.synchronize() + return a.elapsed_time(b) / n + + +for rows, start in ((16384, 15000), (524288, 15000), (524288, 400000)): + packed = pack_rows((torch.randn(rows, 128, device="cuda") * 0.5).half(), False) + q = (torch.randn(T, H, 128, device="cuda") * 0.3).half() + w = (torch.randn(T, H, device="cuda") * 0.2).half() + pos = torch.arange(start, start + T, device="cuda") + valid = torch.ones(T, dtype=torch.bool, device="cuda") + scores = torch.empty(T, rows, dtype=torch.float16, device="cuda") + ident = lambda x: x + f = lambda: ops.select_packed_positions(q, w, packed, pos, valid, 1, ident) + print( + "rows %d start %d select_packed (decode, incl topk): %.3f ms" + % (rows, start, tm(f)) + ) + res = torch.empty(T, 512, dtype=torch.int32, device="cuda") + print( + " index_topk: %.3f ms" + % tm(lambda: index_topk(res, scores, pos, valid, 1, 512)) + ) diff --git a/benchmarks/deepseek_v41/tests/test_fused_sparse.py b/benchmarks/deepseek_v41/tests/test_fused_sparse.py new file mode 100644 index 0000000000..efbfdf0350 --- /dev/null +++ b/benchmarks/deepseek_v41/tests/test_fused_sparse.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused sparse attention vs the gather + bmm path (ring window + compressed top-k).""" + +import torch + +import vllm.models.deepseek_v41.attention_ops as ops + +torch.manual_seed(0) +dev = "cuda" +H, BS, R = 16, 64, 2 +ok = True +for T, start, comp_rows in ((300, 0, 0), (512, 700, 600), (2048, 5000, 2600)): + total = start + T + ring = (torch.randn(ops.WINDOW_RING, 512, device=dev) * 0.3).half() + values = (torch.randn(T, 512, device=dev) * 0.3).half() + q = (torch.randn(T, H, 512, device=dev) * 0.3).half() + pos = torch.arange(start, total, device=dev) + valid = torch.ones(T, dtype=torch.bool, device=dev) + valid[5] = False + sink = torch.randn(H, device=dev) + table = torch.arange(total // BS + 2, dtype=torch.int32, device=dev)[None] + kwargs = dict(window_block=BS, window_ring=ring, window_values=values) + if comp_rows: + ccache = (torch.randn(total // BS + 2, BS // R, 512, device=dev) * 0.3).half() + n_c = (pos + 1) // R + topk = (torch.rand(T, 512, device=dev) * n_c[:, None]).long() + topk = torch.where(torch.rand(T, 512, device=dev) < 0.05, -1, topk).int() + kwargs.update( + compressed_cache=ccache, + compressed_table=table, + compressed_block=BS, + ratio=R, + topk=topk.contiguous(), + ) + ops._FUSED_SPARSE = False + ref = ops.attend_paged_latents(q, pos, valid, None, table, sink, **kwargs) + ops._FUSED_SPARSE = True + new = ops.attend_paged_latents(q, pos, valid, None, table, sink, **kwargs) + diff = (new.float() - ref.float()).abs() + rel = (diff.norm() / ref.float().norm()).item() + good = torch.isfinite(new).all().item() and rel < 5e-3 + ok &= good + for fused in (False, True): + ops._FUSED_SPARSE = fused + for _ in range(3): + ops.attend_paged_latents(q, pos, valid, None, table, sink, **kwargs) + torch.accelerator.synchronize() + a, b = ( + torch.cuda.Event(enable_timing=True), + torch.cuda.Event(enable_timing=True), + ) + a.record() + for _ in range(5): + ops.attend_paged_latents(q, pos, valid, None, table, sink, **kwargs) + b.record() + torch.accelerator.synchronize() + print( + " T=%d compressed=%s fused=%s %.2f ms" + % (T, bool(comp_rows), fused, a.elapsed_time(b) / 5) + ) + print("T=%d rel=%.2e max=%.3g ok=%s" % (T, rel, diff.max().item(), good)) +print("ALL OK" if ok else "FAILED") diff --git a/benchmarks/deepseek_v41/tests/test_g1mem.py b/benchmarks/deepseek_v41/tests/test_g1mem.py new file mode 100644 index 0000000000..e968df999f --- /dev/null +++ b/benchmarks/deepseek_v41/tests/test_g1mem.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Micro-test G1 memory changes against the previous implementations.""" + +import torch + +import vllm.model_executor.kernels.mhc.deepseek_v41_fused as fused +from vllm.model_executor.kernels.mhc.deepseek_v41 import expand_hc +from vllm.model_executor.layers.quantization.gguf_expert_cache import _unroute +from vllm.model_executor.layers.quantization.gguf_shared_prefill import ( + _v41_prefill_accumulate, +) +from vllm.triton_utils import triton + +torch.manual_seed(0) +dev = "cuda" +H = 5120 +ok = True + +# 1. expand_hc in place == out of place. +for T in (1, 7, 300, 3000): + res = torch.randn(T, 4, H, device=dev) * 3 + x = torch.randn(T, H, device=dev).half() + post = torch.rand(T, 4, device=dev) * 2 + comb = torch.rand(T, 4, 4, device=dev) + ref = expand_hc(x, res, post, comb) + out = expand_hc(x, res.clone(), post, comb, inplace=True) + same = torch.equal(ref, out) + ok &= same + print( + "expand T=%d equal=%s maxdiff=%.3g" % (T, same, (ref - out).abs().max().item()) + ) + +# 2. windowed hc_transition == single window. +for T, mix in ((3000, True), (5000, False), (4096, True)): + res = torch.randn(T, 4, H, device=dev) + y = torch.randn(T, H, device=dev).half() + post = torch.rand(T, 4, device=dev) * 2 + comb = torch.rand(T, 4, 4, device=dev) + cpre = torch.rand(T, 4, device=dev) + norm = torch.rand(H, device=dev) + fn = torch.randn(24, 4 * H, device=dev) * 0.01 if mix else None + scale = torch.rand(3, device=dev) if mix else None + base = torch.randn(24, device=dev) if mix else None + outs = [] + for window in (1 << 30, 2048, 1000): + fused._WINDOW = window + r = res.clone() + o, m = fused.hc_transition(r, y, post, comb, cpre, norm, fn, scale, base) + outs.append((r, o, m)) + for r, o, m in outs[1:]: + same = torch.equal(r, outs[0][0]) and torch.equal(o, outs[0][1]) + if mix: + same &= all(torch.equal(a, b) for a, b in zip(m, outs[0][2])) + ok &= same + print("transition T=%d mix=%s equal=%s" % (T, mix, same)) +fused._WINDOW = 2048 + +# 3. per-tile accumulation vs the old full down + _unroute. +E, TILE = 384, 32 +for T in (1, 13, 2048, 4096): + ids = torch.stack([torch.randperm(E, device=dev)[:6] for _ in range(T)]).int() + flat = ids.reshape(-1) + sorted_ids, order = flat.sort(stable=True) + offsets = torch.searchsorted( + sorted_ids, + torch.arange(E + 1, device=dev, dtype=sorted_ids.dtype), + out_int32=True, + ) + inverse = torch.empty_like(order, dtype=torch.int32) + inverse[order] = torch.arange(order.numel(), device=dev, dtype=torch.int32) + down = (torch.randn(T * 6, H, device=dev) * 0.5).half() # sorted-order rows + ref = torch.empty(T, H, device=dev, dtype=torch.float16) + ones = torch.ones(T, 6, device=dev) + _unroute[(T, triton.cdiv(H, 1024))](down, ones, inverse, ref, H, TOPK=6, BLOCK=1024) + acc = torch.zeros(T, H, device=dev) + for f in range(0, E, TILE): + bounds = offsets[f : f + TILE + 1] + scratch = torch.full_like(down, float("nan")) + lo, hi = bounds[0].item(), bounds[-1].item() + scratch[lo:hi] = down[lo:hi] # only this tile's rows are valid + _v41_prefill_accumulate[(T, triton.cdiv(H, 1024))]( + scratch, inverse, bounds, acc, T, EXPERTS=TILE, TOPK=6, H=H, BLOCK=1024 + ) + new = acc.half() + diff = (new.float() - ref.float()).abs() + rel = diff.max().item() / ref.float().abs().max().item() + mism = (new != ref).float().mean().item() + good = torch.isfinite(new).all().item() and rel < 1e-3 + ok &= good + print( + "accumulate T=%d max|d|=%.3g rel=%.2g mismatch-frac=%.4f" + % (T, diff.max().item(), rel, mism) + ) + +print("ALL OK" if ok else "FAILED") diff --git a/benchmarks/deepseek_v41/tests/test_g3_down.py b/benchmarks/deepseek_v41/tests/test_g3_down.py new file mode 100644 index 0000000000..64d83f2a8f --- /dev/null +++ b/benchmarks/deepseek_v41/tests/test_g3_down.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""G3 micro-test: row-split down (gate_up7 + down_rows) vs the K-split moe7, 4 ranks simulated.""" + +import glob +import os +import sys + +import gguf +import numpy as np +import torch + +lib = ( + sys.argv[1] + if len(sys.argv) > 1 + else "/home/ymzx/ds41-dev/tools/v41k_g3/build/v41k.so" +) +torch.ops.load_library(lib) +dev = torch.device("cuda") +SIZES = {16: (256, 66), 17: (256, 74), 29: (256, 56)} + +tensors = {} +for f in sorted( + glob.glob("/home/ymzx/ds41-models/gguf/DeepSeek-V4.1-Flash-*-0000*-of-00004.gguf") +): + for t in gguf.GGUFReader(f).tensors: + if "exps" in t.name: + tensors[t.name] = t +picked = {} +for name, t in tensors.items(): + if name.endswith("ffn_down_exps.weight"): + layer = int(name.split(".")[1]) + kind = int(t.tensor_type) + picked.setdefault(kind, layer) +print("layers by down type", picked) + +E = 12 +torch.manual_seed(0) +ok = True +for dkind, layer in sorted(picked.items()): + g = tensors[f"blk.{layer}.ffn_gate_exps.weight"] + u = tensors[f"blk.{layer}.ffn_up_exps.weight"] + d = tensors[f"blk.{layer}.ffn_down_exps.weight"] + gkind = int(g.tensor_type) + experts = np.random.default_rng(layer).choice(g.data.shape[0], E, replace=False) + gate = torch.from_numpy(np.ascontiguousarray(g.data[experts])).to( + dev + ) # [E, 2304, rb] + up = torch.from_numpy(np.ascontiguousarray(u.data[experts])).to(dev) + down = torch.from_numpy(np.ascontiguousarray(d.data[experts])).to( + dev + ) # [E, 5120, rb] + block, bb = SIZES[dkind] + gbook = torch.ops.v41k.table(gkind, dev) + dbook = torch.ops.v41k.table(dkind, dev) + for m in (1, 3, 6): + x = (torch.randn(m, 5120, device=dev) * 0.5).half() + # Distinct per token; expert 0 shared by every token (<= 6 routes). + ids = torch.stack( + [ + torch.cat( + [ + torch.zeros(1, device=dev, dtype=torch.long), + torch.randperm(E - 1, device=dev)[:5] + 1, + ] + ) + for _ in range(m) + ] + ).int() + prob = torch.rand(m, 6, device=dev) + prob = (prob / prob.sum(1, keepdim=True) * 1.5).contiguous() + ids = ids.contiguous() + xq = torch.empty((m, 160, 32), dtype=torch.int8, device=dev) + xd = torch.empty((m, 160), dtype=torch.float32, device=dev) + torch.ops.v41k.q8_aligned(xq, xd, x) + routes = m * 6 + stride = (routes * 648 + 255) // 256 * 256 + gathered = torch.zeros(4, stride, dtype=torch.uint8, device=dev) + ref = torch.zeros(m, 5120, device=dev) + new = torch.zeros(m, 5120, device=dev) + rows, warps = (16, 8) if dkind == 29 else (8, 4) + nrows, nwarps = [int(v) for v in os.environ.get("ROWS_TILE", "4,8").split(",")] + for r in range(4): + gr = gate[:, 576 * r : 576 * (r + 1)].contiguous() + ur = up[:, 576 * r : 576 * (r + 1)].contiguous() + first, end = 576 * r // block, (576 * (r + 1) + block - 1) // block + left = 576 * r - first * block + dc = torch.zeros( + E, 5120, (end - first) * bb + 512, dtype=torch.uint8, device=dev + ) + dc = dc[:, :, : (end - first) * bb] + dc.copy_(down[:, :, first * bb : end * bb]) + dc = dc.contiguous() + hq = torch.empty((routes, 18, 32), dtype=torch.int8, device=dev) + hd = torch.empty((routes, 18), dtype=torch.float32, device=dev) + out = torch.empty(m, 5120, dtype=torch.float16, device=dev) + plan = torch.zeros(512, dtype=torch.int32, device=dev) + torch.ops.v41k.moe7( + out, + xq, + xd, + hq, + hd, + ids, + prob, + gr, + ur, + dc, + gkind, + dkind, + left, + gbook, + dbook, + rows, + warps, + plan, + ) + ref += out.float() + # Row mode: gate/up into this rank's packed slot. + q = gathered[r, : routes * 576].view(torch.int8).view(routes, 18, 32) + dd = ( + gathered[r, routes * 576 : routes * 648] + .view(torch.float32) + .view(routes, 18) + ) + plan2 = torch.zeros(512, dtype=torch.int32, device=dev) + torch.ops.v41k.gate_up7( + xq, xd, q, dd, ids, prob, gr, ur, gkind, gbook, plan2 + ) + same_h = torch.equal(q, hq) and torch.equal(dd, hd) + ok &= same_h + for r in range(4): + dr = torch.zeros( + E, 1280, down.shape[2] + 64, dtype=torch.uint8, device=dev + )[:, :, : down.shape[2]] + dr.copy_(down[:, 1280 * r : 1280 * (r + 1)]) + dr = dr.contiguous() + out = torch.zeros(m, 5120, dtype=torch.float16, device=dev) + plan2 = torch.zeros(512, dtype=torch.int32, device=dev) + torch.ops.v41k.gate_up7( + xq, + xd, + torch.empty((routes, 18, 32), dtype=torch.int8, device=dev), + torch.empty((routes, 18), dtype=torch.float32, device=dev), + ids, + prob, + gate[:, :576].contiguous(), + up[:, :576].contiguous(), + gkind, + gbook, + plan2, + ) # only fills the plan + torch.ops.v41k.down_rows( + out, + gathered, + stride, + routes, + dr, + dkind, + 1280 * r, + dbook, + nrows, + nwarps, + plan2, + ) + others = torch.ones(5120, dtype=torch.bool, device=dev) + others[1280 * r : 1280 * (r + 1)] = False + ok &= bool((out[:, others] == 0).all()) + new += out.float() + diff = (new - ref).abs() + rel = (diff.norm() / ref.norm()).item() + good = torch.isfinite(new).all().item() and rel < 2e-3 + ok &= good + print( + "down %d gate %d layer %d m=%d same_h=%s rel=%.2e max=%.3g |ref|max=%.3g" + % ( + dkind, + gkind, + layer, + m, + same_h, + rel, + diff.max().item(), + ref.abs().max().item(), + ) + ) +print("ALL OK" if ok else "FAILED") diff --git a/benchmarks/deepseek_v41/tests/test_hc_tok.py b/benchmarks/deepseek_v41/tests/test_hc_tok.py new file mode 100644 index 0000000000..82cf0952e3 --- /dev/null +++ b/benchmarks/deepseek_v41/tests/test_hc_tok.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Token-blocked HC transition vs the per-token kernels (T > 8).""" + +import torch + +import vllm.model_executor.kernels.mhc.deepseek_v41_fused as fused + +torch.manual_seed(0) +dev = "cuda" +H = 5120 +ok = True +for T, mix in ((37, True), (2048, True), (8192, True), (8192, False)): + res = torch.randn(T, 4, H, device=dev) + y = torch.randn(T, H, device=dev).half() + post = torch.rand(T, 4, device=dev) * 2 + comb = torch.rand(T, 4, 4, device=dev) + cpre = torch.rand(T, 4, device=dev) + norm = torch.rand(H, device=dev) + fn = torch.randn(24, 4 * H, device=dev) * 0.01 if mix else None + scale = torch.rand(3, device=dev) if mix else None + base = torch.randn(24, device=dev) if mix else None + outs, times = [], [] + for blocked in (False, True): + fused._TOKEN_BLOCKED = blocked + r = res.clone() + o, m = fused.hc_transition(r, y, post, comb, cpre, norm, fn, scale, base) + outs.append((r, o, m)) + rr = res.clone() + torch.accelerator.synchronize() + a, b = ( + torch.cuda.Event(enable_timing=True), + torch.cuda.Event(enable_timing=True), + ) + a.record() + for _ in range(5): + fused.hc_transition(rr, y, post, comb, cpre, norm, fn, scale, base) + b.record() + torch.accelerator.synchronize() + times.append(a.elapsed_time(b) / 5) + (r0, o0, m0), (r1, o1, m1) = outs + same_res = torch.equal(r0, r1) + dout = (o0.float() - o1.float()).abs().max().item() + dm = max((a - b).abs().max().item() for a, b in zip(m0, m1)) if mix else 0.0 + good = same_res and dout < 2e-2 and dm < 1e-4 + ok &= good + print( + "T=%d mix=%s residual_equal=%s out_maxdiff=%.2e mix_maxdiff=%.2e old %.2f ms new %.2f ms" + % (T, mix, same_res, dout, dm, times[0], times[1]) + ) +print("ALL OK" if ok else "FAILED") diff --git a/benchmarks/deepseek_v41/tests/test_index_packed.py b/benchmarks/deepseek_v41/tests/test_index_packed.py new file mode 100644 index 0000000000..aaf0a25c75 --- /dev/null +++ b/benchmarks/deepseek_v41/tests/test_index_packed.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Packed fused indexer vs the gather + mm path: top-512 agreement and timing.""" + +import os + +import torch + +os.environ.setdefault( + "VLLM_V41K_LIB", "/home/ymzx/ds41-dev/tools/v41k_g3/build/v41k.so" +) +import vllm.models.deepseek_v41.attention_ops as ops +from vllm.models.deepseek_v41.kv_pack import pack_rows +from vllm.models.deepseek_v41.quant_rope import attention_qdq + +torch.manual_seed(0) +H = 8 +ident = lambda x: x +for T, start, ratio, rows in ( + (6, 9000, 2, 8192), + (6, 15000, 1, 16384), + (512, 3000, 2, 8192), + (2048, 12000, 1, 16384), + (6, 400000, 1, 524288), +): + raw = (torch.randn(rows, 128, device="cuda") * 0.5).half() + ref_keys = attention_qdq(raw, fp4=True, group=32) + packed = pack_rows(raw, False) + q = attention_qdq( + (torch.randn(T, H, 128, device="cuda") * 0.5) + .half() + .reshape(T * H, 128) + .contiguous(), + fp4=True, + group=32, + ).view(T, H, 128) + w = (torch.randn(T, H, device="cuda") * 0.2).half() + pos = torch.arange(start, start + T, device="cuda") + valid = torch.ones(T, dtype=torch.bool, device="cuda") + old = ( + ops.select_compressed_positions(q, w, ref_keys, pos, valid, ratio, ident) + if rows <= 16384 + else None + ) + new = ops.select_packed_positions(q, w, packed, pos, valid, ratio, ident) + if old is not None: + agree = [] + for r in range(T): + a, b = set(old[r].tolist()) - {-1}, set(new[r].tolist()) - {-1} + agree.append(len(a & b) / max(1, len(a))) + print( + "T=%d start=%d ratio=%d: top-512 overlap min %.4f mean %.5f exact rows %d/%d" + % ( + T, + start, + ratio, + min(agree), + sum(agree) / T, + sum(1 for x in agree if x == 1.0), + T, + ) + ) + for name, fn in ( + ( + "old", + lambda: ops.select_compressed_positions( + q, w, ref_keys, pos, valid, ratio, ident + ), + ) + if rows <= 16384 + else ("old", None), + ( + "new", + lambda: ops.select_packed_positions(q, w, packed, pos, valid, ratio, ident), + ), + ): + if fn is None: + continue + fn() + torch.accelerator.synchronize() + a, b = ( + torch.cuda.Event(enable_timing=True), + torch.cuda.Event(enable_timing=True), + ) + a.record() + for _ in range(5): + fn() + b.record() + torch.accelerator.synchronize() + print(" %s %.3f ms" % (name, a.elapsed_time(b) / 5)) diff --git a/benchmarks/deepseek_v41/tests/test_kv_pack.py b/benchmarks/deepseek_v41/tests/test_kv_pack.py new file mode 100644 index 0000000000..3ff0920551 --- /dev/null +++ b/benchmarks/deepseek_v41/tests/test_kv_pack.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Packed KV rows reproduce attention_qdq FP16 values bit for bit.""" + +import torch + +from vllm.models.deepseek_v41.kv_pack import pack_rows, unpack_rows +from vllm.models.deepseek_v41.quant_rope import attention_qdq + +torch.manual_seed(0) +ok = True +for scale in (1e-6, 1e-3, 0.05, 1.0, 30.0, 3000.0): + for compressed in (True, False): + dim = 512 if compressed else 128 + x = (torch.randn(4096, dim, device="cuda") * scale).half() + x[0].zero_() + x[1, :17] = 0 # zero groups + ref = attention_qdq( + x, fp4=True, group=16 if compressed else 32, e4m3_scale=compressed + ) + got = unpack_rows(pack_rows(x, compressed), compressed) + same = torch.equal(ref.view(torch.int16), got.view(torch.int16)) + ok &= same + print( + "scale %g compressed=%s bit-exact=%s mismatches=%d" + % ( + scale, + compressed, + same, + int((ref.view(torch.int16) != got.view(torch.int16)).sum()), + ) + ) +print("ALL OK" if ok else "FAILED") diff --git a/benchmarks/deepseek_v41/tests/test_packed_attn.py b/benchmarks/deepseek_v41/tests/test_packed_attn.py new file mode 100644 index 0000000000..695c9b3f32 --- /dev/null +++ b/benchmarks/deepseek_v41/tests/test_packed_attn.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Packed compressed rows (hot ring + staging / selected buffer) vs the FP16 paged cache.""" + +from types import SimpleNamespace + +import torch + +import vllm.models.deepseek_v41.attention_ops as ops +from vllm.models.deepseek_v41.kv_pack import pack_rows +from vllm.models.deepseek_v41.kv_store import _fill_selected +from vllm.models.deepseek_v41.quant_rope import attention_qdq + +torch.manual_seed(0) +dev = "cuda" +H, BS, R = 16, 64, 2 +ok = True +for T, start, hot_rows in ( + (300, 5000, 1024), + (6, 9000, 1024), + (6, 3000, 4096), + (2048, 12000, 2048), +): + total = start + T + rows = total // R + 1 + raw = (torch.randn(rows, 512, device=dev) * 0.3).half() + fp16 = attention_qdq(raw, fp4=True, group=16, e4m3_scale=True) + packed = pack_rows(raw, True) + # paged FP16 reference cache: identity block table + ccache = torch.zeros( + rows // (BS // R) + 2, BS // R, 512, dtype=torch.float16, device=dev + ) + ccache.view(-1, 512)[:rows] = fp16 + table = torch.arange(ccache.shape[0], dtype=torch.int32, device=dev)[None] + ring = (torch.randn(ops.WINDOW_RING, 512, device=dev) * 0.3).half() + values = (torch.randn(T, 512, device=dev) * 0.3).half() + q = (torch.randn(T, H, 512, device=dev) * 0.3).half() + pos = torch.arange(start, total, device=dev) + valid = torch.ones(T, dtype=torch.bool, device=dev) + sink = torch.randn(H, device=dev) + n_c = (pos + 1) // R + topk = (torch.rand(T, 512, device=dev) * n_c[:, None]).long() + topk = ( + torch.where(torch.rand(T, 512, device=dev) < 0.05, -1, topk).int().contiguous() + ) + ref = ops.attend_paged_latents( + q, + pos, + valid, + None, + table, + sink, + window_block=BS, + window_ring=ring, + window_values=values, + compressed_cache=ccache, + compressed_table=table, + compressed_block=BS, + ratio=R, + topk=topk, + ) + newest = int(total) // R - 1 + hot = torch.zeros(hot_rows, 288, dtype=torch.uint8, device=dev) + idx = torch.arange(max(0, newest - hot_rows + 1), newest + 1, device=dev) + hot[idx % hot_rows] = packed[idx] + old = newest - hot_rows + 1 + staging = torch.zeros(max(rows, 1), 288, dtype=torch.uint8, device=dev) + if old > 0: + staging[:old] = packed[:old] + selected = torch.zeros(8, 512, 288, dtype=torch.uint8, device=dev) + store = SimpleNamespace( + hot={0: hot}, staging=staging, hot_rows={0: hot_rows}, selected=selected + ) + if T <= 8: + _fill_selected[(T, 8)]( + topk, pos, hot, packed, selected, T, hot_rows, RATIO=R, TOPK=512, BLOCK=64 + ) + new = ops.attend_paged_latents( + q, + pos, + valid, + None, + table, + sink, + window_block=BS, + window_ring=ring, + window_values=values, + ratio=R, + topk=topk, + packed=(store, 0), + ) + same = torch.equal(ref, new) + ok &= same + print( + "T=%d start=%d hot_rows=%d mode=%s bit-exact=%s maxdiff=%.3g" + % ( + T, + start, + hot_rows, + "selected" if T <= 8 else "hot+staging", + same, + (ref.float() - new.float()).abs().max().item(), + ) + ) +print("ALL OK" if ok else "FAILED") diff --git a/benchmarks/deepseek_v41/tests/test_peer_gather.py b/benchmarks/deepseek_v41/tests/test_peer_gather.py new file mode 100644 index 0000000000..7baa79619a --- /dev/null +++ b/benchmarks/deepseek_v41/tests/test_peer_gather.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""4-process test of v41k.peer_gather: correctness over many epochs + latency (eager and graph).""" + +import os + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +LIB = "/home/ymzx/ds41-dev/tools/v41k_g3/build/v41k.so" + + +class View: + def __init__(self, ptr, n): + self.__cuda_array_interface__ = { + "shape": (n,), + "typestr": "|u1", + "data": (ptr, False), + "version": 3, + } + + +def worker(rank): + os.environ.update(MASTER_ADDR="127.0.0.1", MASTER_PORT="29533") + torch.accelerator.set_device_index(rank) + dist.init_process_group("gloo", rank=rank, world_size=4) + torch.ops.load_library(LIB) + from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce + + stride_max = (36 * 648 + 255) // 256 * 256 + flags = 4 * stride_max + nbytes = flags + 256 + ptrs = CustomAllreduce.create_shared_buffer(nbytes, None) + own = torch.as_tensor(View(ptrs[rank], nbytes), device=f"cuda:{rank}") + own.zero_() + torch.accelerator.synchronize() + dist.barrier() + epoch = torch.zeros(1, dtype=torch.int32, device=f"cuda:{rank}") + ok = True + for it in range(300): + routes = 6 * (1 + it % 6) + stride = (routes * 648 + 255) // 256 * 256 + own[rank * stride : (rank + 1) * stride].fill_((it * 7 + rank * 31) % 251) + torch.ops.v41k.peer_gather(*ptrs, rank, stride, flags, epoch) + for r in range(4): + want = (it * 7 + r * 31) % 251 + if not bool((own[r * stride : (r + 1) * stride] == want).all()): + ok = False + # Mimic the MoE output all-reduce that bounds rank skew. + torch.accelerator.synchronize() + dist.barrier() + stride = (36 * 648 + 255) // 256 * 256 + for _ in range(50): + torch.ops.v41k.peer_gather(*ptrs, rank, stride, flags, epoch) + torch.accelerator.synchronize() + dist.barrier() + a, b = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + a.record() + for _ in range(1000): + torch.ops.v41k.peer_gather(*ptrs, rank, stride, flags, epoch) + b.record() + torch.accelerator.synchronize() + eager = a.elapsed_time(b) + dist.barrier() + g = torch.cuda.CUDAGraph() + s = torch.cuda.Stream() + with torch.cuda.stream(s), torch.cuda.graph(g): + for _ in range(40): + torch.ops.v41k.peer_gather(*ptrs, rank, stride, flags, epoch) + torch.accelerator.synchronize() + dist.barrier() + a.record() + for _ in range(25): + g.replay() + b.record() + torch.accelerator.synchronize() + graph = a.elapsed_time(b) / 1000 * 1000 + print( + f"rank {rank} ok={ok} eager {eager:.2f} us/call graph {graph:.2f} us/call epoch={int(epoch.item())}", + flush=True, + ) + dist.barrier() + + +if __name__ == "__main__": + mp.spawn(worker, nprocs=4) diff --git a/benchmarks/deepseek_v41/tests/test_rope_tail.py b/benchmarks/deepseek_v41/tests/test_rope_tail.py new file mode 100644 index 0000000000..54e680b4d8 --- /dev/null +++ b/benchmarks/deepseek_v41/tests/test_rope_tail.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch + +from vllm.models.deepseek_v41.quant_rope import apply_rope, rope_table + +torch.manual_seed(0) +for compressed in (False, True): + small = rope_table(compressed=compressed) # 16384 rows + freq row + full = rope_table(600000, compressed=compressed) # reference big table + x = torch.randn(4096, 1, 512, device="cuda").half() + for lo in (0, 12288, 16380, 100000, 520000): + pos = torch.arange(lo, lo + 4096, device="cuda") + a = apply_rope(x, pos, small) + b = apply_rope(x, pos, full) + print( + "compressed=%s pos %d..: bit-exact=%s maxdiff=%.3g" + % ( + compressed, + lo, + torch.equal(a, b), + (a.float() - b.float()).abs().max().item(), + ) + ) diff --git a/benchmarks/deepseek_v41/tests/test_topk_par.py b/benchmarks/deepseek_v41/tests/test_topk_par.py new file mode 100644 index 0000000000..1f0b6191a5 --- /dev/null +++ b/benchmarks/deepseek_v41/tests/test_topk_par.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch + +from vllm.models.deepseek_v41.index_topk import index_topk, index_topk_parallel + +torch.manual_seed(0) + + +def tm(fn, n=10): + fn() + torch.accelerator.synchronize() + a, b = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + a.record() + for _ in range(n): + fn() + b.record() + torch.accelerator.synchronize() + return a.elapsed_time(b) / n + + +ok = True +for rows, width, start in ( + (6, 524288, 400000), + (6, 524288, 300), + (3, 131072, 100000), + (6, 65536, 65000), +): + sc = torch.randn(rows, width, device="cuda").half() + sc[:, ::7] = 0.5 # ties + pos = torch.arange(start, start + rows, device="cuda") + valid = torch.ones(rows, dtype=torch.bool, device="cuda") + a = torch.empty(rows, 512, dtype=torch.int32, device="cuda") + b = torch.empty_like(a) + index_topk(a, sc, pos, valid, 1, 512) + index_topk_parallel(b, sc, pos, valid, 1, 512) + same = torch.equal(a, b) + ok &= same + print( + "rows %d width %d start %d equal=%s single %.3f ms parallel %.3f ms" + % ( + rows, + width, + start, + same, + tm(lambda: index_topk(a, sc, pos, valid, 1, 512)), + tm(lambda: index_topk_parallel(b, sc, pos, valid, 1, 512)), + ) + ) +print("ALL OK" if ok else "FAILED") diff --git a/benchmarks/deepseek_v41/tests/test_window_ring.py b/benchmarks/deepseek_v41/tests/test_window_ring.py new file mode 100644 index 0000000000..39a73b3014 --- /dev/null +++ b/benchmarks/deepseek_v41/tests/test_window_ring.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""G4 window ring vs paged window cache: identical attention over chunked prefill + decode.""" + +import torch + +from vllm.models.deepseek_v41.attention_ops import ( + WINDOW_RING, + attend_paged_latents, + insert_latents, + insert_window_ring, +) + +torch.manual_seed(0) +dev = "cuda" +H, BS = 16, 64 +total = 1200 +cache = torch.zeros(total // BS + 2, BS, 512, dtype=torch.float16, device=dev) +table = torch.arange(cache.shape[0], dtype=torch.int32, device=dev)[None] +ring = torch.zeros(WINDOW_RING, 512, dtype=torch.float16, device=dev) +sink = torch.randn(H, device=dev) +ok, start = True, 0 +for T in (300, 200, 6, 6, 3, 400, 1, 5): + pos = torch.arange(start, start + T, device=dev) + valid = torch.ones(T, dtype=torch.bool, device=dev) + values = (torch.randn(T, 512, device=dev) * 0.3).half() + q = (torch.randn(T, H, 512, device=dev) * 0.3).half() + insert_latents(cache, values, pos, valid, table, block_size=BS) + old = attend_paged_latents(q, pos, valid, cache, table, sink, window_block=BS) + new = attend_paged_latents( + q, + pos, + valid, + None, + table, + sink, + window_block=BS, + window_ring=ring, + window_values=values, + ) + insert_window_ring(ring, values, pos, valid) + same = torch.equal(old, new) + ok &= same + print( + "chunk T=%d start=%d equal=%s maxdiff=%.3g" + % (T, start, same, (old.float() - new.float()).abs().max().item()) + ) + start += T +print("ALL OK" if ok else "FAILED") diff --git a/benchmarks/deepseek_v41/tools/api_multiturn.py b/benchmarks/deepseek_v41/tools/api_multiturn.py new file mode 100644 index 0000000000..1f6d08ae3c --- /dev/null +++ b/benchmarks/deepseek_v41/tools/api_multiturn.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Multi-turn chat over the OpenAI API (streaming): TTFT and decode tok/s per turn.""" + +import json +import sys +import time +import urllib.request + +URL = sys.argv[1] if len(sys.argv) > 1 else "http://127.0.0.1:18000/v1/chat/completions" +DOC = open("/home/ymzx/ds41-dev/src/README.md", encoding="utf-8").read()[:12000] + + +def chat(messages, max_tokens=200): + body = json.dumps( + dict( + model="deepseek-v4.1-flash", + messages=messages, + max_tokens=max_tokens, + temperature=0, + stream=True, + stream_options={"include_usage": True}, + ) + ).encode() + req = urllib.request.Request(URL, body, {"Content-Type": "application/json"}) + t0 = time.perf_counter() + first = None + text = [] + usage = {} + with urllib.request.urlopen(req, timeout=600) as r: + for line in r: + line = line.decode().strip() + if not line.startswith("data: ") or line == "data: [DONE]": + continue + d = json.loads(line[6:]) + if d.get("usage"): + usage = d["usage"] + for c in d.get("choices", []): + piece = c.get("delta", {}).get("content") + if piece: + if first is None: + first = time.perf_counter() - t0 + text.append(piece) + total = time.perf_counter() - t0 + n = usage.get("completion_tokens", 0) + return ( + "".join(text), + first, + (n - 1) / (total - first) if first and n > 1 else 0, + usage, + ) + + +for name, opener in ( + ("short", "你好，请用三句话介绍一下长城。"), + ("long", "请阅读下面的文档：\n\n" + DOC + "\n\n这份文档主要讲了什么？"), +): + messages = [{"role": "user", "content": opener}] + for turn, follow in enumerate( + ("请再简要总结成一句话。", "再列出三个关键词。", None) + ): + text, ttft, tps, usage = chat(messages) + print( + "%s turn%d prompt=%s ttft=%.2fs decode=%.1f tok/s out=%r" + % (name, turn + 1, usage.get("prompt_tokens"), ttft or -1, tps, text[:60]), + flush=True, + ) + if follow is None: + break + messages += [ + {"role": "assistant", "content": text}, + {"role": "user", "content": follow}, + ] diff --git a/benchmarks/deepseek_v41/tools/api_smoke.py b/benchmarks/deepseek_v41/tools/api_smoke.py new file mode 100644 index 0000000000..4605f8b4a4 --- /dev/null +++ b/benchmarks/deepseek_v41/tools/api_smoke.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json +import sys +import urllib.request + +base = sys.argv[1] +key = open("/home/ymzx/ds41-dev/state/api_key.txt").read().strip() + + +def post(path, body, auth=True): + req = urllib.request.Request( + base + path, + json.dumps(body).encode(), + { + "Content-Type": "application/json", + **({"Authorization": "Bearer " + key} if auth else {}), + }, + ) + with urllib.request.urlopen(req, timeout=300) as r: + return json.loads(r.read()) + + +d = post( + "/v1/chat/completions", + { + "model": "deepseek-v4.1-flash", + "messages": [{"role": "user", "content": "用一句话介绍你自己"}], + "max_tokens": 64, + }, +) +print("chat:", d["choices"][0]["message"]["content"], d["usage"]) +tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "查询城市天气", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } +] +d = post( + "/v1/chat/completions", + { + "model": "deepseek-v4.1-flash", + "messages": [{"role": "user", "content": "北京今天天气怎么样？"}], + "tools": tools, + "tool_choice": "auto", + "max_tokens": 128, + }, +) +m = d["choices"][0]["message"] +print( + "tool_calls:", + json.dumps(m.get("tool_calls"), ensure_ascii=False), + "finish:", + d["choices"][0]["finish_reason"], +) +if m.get("tool_calls"): + call = m["tool_calls"][0] + msgs = [ + {"role": "user", "content": "北京今天天气怎么样？"}, + { + "role": "assistant", + "content": m.get("content") or "", + "tool_calls": m["tool_calls"], + }, + { + "role": "tool", + "tool_call_id": call["id"], + "content": json.dumps( + {"city": "北京", "weather": "晴", "temp_c": 18}, ensure_ascii=False + ), + }, + ] + d = post( + "/v1/chat/completions", + { + "model": "deepseek-v4.1-flash", + "messages": msgs, + "tools": tools, + "max_tokens": 96, + }, + ) + print("after tool:", d["choices"][0]["message"]["content"]) +try: + post( + "/v1/chat/completions", + { + "model": "deepseek-v4.1-flash", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 4, + }, + auth=False, + ) + print("no-key: ACCEPTED (bad)") +except urllib.error.HTTPError as e: + print("no-key rejected:", e.code) diff --git a/benchmarks/deepseek_v41/tools/api_soak.py b/benchmarks/deepseek_v41/tools/api_soak.py new file mode 100644 index 0000000000..feb333dc9b --- /dev/null +++ b/benchmarks/deepseek_v41/tools/api_soak.py @@ -0,0 +1,213 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""API soak: streaming chat requests of mixed lengths against the OpenAI server. + +Per request: TTFT, decode tokens/s, output tokens, finish reason, errors. +Mix: mixed-240 questions, long documents (2K/8K/15K prompt tokens), growing +multi-turn chats, and sampled (temperature 0.7) requests. +""" + +import argparse +import json +import random +import time +import urllib.request +from pathlib import Path + +BASE = "http://127.0.0.1:18000" +MODEL = "deepseek-v4.1-flash" + + +def stream(messages, max_tokens, temperature, timeout=1800): + body = dict( + model=MODEL, + messages=messages, + max_tokens=max_tokens, + temperature=temperature, + stream=True, + stream_options={"include_usage": True}, + chat_template_kwargs={"enable_thinking": False}, + ) + req = urllib.request.Request( + BASE + "/v1/chat/completions", + data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}, + ) + start = time.monotonic() + first = None + text, usage, finish, chunks = [], None, None, 0 + with urllib.request.urlopen(req, timeout=timeout) as resp: + for raw in resp: + line = raw.decode().strip() + if not line.startswith("data:"): + continue + data = line[5:].strip() + if data == "[DONE]": + break + obj = json.loads(data) + if obj.get("usage"): + usage = obj["usage"] + for choice in obj.get("choices", []): + delta = choice.get("delta", {}).get("content") + if delta: + if first is None: + first = time.monotonic() + text.append(delta) + chunks += 1 + finish = choice.get("finish_reason") or finish + end = time.monotonic() + out_tokens = usage["completion_tokens"] if usage else None + decode = (end - first) if first else None + return dict( + ttft=round((first or end) - start, 3), + total=round(end - start, 3), + prompt_tokens=usage["prompt_tokens"] if usage else None, + out_tokens=out_tokens, + decode_tok_s=round((out_tokens - 1) / decode, 1) + if decode and out_tokens and out_tokens > 1 + else None, + finish=finish, + chunks=chunks, + text="".join(text), + ) + + +def long_doc(src_root, approx_tokens): + files = sorted(Path(src_root).rglob("*.py")) + random.shuffle(files) + parts, size = [], 0 + for f in files: + try: + t = f.read_text() + except Exception: + continue + parts.append(t) + size += len(t) // 3 # ~3 characters per token for code + if size >= approx_tokens: + break + doc = "".join(parts)[: approx_tokens * 3] + return doc + "\n\nSummarize what the code above does in five bullet points." + + +def main(args): + random.seed(7) + sys_path = Path(__file__).resolve().parent + import sys + + sys.path.insert(0, str(sys_path)) + from qgrade import prompt_of + + items = json.load(open(args.items)) + plan = [] + for i, item in enumerate(items): + plan.append( + ( + "q", + dict( + messages=[{"role": "user", "content": prompt_of(item)}], + max_tokens=512, + temperature=0.0, + ), + ) + ) + if i % 10 == 5: + plan.append( + ( + "sample", + dict( + messages=[{"role": "user", "content": prompt_of(item)}], + max_tokens=256, + temperature=0.7, + ), + ) + ) + if i % 40 == 20: + for n in (2048, 8192, 15000): + plan.append( + ( + f"long{n}", + dict( + messages=[ + {"role": "user", "content": long_doc(args.src_root, n)} + ], + max_tokens=256, + temperature=0.0, + ), + ) + ) + if i % 60 == 30: + plan.append(("multiturn", None)) + if args.sample_stress: + plan = [ + ( + "sample", + dict( + messages=[ + {"role": "user", "content": prompt_of(items[i % len(items)])} + ], + max_tokens=256, + temperature=0.8, + ), + ) + for i in range(args.sample_stress) + ] + out = open(args.output, "w") + for k, (kind, req) in enumerate(plan[: args.limit]): + rec = dict(k=k, kind=kind) + try: + if kind == "multiturn": + msgs, turns = [], [] + for turn in ( + "给我讲讲长城的历史。", + "那它在明朝有什么变化？", + "再总结成三点。", + "用英文翻译你的总结。", + ): + msgs.append({"role": "user", "content": turn}) + r = stream(msgs, 384, 0.0) + msgs.append({"role": "assistant", "content": r["text"]}) + r.pop("text") + turns.append(r) + rec.update( + turns=turns, + ok=all(t["finish"] in ("stop", "length") for t in turns), + ) + else: + r = stream(**req) + rec.update( + {k2: v for k2, v in r.items() if k2 != "text"}, + head=r["text"][:120], + ok=r["finish"] in ("stop", "length") and r["out_tokens"], + ) + except Exception as error: + rec.update(ok=False, error=repr(error)[:500]) + out.write(json.dumps(rec, ensure_ascii=False) + "\n") + out.flush() + print( + json.dumps( + { + x: rec.get(x) + for x in ( + "k", + "kind", + "ok", + "ttft", + "decode_tok_s", + "prompt_tokens", + "out_tokens", + "error", + ) + } + ), + flush=True, + ) + + +if __name__ == "__main__": + p = argparse.ArgumentParser() + p.add_argument("--items", default="/home/ymzx/ds41-dev/evaldata/mixed240.json") + p.add_argument("--src-root", default="/home/ymzx/ds41-dev/src/vllm") + p.add_argument("--output", required=True) + p.add_argument("--limit", type=int, default=10000) + p.add_argument("--sample-stress", type=int, default=0) + main(p.parse_args()) diff --git a/benchmarks/deepseek_v41/tools/api_think.py b/benchmarks/deepseek_v41/tools/api_think.py new file mode 100644 index 0000000000..33696ee2fc --- /dev/null +++ b/benchmarks/deepseek_v41/tools/api_think.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json +import sys +import time +import urllib.request + +base, key = sys.argv[1], sys.argv[2] + + +def post(body): + req = urllib.request.Request( + base + "/v1/chat/completions", + json.dumps(body).encode(), + {"Content-Type": "application/json", "Authorization": "Bearer " + key}, + ) + t = time.time() + with urllib.request.urlopen(req, timeout=600) as r: + d = json.loads(r.read()) + return d, time.time() - t + + +M = "deepseek-v4.1-flash" +d, dt = post( + { + "model": M, + "messages": [ + { + "role": "user", + "content": "一个笼子里有鸡和兔共35只，脚共94只，鸡兔各几只？", + } + ], + "max_tokens": 2048, + } +) +m = d["choices"][0]["message"] +print("[thinking default] %.1fs usage=%s" % (dt, d["usage"])) +print( + " reasoning_content:", + (m.get("reasoning_content") or m.get("reasoning") or "")[:200].replace("\n", " "), +) +print(" content:", (m.get("content") or "")[:200].replace("\n", " ")) +d, dt = post( + { + "model": M, + "messages": [ + { + "role": "user", + "content": "一个笼子里有鸡和兔共35只，脚共94只，鸡兔各几只？", + } + ], + "max_tokens": 512, + "chat_template_kwargs": {"thinking": False}, + } +) +m = d["choices"][0]["message"] +print( + "[thinking off] %.1fs usage=%s reasoning=%r" + % (dt, d["usage"], (m.get("reasoning_content") or "")[:40]) +) +print(" content:", (m.get("content") or "")[:160].replace("\n", " ")) +tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "查询城市天气", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } +] +d, dt = post( + { + "model": M, + "messages": [{"role": "user", "content": "上海明天适合跑步吗？先查一下天气。"}], + "tools": tools, + "max_tokens": 1024, + } +) +m = d["choices"][0]["message"] +print( + "[thinking + tools] %.1fs finish=%s tool_calls=%s" + % ( + dt, + d["choices"][0]["finish_reason"], + json.dumps(m.get("tool_calls"), ensure_ascii=False)[:200], + ) +) +print(" reasoning:", (m.get("reasoning_content") or "")[:120].replace("\n", " ")) diff --git a/benchmarks/deepseek_v41/tools/decode_loop.py b/benchmarks/deepseek_v41/tools/decode_loop.py new file mode 100644 index 0000000000..76b8802ef7 --- /dev/null +++ b/benchmarks/deepseek_v41/tools/decode_loop.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Long sync decode loop for CPU profiling (prints DECODE START, then 4 x 1500 tokens).""" + +import os +import time +from pathlib import Path + +os.environ.setdefault("VLLM_V41_ASYNC", "0") + + +def main(): + from benchmarks.deepseek_v41.main_smoke import create_llm + from benchmarks.deepseek_v41.performance import make_prompt + from vllm import SamplingParams + + llm = create_llm( + Path("/home/ymzx/ds41-models"), + 8192, + Path("/home/ymzx/ds41-dev/dspark-v41"), + 0.33, + kv_cache_mib=256, + decode_graph=True, + draft_graph=True, + custom_allreduce=True, + ) + tok = llm.get_tokenizer() + prompt = make_prompt(tok, 128) + sp = SamplingParams(max_tokens=1500, temperature=0, ignore_eos=True) + llm.generate([{"prompt_token_ids": prompt}], sp, use_tqdm=False) + print("DECODE START", flush=True) + for i in range(4): + t = time.perf_counter() + out = llm.generate([{"prompt_token_ids": prompt}], sp, use_tqdm=False) + n = len(out[0].outputs[0].token_ids) + print( + "rep %d tokens %d %.2f s -> %.2f ms/token" + % (i, n, time.perf_counter() - t, (time.perf_counter() - t) / n * 1000), + flush=True, + ) + print("DECODE END", flush=True) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/deepseek_v41/tools/fresh_prompts.py b/benchmarks/deepseek_v41/tools/fresh_prompts.py new file mode 100644 index 0000000000..42c9316990 --- /dev/null +++ b/benchmarks/deepseek_v41/tools/fresh_prompts.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""API-like first-seen prompt corpus for the fresh-text decode metric (never reused as warm-up).""" + +WARMUP = [ + "你好，今天天气怎么样？", + "Write a haiku about the sea.", + "用一句话解释什么是机器学习。", + "List three fruits that are red.", +] +FRESH = [ + ("zh-history", "简要介绍一下唐朝的科举制度，以及它对后世的影响。"), + ( + "en-science", + "Explain how photosynthesis works, step by step, for a high school student.", + ), + ("zh-advice", "我下个月要去面试一家互联网公司的产品经理岗位，应该怎么准备？"), + ( + "en-email", + "Draft a polite email to a client explaining that the project delivery will be delayed by two weeks.", + ), + ( + "code-js", + "Write a JavaScript function that debounces another function, with comments.", + ), + ( + "code-sql", + "写一个 SQL 查询：找出每个部门工资最高的三名员工，表结构为 employees(id, name, dept_id, salary)。", + ), + ( + "code-rust", + "Implement a thread-safe counter in Rust using Arc and Mutex, and explain the ownership rules involved.", + ), + ( + "code-debug", + "下面的 Python 代码为什么会报错？如何修复？\nnums = [1, 2, 3]\nfor i in range(len(nums) + 1):\n print(nums[i])", + ), + ("math-proof", "证明：任意两个奇数之和是偶数。"), + ( + "math-word", + "A train leaves at 9:00 traveling 80 km/h; another leaves the same station at 10:00 at 100 km/h on the same track. When does the second catch up?", + ), + ("zh-poem", "以“春夜”为题写一首七言绝句，并解释每一句的意境。"), + ( + "en-story", + "Write a short science-fiction story about a robot who learns to paint, about 300 words.", + ), + ("zh-story", "写一个关于外婆家老槐树的温情小故事，约300字。"), + ( + "translate-en", + "Translate into Chinese: The rapid development of renewable energy is reshaping the global economy, creating new industries while challenging traditional ones.", + ), + ( + "translate-zh", + "把这句话翻译成日语和法语：学习一门新语言最重要的是坚持每天练习。", + ), + ( + "json-api", + "Design a JSON schema for a REST API response that returns a paginated list of orders, with an example.", + ), + ( + "yaml-config", + "写一个 docker-compose.yml，包含 nginx、一个 Python Flask 应用和 PostgreSQL 数据库，并加上注释。", + ), + ("zh-econ", "通货膨胀对普通家庭有哪些影响？普通人可以如何应对？"), + ("en-compare", "Compare TCP and UDP in a table and explain when to use each."), + ("zh-health", "长期熬夜对身体有哪些危害？给出改善作息的具体建议。"), + ( + "en-philosophy", + "What is the trolley problem, and what do utilitarian and deontological ethics say about it?", + ), + ("zh-travel", "帮我规划一个北京三日游的行程，包括交通和美食推荐。"), + ( + "en-summary", + "Summarize the main causes and consequences of the 2008 financial crisis in five bullet points.", + ), + ("zh-tech", "解释一下 Transformer 模型中的注意力机制，最好举一个直观的例子。"), +] diff --git a/benchmarks/deepseek_v41/tools/longctx_test.py b/benchmarks/deepseek_v41/tools/longctx_test.py new file mode 100644 index 0000000000..7b54b4722d --- /dev/null +++ b/benchmarks/deepseek_v41/tools/longctx_test.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Long-context check: needle retrieval, TTFT and decode speed at large prompt lengths.""" + +import json +import os +import sys +import time +from pathlib import Path + + +def main(out_path, lengths): + from v41_e2e import long_prompt + + from benchmarks.deepseek_v41.main_smoke import create_llm + from vllm import SamplingParams + + llm = create_llm( + Path("/home/ymzx/ds41-models"), + 8192, + Path("/home/ymzx/ds41-dev/dspark-v41"), + 0.33, + kv_cache_mib=int(os.environ.get("KV_MIB", "512")), + decode_graph=True, + draft_graph=True, + custom_allreduce=True, + ) + tok = llm.get_tokenizer() + engine = llm.llm_engine + report = [] + for n in lengths: + needle = "【重要】本次测试的秘密口令是：翡翠-7381-QX。请记住它。" + question = "\n\n问题：上文中提到的秘密口令是什么？请只回答口令本身。" + body = long_prompt(tok, Path("/home/ymzx/ds41-dev/src"), n) + text = tok.decode(body) + depth = len(text) // 3 + text = text[:depth] + "\n" + needle + "\n" + text[depth:] + prompt = tok.apply_chat_template( + [{"role": "user", "content": text + question}], + tokenize=True, + add_generation_prompt=True, + ) + engine.add_request( + f"long-{n}", + {"prompt_token_ids": prompt}, + SamplingParams(max_tokens=48, temperature=0), + ) + t = time.perf_counter() + first = None + out = None + stamps = [] + while engine.has_unfinished_requests(): + for o in engine.step(): + now = time.perf_counter() + if o.outputs and o.outputs[0].token_ids: + if first is None: + first = now - t + stamps.append((now, len(o.outputs[0].token_ids))) + if o.finished: + out = o.outputs[0] + answer = out.text + decode = ( + (stamps[-1][1] - stamps[0][1]) / (stamps[-1][0] - stamps[0][0]) + if len(stamps) > 1 + else 0 + ) + entry = dict( + prompt_tokens=len(prompt), + ttft=round(first, 2), + decode_tok_s=round(decode, 1), + found="7381" in answer, + answer=answer[:80], + ) + # Follow-up turn: prefix-cache hit at long context + a long decode. + follow = tok.encode( + "\n\n请用中文写一段约 300 字的说明，介绍这份文档的主要内容。", + add_special_tokens=False, + ) + p2 = prompt + list(out.token_ids) + follow + engine.add_request( + f"long2-{n}", + {"prompt_token_ids": p2}, + SamplingParams(max_tokens=300, temperature=0, ignore_eos=True), + ) + t = time.perf_counter() + first2 = None + stamps = [] + while engine.has_unfinished_requests(): + for o in engine.step(): + now = time.perf_counter() + if o.outputs and o.outputs[0].token_ids: + if first2 is None: + first2 = now - t + stamps.append((now, len(o.outputs[0].token_ids))) + if len(stamps) > 2: + entry["turn2_ttft"] = round(first2, 2) + entry["turn2_decode_tok_s"] = round( + (stamps[-1][1] - stamps[0][1]) / (stamps[-1][0] - stamps[0][0]), 1 + ) + entry["turn2_ms_per_round"] = round( + (stamps[-1][0] - stamps[0][0]) * 1000 / max(1, len(stamps) - 1), 1 + ) + report.append(entry) + print("LONG", json.dumps(entry, ensure_ascii=False), flush=True) + Path(out_path).write_text(json.dumps(report, ensure_ascii=False)) + + +if __name__ == "__main__": + main(sys.argv[1], [int(x) for x in sys.argv[2:]]) diff --git a/benchmarks/deepseek_v41/tools/memlog2.sh b/benchmarks/deepseek_v41/tools/memlog2.sh new file mode 100755 index 0000000000..42c7566426 --- /dev/null +++ b/benchmarks/deepseek_v41/tools/memlog2.sh @@ -0,0 +1,13 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# memlog2.sh : host memcg + per-GPU free VRAM every 20 s +out=$1; prog=$2 +while true; do + a=$(docker exec ds41-compute-20261024 sh -c "grep -E \"^(anon|shmem) \" /sys/fs/cgroup/memory.stat | tr \"\n\" \" \"") + avail=$(grep MemAvailable /proc/meminfo | awk '{print $2}') + gfree=$(nvidia-smi --id=0,1,2,3 --query-gpu=memory.free --format=csv,noheader,nounits | tr "\n" " ") + q=$(wc -l < "$prog" 2>/dev/null) + echo "$(date +%H:%M:%S) n=$q avail_mb=$((avail/1024)) $a gpu_free_mib: $gfree" >> "$out" + sleep 20 +done diff --git a/benchmarks/deepseek_v41/tools/prefill_profile.py b/benchmarks/deepseek_v41/tools/prefill_profile.py new file mode 100644 index 0000000000..26a22fbc9e --- /dev/null +++ b/benchmarks/deepseek_v41/tools/prefill_profile.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Rank-0 torch.profiler trace of one long prefill (max_tokens=1) + wall TTFTs.""" + +import argparse +import json +import time +from pathlib import Path + + +def mem_start(worker): + if worker.local_rank != 0: + return + import torch + + torch.cuda.memory._record_memory_history(max_entries=400000, stacks="python") + + +def mem_dump(worker, path): + if worker.local_rank != 0: + return None + import torch + + torch.cuda.memory._dump_snapshot(path) + torch.cuda.memory._record_memory_history(enabled=None) + return "snapshot %s peak allocated %.0f MiB reserved %.0f MiB" % ( + path, + torch.accelerator.max_memory_allocated() / 2**20, + torch.accelerator.max_memory_reserved() / 2**20, + ) + + +def main(args): + import profile_torch + from fresh_prompts import WARMUP + from v41_e2e import chat, long_prompt + + from benchmarks.deepseek_v41.main_smoke import create_llm, trim_allocator + from vllm import SamplingParams + + llm = create_llm( + Path("/home/ymzx/ds41-models"), + args.chunk, + Path("/home/ymzx/ds41-dev/dspark-v41"), + args.offload, + kv_cache_mib=args.kv_mib, + decode_graph=True, + draft_graph=True, + custom_allreduce=True, + ) + tok = llm.get_tokenizer() + engine = llm.llm_engine + llm.collective_rpc(trim_allocator) + + def run(prompt, name): + engine.add_request( + name, + {"prompt_token_ids": prompt}, + SamplingParams(max_tokens=1, temperature=0), + ) + t = time.perf_counter() + while engine.has_unfinished_requests(): + engine.step() + return time.perf_counter() - t + + for i, text in enumerate(WARMUP[:2]): + run(chat(tok, text), f"w{i}") + report = {} + for n in args.lengths: + if n < 1024: + from benchmarks.deepseek_v41.performance import make_prompt + + prompt = make_prompt(tok, n) + else: + prompt = long_prompt(tok, Path("/home/ymzx/ds41-dev/src"), n) + report[n] = round(run(prompt, f"warm-{n}"), 2) + print("untraced TTFT", n, report[n], flush=True) + if args.memsnap: + prompt = long_prompt(tok, Path("/home/ymzx/ds41-dev/src"), args.trace_len) + llm.collective_rpc(mem_start) + run(prompt, "memsnap") + print(llm.collective_rpc(mem_dump, args=(args.memsnap,))[0], flush=True) + return + if not args.trace: + Path(args.output).write_text(json.dumps(dict(untraced=report))) + return + prompt = long_prompt(tok, Path("/home/ymzx/ds41-dev/src"), args.trace_len) + llm.collective_rpc(profile_torch.start, args=(args.trace,)) + ttft = run(prompt, "traced") + llm.collective_rpc(profile_torch.stop) + print("traced TTFT", args.trace_len, round(ttft, 2), flush=True) + Path(args.output).write_text(json.dumps(dict(untraced=report, traced=ttft))) + + +if __name__ == "__main__": + p = argparse.ArgumentParser() + p.add_argument("--chunk", type=int, default=2048) + p.add_argument("--lengths", type=int, nargs="*", default=[128, 2048, 6144]) + p.add_argument("--trace-len", type=int, default=6144) + p.add_argument("--trace", default="") + p.add_argument("--memsnap", default="") + p.add_argument("--output", required=True) + p.add_argument("--kv-mib", type=int, default=384) + p.add_argument("--offload", type=float, default=0.34) + main(p.parse_args()) diff --git a/benchmarks/deepseek_v41/tools/profile_torch.py b/benchmarks/deepseek_v41/tools/profile_torch.py new file mode 100644 index 0000000000..182547dc76 --- /dev/null +++ b/benchmarks/deepseek_v41/tools/profile_torch.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Rank-0 torch.profiler trace of steady DSpark rounds (no nsys, low overhead).""" + +import argparse +import json +import time +from pathlib import Path + + +def start(worker, path): + if worker.local_rank != 0: + return + import torch + + runner = worker.model_runner + originals = worker._tp_originals = [] + + def wrap(owner, name, label): + original = getattr(owner, name) + originals.append((owner, name, original)) + + def traced(*a, **k): + with torch.profiler.record_function(f"V41::{label}"): + return original(*a, **k) + + setattr(owner, name, traced) + + wrap(runner, "execute_model", "execute_model") + wrap(runner, "sample_tokens", "sample_tokens") if hasattr( + runner, "sample_tokens" + ) else None + wrap(runner, "_model_forward", "target_forward") + wrap(runner, "_sample", "target_sampling") + wrap(runner.model, "compute_logits", "target_head") + wrap(runner.model, "commit_speculative_state", "commit") + wrap(runner.drafter, "propose", "draft_propose") if hasattr( + runner.drafter, "propose" + ) else None + wrap(runner.drafter.model, "forward", "draft_forward") + wrap(runner.drafter, "_sample_draft_tokens", "draft_sampling") + worker._tp = torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ] + ) + worker._tp_path = path + worker._tp.__enter__() + + +def stop(worker): + if worker.local_rank != 0: + return None + worker._tp.__exit__(None, None, None) + worker._tp.export_chrome_trace(worker._tp_path) + for owner, name, original in reversed(worker._tp_originals): + setattr(owner, name, original) + return worker._tp_path + + +def main(args): + from benchmarks.deepseek_v41.acceptance import request + from benchmarks.deepseek_v41.main_smoke import create_llm, trim_allocator + from benchmarks.deepseek_v41.performance import make_prompt + + llm = create_llm( + Path("/home/ymzx/ds41-models"), + args.chunk, + Path("/home/ymzx/ds41-dev/dspark-v41"), + args.offload, + kv_cache_mib=384, + decode_graph=True, + draft_graph=True, + custom_allreduce=True, + ) + try: + llm.collective_rpc(trim_allocator) + prompt = make_prompt(llm.get_tokenizer(), 128) + for _ in range(2): + request(llm, prompt, 128) + llm.collective_rpc(start, args=(str(args.output),)) + engine = llm.llm_engine + original = engine.step + walls = [] + + def step(): + t = time.perf_counter_ns() + out = original() + walls.append((time.perf_counter_ns() - t) / 1e6) + return out + + engine.step = step + result = request(llm, prompt, 64) + engine.step = original + llm.collective_rpc(stop) + Path(str(args.output) + ".steps.json").write_text( + json.dumps(dict(walls=walls, result=result)) + ) + finally: + llm.llm_engine.engine_core.shutdown() + + +if __name__ == "__main__": + p = argparse.ArgumentParser() + p.add_argument("--output", type=Path, required=True) + p.add_argument("--chunk", type=int, default=2048) + p.add_argument("--offload", type=float, default=0.34) + main(p.parse_args()) diff --git a/benchmarks/deepseek_v41/tools/real_decode.py b/benchmarks/deepseek_v41/tools/real_decode.py new file mode 100644 index 0000000000..da98f3dd9c --- /dev/null +++ b/benchmarks/deepseek_v41/tools/real_decode.py @@ -0,0 +1,188 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Real-text decode diagnosis: per-step wall time, tokens/step, expert misses/step; optional trace.""" + +import argparse +import gc +import json +import time +from pathlib import Path + + +def stats(worker, reset=False): + from vllm.model_executor.layers.quantization.gguf_shared_runtime import ( + SharedExpertRank, + ) + + ranks = [o for o in gc.get_objects() if type(o) is SharedExpertRank] + out = [] + for r in ranks: + s = r.statistics + if reset: + s.zero_() + else: + out.append(s.cpu().tolist()) + if not reset and ranks and getattr(ranks[0], "miss_hist", None) is not None: + print("MISS_HIST", ranks[0].miss_hist.cpu().tolist(), flush=True) + return out + + +def main(args): + import profile_torch + from v41_e2e import REAL_TEXT as CASES + from v41_e2e import chat + + from benchmarks.deepseek_v41.main_smoke import create_llm, trim_allocator + from benchmarks.deepseek_v41.performance import make_prompt + + llm = create_llm( + Path("/home/ymzx/ds41-models"), + 2048, + Path("/home/ymzx/ds41-dev/dspark-v41"), + 0.34, + kv_cache_mib=384, + decode_graph=True, + draft_graph=True, + custom_allreduce=True, + ) + tok = llm.get_tokenizer() + engine = llm.llm_engine + from vllm import SamplingParams + + report = [] + try: + llm.collective_rpc(trim_allocator) + if args.corpus == "fresh": + from fresh_prompts import FRESH, WARMUP + + for i, text in enumerate(WARMUP): + engine.add_request( + f"warm-{i}", + {"prompt_token_ids": chat(tok, text)}, + SamplingParams(max_tokens=64, temperature=0), + ) + while engine.has_unfinished_requests(): + engine.step() + prompts = [(n, chat(tok, t)) for n, t in FRESH[: args.cases]] + reps = 1 + else: + prompts = [("count-128", make_prompt(tok, 128))] + [ + (n, chat(tok, t)) for n, t in CASES[: args.cases] + ] + reps = 2 + for name, prompt in prompts: + for rep in range(reps): # second pass = experts warmed by the same text + llm.collective_rpc(stats, args=(True,)) + engine.add_request( + f"{name}-{rep}-{time.monotonic_ns()}", + {"prompt_token_ids": prompt}, + SamplingParams( + max_tokens=args.tokens, + temperature=0, + logprobs=0 if args.logprobs else None, + ), + ) + steps = [] + count = 0 + ids = [] + last = None + while engine.has_unfinished_requests(): + t = time.perf_counter() + outs = engine.step() + dt = time.perf_counter() - t + n = ( + sum(len(o.outputs[0].token_ids) for o in outs) + if outs + else count + ) + if outs and outs[0].outputs[0].token_ids: + ids = list(outs[0].outputs[0].token_ids) + last = outs[0].outputs[0] + steps.append((dt, max(0, n - count))) + count = max(count, n) + st = llm.collective_rpc(stats)[0] + req = sum(sum(l[0] for l in s) for s in st[:1]) + miss = sum(sum(l[1] for l in s) for s in st[:1]) + dec = steps[2:] # drop prefill steps + wall = sum(d for d, _ in dec) + toks = sum(k for _, k in dec) + entry = dict( + name=name, + rep=rep, + prompt=len(prompt), + tokens=count, + decode_steps=len(dec), + decode_wall_s=round(wall, 3), + ms_per_step=round(1000 * wall / max(1, len(dec)), 2), + tok_per_s=round(toks / max(wall, 1e-9), 1), + tok_per_step=round(toks / max(1, len(dec)), 2), + layer_requests=req, + layer_misses=miss, + misses_per_step=round(miss / max(1, len(dec)), 2), + ttft_s=round(sum(d for d, _ in steps[:2]), 3), + ids=ids, + ) + if args.logprobs and last is not None and last.logprobs is not None: + entry["lp"] = [ + round(float(e[t].logprob), 4) if e and t in e else None + for t, e in zip(last.token_ids, last.logprobs) + ] + print(json.dumps(entry), flush=True) + report.append(entry) + fresh = [e for e in report if e["rep"] == 0] + steps_total = sum(e["decode_steps"] for e in fresh) + wall_total = sum(e["decode_wall_s"] for e in fresh) + tok_total = sum(e["tok_per_step"] * e["decode_steps"] for e in fresh) + summary = dict( + summary=True, + corpus=args.corpus, + prompts=len(fresh), + ms_per_round=round(1000 * wall_total / max(1, steps_total), 2), + tok_per_s=round(tok_total / max(wall_total, 1e-9), 1), + tok_per_round=round(tok_total / max(1, steps_total), 2), + misses_per_round=round( + sum(e["layer_misses"] for e in fresh) / max(1, steps_total), 2 + ), + ) + print(json.dumps(summary), flush=True) + report.append(summary) + if args.trace: + llm.collective_rpc(profile_torch.start, args=(str(args.trace),)) + if args.corpus == "fresh": + from fresh_prompts import FRESH + + name, text = FRESH[min(args.cases, len(FRESH) - 1)] + prompt = chat(tok, text) + else: + name, prompt = prompts[1] + engine.add_request( + f"trace-{time.monotonic_ns()}", + {"prompt_token_ids": prompt}, + SamplingParams(max_tokens=64, temperature=0), + ) + while engine.has_unfinished_requests(): + engine.step() + llm.collective_rpc(profile_torch.stop) + + def pred(worker): + from vllm.models.deepseek_v41 import model as m + + return m.prediction_report() + + rep = llm.collective_rpc(pred)[0] + if rep: + print(json.dumps(dict(prediction=rep)), flush=True) + Path(args.output).write_text(json.dumps(report, indent=1)) + finally: + engine.engine_core.shutdown() + + +if __name__ == "__main__": + p = argparse.ArgumentParser() + p.add_argument("--output", required=True) + p.add_argument("--trace") + p.add_argument("--cases", type=int, default=4) + p.add_argument("--corpus", default="cases", choices=["cases", "fresh"]) + p.add_argument("--tokens", type=int, default=200) + p.add_argument("--logprobs", action="store_true") + main(p.parse_args()) diff --git a/benchmarks/deepseek_v41/tools/run_decode_spy.sh b/benchmarks/deepseek_v41/tools/run_decode_spy.sh new file mode 100644 index 0000000000..24537db89c --- /dev/null +++ b/benchmarks/deepseek_v41/tools/run_decode_spy.sh @@ -0,0 +1,20 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# run_decode_spy.sh NAME: sync decode loop + py-spy of EngineCore and Worker_TP0 (inside ds41-compute, as root) +set -uo pipefail +dev_root=/home/ymzx/ds41-dev; state=$dev_root/state; name=$1 +cd "$dev_root/src" +export PYTHONPATH="$dev_root/src:$dev_root/tools" XDG_CACHE_HOME="$dev_root/cache" XDG_CONFIG_HOME="$dev_root/cache/config" +export FLASHINFER_WORKSPACE_BASE="$dev_root/cache" TRITON_CACHE_DIR="$dev_root/cache/triton" VLLM_NO_USAGE_STATS=1 VLLM_V41_ASYNC=0 +cleanup() { status=$?; pkill -9 -f "VLLM::Worker" 2>/dev/null; pkill -9 -f "VLLM::EngineCore" 2>/dev/null; printf "%s\n" "$status" > "$state/$name.exit"; } +trap cleanup EXIT +.venv/bin/python "$dev_root/tools/decode_loop.py" > "$state/$name.log" 2>&1 & +main=$! +until grep -q "DECODE START" "$state/$name.log" 2>/dev/null || ! kill -0 $main 2>/dev/null; do sleep 2; done +sleep 3 +ec=$(pgrep -f "VLLM::EngineCore" | head -1); w0=$(pgrep -f "VLLM::Worker_TP0" | head -1) +.venv/bin/py-spy record -p "$ec" -d 12 -r 200 -f raw -o "$state/$name.engine.txt" --nonblocking > "$state/$name.spy1.log" 2>&1 & +.venv/bin/py-spy record -p "$w0" -d 12 -r 200 -f raw -o "$state/$name.worker.txt" --nonblocking > "$state/$name.spy2.log" 2>&1 & +.venv/bin/py-spy record -p "$main" -d 12 -r 200 -f raw -o "$state/$name.main.txt" --nonblocking > "$state/$name.spy3.log" 2>&1 & +wait $main diff --git a/benchmarks/deepseek_v41/tools/run_e2e2.sh b/benchmarks/deepseek_v41/tools/run_e2e2.sh new file mode 100755 index 0000000000..7e2cc271a9 --- /dev/null +++ b/benchmarks/deepseek_v41/tools/run_e2e2.sh @@ -0,0 +1,18 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Usage: run_e2e2.sh [v41_e2e.py args]; inside ds41-compute. Cleans up workers. +set -uo pipefail +dev_root=/home/ymzx/ds41-dev +state=$dev_root/state +name=${1:?name}; shift +cd "$dev_root/src" +export PYTHONPATH="$dev_root/src:$dev_root/tools" +export XDG_CACHE_HOME="$dev_root/cache" XDG_CONFIG_HOME="$dev_root/cache/config" +export FLASHINFER_WORKSPACE_BASE="$dev_root/cache" TRITON_CACHE_DIR="$dev_root/cache/triton" +export VLLM_NO_USAGE_STATS=1 +cleanup() { status=$?; pkill -9 -f "VLLM::Worker" 2>/dev/null; pkill -9 -f "VLLM::EngineCore" 2>/dev/null; printf "%s\n" "$status" > "$state/$name.exit"; } +trap cleanup EXIT +pkill -9 -f "VLLM::Worker" 2>/dev/null; sleep 2 +[[ -z $(nvidia-smi --id=0,1,2,3 --query-compute-apps=pid --format=csv,noheader) ]] || { echo "GPU busy" > "$state/$name.log"; exit 3; } +.venv/bin/python "$dev_root/tools/v41_e2e.py" --label "$name" --output "$state/$name.json" "$@" > "$state/$name.log" 2>&1 diff --git a/benchmarks/deepseek_v41/tools/run_longctx.sh b/benchmarks/deepseek_v41/tools/run_longctx.sh new file mode 100644 index 0000000000..fa6975fbd0 --- /dev/null +++ b/benchmarks/deepseek_v41/tools/run_longctx.sh @@ -0,0 +1,11 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +set -uo pipefail +dev_root=/home/ymzx/ds41-dev; state=$dev_root/state; name=$1; shift +cd "$dev_root/src" +export PYTHONPATH="$dev_root/src:$dev_root/tools" XDG_CACHE_HOME="$dev_root/cache" XDG_CONFIG_HOME="$dev_root/cache/config" +export FLASHINFER_WORKSPACE_BASE="$dev_root/cache" TRITON_CACHE_DIR="$dev_root/cache/triton" VLLM_NO_USAGE_STATS=1 VLLM_V41_ASYNC=${VLLM_V41_ASYNC:-1} +cleanup() { status=$?; pkill -9 -f "VLLM::Worker" 2>/dev/null; pkill -9 -f "VLLM::EngineCore" 2>/dev/null; printf "%s\n" "$status" > "$state/$name.exit"; } +trap cleanup EXIT +.venv/bin/python "$dev_root/tools/longctx_test.py" "$state/$name.json" "$@" > "$state/$name.log" 2>&1 diff --git a/benchmarks/deepseek_v41/tools/run_prefill_profile.sh b/benchmarks/deepseek_v41/tools/run_prefill_profile.sh new file mode 100755 index 0000000000..84296edc4c --- /dev/null +++ b/benchmarks/deepseek_v41/tools/run_prefill_profile.sh @@ -0,0 +1,11 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +set -uo pipefail +dev_root=/home/ymzx/ds41-dev; state=$dev_root/state; name=$1; shift +cd "$dev_root/src" +export PYTHONPATH="$dev_root/src:$dev_root/tools" XDG_CACHE_HOME="$dev_root/cache" XDG_CONFIG_HOME="$dev_root/cache/config" +export FLASHINFER_WORKSPACE_BASE="$dev_root/cache" TRITON_CACHE_DIR="$dev_root/cache/triton" VLLM_NO_USAGE_STATS=1 VLLM_V41_ASYNC=1 +cleanup() { status=$?; pkill -9 -f "VLLM::Worker" 2>/dev/null; pkill -9 -f "VLLM::EngineCore" 2>/dev/null; printf "%s\n" "$status" > "$state/$name.exit"; } +trap cleanup EXIT +.venv/bin/python "$dev_root/tools/prefill_profile.py" --output "$state/$name.json" "$@" > "$state/$name.log" 2>&1 diff --git a/benchmarks/deepseek_v41/tools/run_real.sh b/benchmarks/deepseek_v41/tools/run_real.sh new file mode 100755 index 0000000000..f7316f3bd4 --- /dev/null +++ b/benchmarks/deepseek_v41/tools/run_real.sh @@ -0,0 +1,12 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +set -uo pipefail +dev_root=/home/ymzx/ds41-dev; state=$dev_root/state; name=$1; shift +cd "$dev_root/src" +export PYTHONPATH="$dev_root/src:$dev_root/tools" XDG_CACHE_HOME="$dev_root/cache" XDG_CONFIG_HOME="$dev_root/cache/config" +export FLASHINFER_WORKSPACE_BASE="$dev_root/cache" TRITON_CACHE_DIR="$dev_root/cache/triton" VLLM_NO_USAGE_STATS=1 VLLM_V41_ASYNC=${VLLM_V41_ASYNC:-1} +export VLLM_V41_ROUTE_TRACE=${VLLM_V41_ROUTE_TRACE:-} +cleanup() { status=$?; pkill -9 -f "VLLM::Worker" 2>/dev/null; pkill -9 -f "VLLM::EngineCore" 2>/dev/null; printf "%s\n" "$status" > "$state/$name.exit"; } +trap cleanup EXIT +.venv/bin/python "$dev_root/tools/real_decode.py" --output "$state/$name.json" "$@" > "$state/$name.log" 2>&1 diff --git a/benchmarks/deepseek_v41/tools/run_serve2.sh b/benchmarks/deepseek_v41/tools/run_serve2.sh new file mode 100644 index 0000000000..9f69c71cbd --- /dev/null +++ b/benchmarks/deepseek_v41/tools/run_serve2.sh @@ -0,0 +1,13 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# run_serve.sh ; inside ds41-compute. OpenAI API on 0.0.0.0:18000 (container bridge only). +set -uo pipefail +dev_root=${DS41_ROOT:-/home/ymzx/ds41-dev}; state=$dev_root/state; name=${1:-serve} +cd "$dev_root/src" +export PYTHONPATH="$dev_root/src:$dev_root/tools" XDG_CACHE_HOME="$dev_root/cache" XDG_CONFIG_HOME="$dev_root/cache/config" +export FLASHINFER_WORKSPACE_BASE="$dev_root/cache" TRITON_CACHE_DIR="$dev_root/cache/triton" VLLM_NO_USAGE_STATS=1 +export VLLM_WORKER_MULTIPROC_METHOD=spawn VLLM_V41_ASYNC=1 OMP_NUM_THREADS=1 +export VLLM_V41_DOWN_ROWS=1 VLLM_V41_SESSION_PREFIX=1 +pkill -9 -f "VLLM::Worker" 2>/dev/null; sleep 2 +exec .venv/bin/python -m vllm.entrypoints.openai.api_server --config benchmarks/deepseek_v41/serve-ds41-phase2.yaml > "$state/$name.log" 2>&1 diff --git a/benchmarks/deepseek_v41/tools/run_serve3.sh b/benchmarks/deepseek_v41/tools/run_serve3.sh new file mode 100644 index 0000000000..443df27165 --- /dev/null +++ b/benchmarks/deepseek_v41/tools/run_serve3.sh @@ -0,0 +1,13 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# run_serve.sh ; inside ds41-compute. OpenAI API on 0.0.0.0:18000 (container bridge only). +set -uo pipefail +dev_root=${DS41_ROOT:-/home/ymzx/ds41-dev}; state=$dev_root/state; name=${1:-serve} +cd "$dev_root/src" +export PYTHONPATH="$dev_root/src:$dev_root/tools" XDG_CACHE_HOME="$dev_root/cache" XDG_CONFIG_HOME="$dev_root/cache/config" +export FLASHINFER_WORKSPACE_BASE="$dev_root/cache" TRITON_CACHE_DIR="$dev_root/cache/triton" VLLM_NO_USAGE_STATS=1 +export VLLM_WORKER_MULTIPROC_METHOD=spawn VLLM_V41_ASYNC=1 OMP_NUM_THREADS=1 +export VLLM_V41_DOWN_ROWS=1 VLLM_V41_SESSION_PREFIX=1 VLLM_V41_PACKED_KV=1 VLLM_V41_INDEX_TILE_MB=32 +pkill -9 -f "VLLM::Worker" 2>/dev/null; sleep 2 +exec .venv/bin/python -m vllm.entrypoints.openai.api_server --config benchmarks/deepseek_v41/serve-ds41-phase3.yaml > "$state/$name.log" 2>&1 diff --git a/benchmarks/deepseek_v41/tools/run_session.sh b/benchmarks/deepseek_v41/tools/run_session.sh new file mode 100644 index 0000000000..f841924023 --- /dev/null +++ b/benchmarks/deepseek_v41/tools/run_session.sh @@ -0,0 +1,11 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +set -uo pipefail +dev_root=/home/ymzx/ds41-dev; state=$dev_root/state; name=$1 +cd "$dev_root/src" +export PYTHONPATH="$dev_root/src:$dev_root/tools" XDG_CACHE_HOME="$dev_root/cache" XDG_CONFIG_HOME="$dev_root/cache/config" +export FLASHINFER_WORKSPACE_BASE="$dev_root/cache" TRITON_CACHE_DIR="$dev_root/cache/triton" VLLM_NO_USAGE_STATS=1 VLLM_V41_ASYNC=${VLLM_V41_ASYNC:-1} +cleanup() { status=$?; pkill -9 -f "VLLM::Worker" 2>/dev/null; pkill -9 -f "VLLM::EngineCore" 2>/dev/null; printf "%s\n" "$status" > "$state/$name.exit"; } +trap cleanup EXIT +.venv/bin/python "$dev_root/tools/session_test.py" "$state/$name.json" > "$state/$name.log" 2>&1 diff --git a/benchmarks/deepseek_v41/tools/session_test.py b/benchmarks/deepseek_v41/tools/session_test.py new file mode 100644 index 0000000000..8e2d37d805 --- /dev/null +++ b/benchmarks/deepseek_v41/tools/session_test.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""G7 session test: turn 2 continues turn 1 (raw token concat); TTFT and output tokens.""" + +import json +import os +import sys +import time +from pathlib import Path + + +def main(out_path): + from v41_e2e import chat, long_prompt + + from benchmarks.deepseek_v41.main_smoke import create_llm + from vllm import SamplingParams + + llm = create_llm( + Path("/home/ymzx/ds41-models"), + 8192, + Path("/home/ymzx/ds41-dev/dspark-v41"), + 0.33, + kv_cache_mib=288, + decode_graph=True, + draft_graph=True, + custom_allreduce=True, + ) + tok = llm.get_tokenizer() + engine = llm.llm_engine + + def run(prompt, name, n): + engine.add_request( + name, + {"prompt_token_ids": prompt}, + SamplingParams(max_tokens=n, temperature=0), + ) + t = time.perf_counter() + first = None + ids = [] + while engine.has_unfinished_requests(): + for o in engine.step(): + if o.outputs and o.outputs[0].token_ids and first is None: + first = time.perf_counter() - t + if o.finished: + ids = list(o.outputs[0].token_ids) + return first, ids + + report = {"prefix": os.environ.get("VLLM_V41_SESSION_PREFIX", "0")} + doc = long_prompt(tok, Path("/home/ymzx/ds41-dev/src"), 6000) + for name, p1 in (("short", chat(tok, "请用三句话介绍一下长城。")), ("long", doc)): + t1, o1 = run(p1, name + "-1", 120) + follow = tok.encode( + "\n\n用户：请把上面的内容再简要总结成一句话。\n\n助手：", + add_special_tokens=False, + ) + p2 = p1 + o1 + follow + t2, o2 = run(p2, name + "-2", 64) + t3, o3 = run(p2 + o2 + follow, name + "-3", 32) + report[name] = dict( + turn1=[len(p1), round(t1, 2)], + turn2=[len(p2), round(t2, 2)], + turn3=[len(p2 + o2 + follow), round(t3, 2)], + out2=o2, + out3=o3, + ) + print( + name, + report[name]["turn1"], + report[name]["turn2"], + report[name]["turn3"], + flush=True, + ) + Path(out_path).write_text(json.dumps(report)) + + +if __name__ == "__main__": + main(sys.argv[1]) diff --git a/benchmarks/deepseek_v41/tools/spy_decode.py b/benchmarks/deepseek_v41/tools/spy_decode.py new file mode 100644 index 0000000000..451b761186 --- /dev/null +++ b/benchmarks/deepseek_v41/tools/spy_decode.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""py-spy raw: decode-only CPU (drop prefill stacks and GPU-wait leaves).""" + +import collections +import sys + +import regex as re + +inc = collections.Counter() +total = 0 +waits = 0 +pre = 0 +WAIT = ( + "parse_output", + "synchronize", + "_to_list", + "item", + "tolist", + "acquire_read", + "sched_yield", + "wait", +) +for line in open(sys.argv[1]): + stack, _, n = line.rstrip().rpartition(" ") + if not stack: + continue + n = int(n) + frames = [ + re.sub(r" \(.*?/([^/]+\.py):(\d+)\)", r" (\1:\2)", f) for f in stack.split(";") + ] + if any("prefill_moe" in f or "_ring_pump" in f for f in frames): + pre += n + continue + leaf = frames[-1].split(" ")[0] + if leaf in WAIT: + waits += n + continue + total += n + seen = set() + for f in frames: + k = re.sub(r":\d+\)", ")", f) + if k not in seen: + inc[k] += n + seen.add(k) +print("decode cpu samples", total, "waits", waits, "prefill", pre) +for f, n in inc.most_common(45): + print("%6d %s" % (n, f[:140])) diff --git a/benchmarks/deepseek_v41/tools/spy_top.py b/benchmarks/deepseek_v41/tools/spy_top.py new file mode 100644 index 0000000000..15e2bb968b --- /dev/null +++ b/benchmarks/deepseek_v41/tools/spy_top.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Summarize py-spy raw (collapsed) output: inclusive and self share per function.""" + +import collections +import sys + +import regex as re + +inc, own = collections.Counter(), collections.Counter() +total = 0 +for line in open(sys.argv[1]): + line = line.rstrip() + if not line: + continue + stack, _, n = line.rpartition(" ") + n = int(n) + total += n + frames = [ + re.sub(r" \(.*?/([^/]+\.py):\d+\)", r" (\1)", f) for f in stack.split(";") + ] + seen = set() + for f in frames: + if f not in seen: + inc[f] += n + seen.add(f) + own[frames[-1]] += n +print("total samples", total) +print("-- self") +for f, n in own.most_common(int(sys.argv[2]) if len(sys.argv) > 2 else 15): + print("%5.1f%% %s" % (100 * n / total, f[:150])) +print("-- inclusive") +for f, n in inc.most_common(int(sys.argv[3]) if len(sys.argv) > 3 else 30): + print("%5.1f%% %s" % (100 * n / total, f[:150])) diff --git a/benchmarks/deepseek_v41/tools/tcp_forward.py b/benchmarks/deepseek_v41/tools/tcp_forward.py new file mode 100644 index 0000000000..87487ceb4a --- /dev/null +++ b/benchmarks/deepseek_v41/tools/tcp_forward.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Minimal TCP forwarder: 0.0.0.0:LISTEN -> TARGET_HOST:TARGET_PORT (API exposure).""" + +import asyncio +import sys + +LISTEN, HOST, PORT = int(sys.argv[1]), sys.argv[2], int(sys.argv[3]) + + +async def pipe(reader, writer): + try: + while data := await reader.read(65536): + writer.write(data) + await writer.drain() + except Exception: + pass + finally: + writer.close() + + +async def handle(client_reader, client_writer): + try: + server_reader, server_writer = await asyncio.open_connection(HOST, PORT) + except Exception: + client_writer.close() + return + await asyncio.gather( + pipe(client_reader, server_writer), pipe(server_reader, client_writer) + ) + + +async def main(): + server = await asyncio.start_server( + handle, sys.argv[4] if len(sys.argv) > 4 else "127.0.0.1", LISTEN + ) + async with server: + await server.serve_forever() + + +asyncio.run(main()) diff --git a/benchmarks/deepseek_v41/tools/top_decode.py b/benchmarks/deepseek_v41/tools/top_decode.py new file mode 100644 index 0000000000..eed2e73831 --- /dev/null +++ b/benchmarks/deepseek_v41/tools/top_decode.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json +import sys +from collections import defaultdict + +import regex as re + +ev = json.load(open(sys.argv[1]))["traceEvents"] +kern = sorted( + (e for e in ev if e.get("ph") == "X" and e.get("cat") == "kernel"), + key=lambda e: e["ts"], +) +ts = [k["ts"] for k in kern] +import bisect + +r = sorted( + e["ts"] + for e in ev + if e.get("ph") == "X" + and e.get("cat") == "user_annotation" + and e["name"] == "V41::target_forward" +) +tot = defaultdict(float) +cnt = defaultdict(int) +n = 0 +spans = [] +for a, b in zip(r, r[1:]): + seg = kern[bisect.bisect_left(ts, a) : bisect.bisect_left(ts, b)] + if ( + not seg + or any(re.search("transcode|prefill|_partition_tile", k["name"]) for k in seg) + or b - a > 60000 + ): + continue + n += 1 + spans.append(b - a) + for k in seg: + nm = re.sub(r"<.*", "<>", re.sub(r"\(.*", "", k["name"]))[:90] + tot[nm] += k["dur"] + cnt[nm] += 1 +spans.sort() +print( + "decode rounds", + n, + "median span us", + spans[len(spans) // 2] if spans else None, + "busy ms", + round(sum(tot.values()) / n / 1e3, 2), + "kernels", + sum(cnt.values()) / n, +) +for k, v in sorted(tot.items(), key=lambda kv: -kv[1])[ + : int(sys.argv[2]) if len(sys.argv) > 2 else 40 +]: + print(f"{v / n / 1e3:7.3f} ms {cnt[k] / n:7.1f}x {k}") diff --git a/benchmarks/deepseek_v41/tools/v41_e2e.py b/benchmarks/deepseek_v41/tools/v41_e2e.py new file mode 100644 index 0000000000..32a7cfa83a --- /dev/null +++ b/benchmarks/deepseek_v41/tools/v41_e2e.py @@ -0,0 +1,335 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""One-load V4.1 batch acceptance: round cost, stage timing, real text, quality. + +Order matters: unhooked performance first, then hooked stage timing, then the +real-text set (unhooked, acceptance from arrival sizes), then traced quality +and state checks. Precision is judged against a saved baseline by +v41_compare.py; this script only records. +""" + +import argparse +import json +import time +from pathlib import Path + +REAL_TEXT = [ + ("chat-zh", "请用三段话介绍一下长城的历史，以及它在今天的意义。"), + ("chat-en", "Explain in a few paragraphs how vaccines train the immune system."), + ( + "advice-zh", + "我是刚毕业的程序员，想在三年内成长为后端架构师，请给我一份具体的学习计划。", + ), + ( + "code-py", + "用Python实现一个线程安全的LRU缓存类，支持get、put和容量上限，并附上简单的单元测试。", + ), + ( + "code-explain", + "Explain what the following does and point out any bugs:\n\n" + "def merge(a, b):\n i = j = 0\n out = []\n while i < len(a) and j < len(b):\n" + " if a[i] < b[j]:\n out.append(a[i]); i += 1\n else:\n" + " out.append(b[j]); j += 1\n return out\n", + ), + ( + "math-zh", + "一个水池有两个进水管，甲管单独注满需要6小时，乙管单独注满需要9小时。" + "两管同时打开，多少小时能注满？请写出详细的推理过程。", + ), + ( + "math-en", + "Prove that the square root of 2 is irrational. Be rigorous but concise.", + ), + ("story-zh", "写一个关于一只在城市里迷路的小猫找到回家路的短篇故事，约400字。"), + ( + "translate", + "把下面这段话翻译成英文：人工智能正在改变我们工作和生活的方式，" + "但我们也需要认真思考它带来的伦理和社会问题。", + ), + ( + "json-zh", + "为一个在线书店设计商品数据的JSON示例，包含书名、作者、价格、库存、" + "分类和评分字段，给出三本书的数据。", + ), +] +LONG_SOURCES = ["README.md", "docs/design/architecture/runtime_reference.md"] + + +def chat(tokenizer, text): + return tokenizer.apply_chat_template( + [{"role": "user", "content": text}], + tokenize=True, + add_generation_prompt=True, + enable_thinking=False, + ) + + +def long_prompt(tokenizer, src_root, length): + material = "\n\n".join( + (src_root / name).read_text(errors="ignore") for name in LONG_SOURCES + ) + head = tokenizer.encode("下面是一份技术文档：\n\n", add_special_tokens=False) + tail_text = "\n\n请用中文分点总结这份文档的主要内容和关键设计。" + tail = chat(tokenizer, "__X__") + marker = tokenizer.decode(tail) + before_text, after_text = marker.split("__X__") + before = tokenizer.encode(before_text, add_special_tokens=False) + after = tokenizer.encode(tail_text + after_text, add_special_tokens=False) + body = tokenizer.encode(material, add_special_tokens=False) + size = length - len(before) - len(head) - len(after) + while len(body) < size: + body = body + body + return before + head + body[:size] + after + + +def run_request(llm, prompt, maximum, logprobs): + from vllm import SamplingParams + + engine = llm.llm_engine + begin = time.monotonic() + params = SamplingParams( + max_tokens=maximum, temperature=0, seed=41, logprobs=0 if logprobs else None + ) + engine.add_request( + f"e2e-{time.monotonic_ns()}", {"prompt_token_ids": prompt}, params + ) + arrivals, previous, final = [], 0, None + while engine.has_unfinished_requests(): + for output in engine.step(): + count = len(output.outputs[0].token_ids) + if count > previous: + arrivals.append( + dict(seconds=time.monotonic() - begin, tokens=count - previous) + ) + previous = count + if output.finished: + final = output + completion = final.outputs[0] + lps = None + if logprobs and completion.logprobs is not None: + lps = [ + float(entry[tok].logprob) if entry and tok in entry else None + for tok, entry in zip(completion.token_ids, completion.logprobs) + ] + first, last = arrivals[0]["seconds"], arrivals[-1]["seconds"] + return dict( + prompt_tokens=len(prompt), + token_ids=list(completion.token_ids), + text=completion.text, + finish_reason=completion.finish_reason, + ttft_seconds=first, + tpot_seconds=(last - first) / (previous - 1) if previous > 1 else None, + arrivals=arrivals, + logprobs=lps, + ) + + +def _validate(validate_commits, ranks, expected_end): + """Per-commit invariants always; under async scheduling the request's + final in-flight speculative step may commit up to one round past the end.""" + import os + + if os.environ.get("VLLM_V41_ASYNC", "0") != "1": + return validate_commits(ranks, expected_end) + validate_commits(ranks, None) + for rank in ranks: + if not rank["commits"]: + raise AssertionError("Request has no state commits") + end = rank["commits"][-1]["hash_history"][3] + if not expected_end <= end <= expected_end + 6: + raise AssertionError((rank["rank"], expected_end, end)) + + +def acceptance(arrivals): + # Arrival 0 is the prefill token; the final arrival may be truncated. + sizes = [a["tokens"] for a in arrivals[1:-1]] + if not sizes: + return None + hits = [sum(1 for s in sizes if s - 1 >= i) / len(sizes) for i in range(1, 6)] + return dict( + rounds=len(sizes), + mean_accepted=sum(s - 1 for s in sizes) / len(sizes), + position_rates=hits, + ) + + +def main(args): + import profile_decode as pd + + from benchmarks.deepseek_v41 import trace + from benchmarks.deepseek_v41.acceptance import check_quality, validate_commits + from benchmarks.deepseek_v41.main_smoke import ( + _memory_report, + create_llm, + trim_allocator, + ) + from benchmarks.deepseek_v41.performance import make_prompt + from benchmarks.deepseek_v41.resources import ResourceMonitor + + monitor = ResourceMonitor() + report = dict( + status="running", + label=args.label, + perf={}, + stages=None, + real_text=[], + quality=[], + errors=[], + ) + + def save(): + report["resources"] = monitor.report() + tmp = args.output.with_suffix(".tmp") + tmp.write_text(json.dumps(report, ensure_ascii=False, indent=1) + "\n") + tmp.replace(args.output) + + llm = None + try: + begin = time.monotonic() + llm = create_llm( + args.model_root, + args.chunk, + args.draft, + args.offload, + kv_cache_mib=args.kv_mib, + decode_graph=True, + draft_graph=True, + custom_allreduce=True, + ) + report["load_seconds"] = time.monotonic() - begin + llm.collective_rpc(trim_allocator) + tok = llm.get_tokenizer() + save() + + # 1. Round cost on the counting load (unhooked). + for length in args.lengths: + prompt = make_prompt(tok, length) + for _ in range(2): + run_request(llm, prompt, 128, False) + report["perf"][str(length)] = [ + run_request(llm, prompt, 128, False) for _ in range(3) + ] + save() + + # 2. Per-stage GPU timing (hooked), 128 input. + prompt = make_prompt(tok, args.lengths[0]) + llm.collective_rpc(pd.install) + stage_runs = [] + for _ in range(2): + run_request(llm, prompt, 128, False) + stage_runs.append(llm.collective_rpc(pd.drain)) + llm.collective_rpc(pd.uninstall) + report["stages"] = stage_runs + save() + + # 3. Real text: acceptance from arrival sizes, logprobs for precision. + quick_names = ( + "chat-zh", + "chat-en", + "code-py", + "math-zh", + "story-zh", + "json-zh", + ) + cases = [ + (n, chat(tok, q)) + for n, q in REAL_TEXT + if not args.quick or n in quick_names + ] + for length in args.long_lengths: + cases.append((f"longdoc-{length}", long_prompt(tok, args.src_root, length))) + for name, prompt in cases: + try: + result = run_request(llm, prompt, args.real_tokens, True) + except Exception as error: # logprobs unsupported under spec decode + report["errors"].append(f"{name}: logprobs failed: {error!r}") + result = run_request(llm, prompt, args.real_tokens, False) + result["name"] = name + result["acceptance"] = acceptance(result["arrivals"]) + report["real_text"].append(result) + save() + + # 4. Quality and committed-state checks (traced). + llm.collective_rpc(trace.install) + prompts = { + "math": "只输出计算结果：17乘以23等于多少？", + "chat": "用一句话解释为什么月亮会有阴晴圆缺。", + "code": "只输出Python代码，写一个返回两个整数之和的add函数。", + "json": "只输出有效JSON，内容是城市北京、温度20，键为city和temperature。", + } + for name, question in prompts.items(): + prompt = chat(tok, question) + result = run_request(llm, prompt, 96, False) + ranks = llm.collective_rpc(trace.drain) + entry = dict(name=name, text=result["text"], ok=True) + try: + _validate( + validate_commits, ranks, len(prompt) + len(result["token_ids"]) - 1 + ) + check_quality(name, result["text"]) + except AssertionError as error: + entry.update(ok=False, error=repr(error)[:2000]) + report["quality"].append(entry) + save() + if args.context_16k: + source = tok.encode( + "这是一段用于验证模型的测试上下文。" * 2048, add_special_tokens=False + ) + question = tok.encode( + "\n请用中文简短回答：2+2等于几？\n答：", add_special_tokens=False + ) + prompt = source[: 16384 - 32 - len(question)] + question + result = run_request(llm, prompt, 32, False) + ranks = llm.collective_rpc(trace.drain) + entry = dict( + name="context-16k", + text=result["text"], + ok=True, + ttft_seconds=result["ttft_seconds"], + ) + try: + _validate( + validate_commits, ranks, len(prompt) + len(result["token_ids"]) - 1 + ) + except AssertionError as error: + entry.update(ok=False, error=repr(error)[:2000]) + report["quality"].append(entry) + save() + report["memory_by_rank"] = llm.collective_rpc(_memory_report) + report["status"] = "completed" + save() + except BaseException as error: + report.update(status="failed", error=repr(error)) + save() + raise + finally: + if llm is not None: + llm.llm_engine.engine_core.shutdown() + monitor.close() + + +if __name__ == "__main__": + p = argparse.ArgumentParser() + p.add_argument("--model-root", type=Path, default=Path("/home/ymzx/ds41-models")) + p.add_argument("--draft", type=Path, default=Path("/home/ymzx/ds41-dev/dspark-v41")) + p.add_argument("--src-root", type=Path, default=Path("/home/ymzx/ds41-dev/src")) + p.add_argument("--lengths", type=int, nargs="+", default=[128, 2048]) + p.add_argument("--long-lengths", type=int, nargs="*", default=[2048, 8192]) + p.add_argument("--real-tokens", type=int, default=256) + p.add_argument("--context-16k", action="store_true") + p.add_argument("--offload", type=float, default=0.30) + p.add_argument("--chunk", type=int, default=512) + p.add_argument("--kv-mib", type=int, default=256) + p.add_argument( + "--quick", + action="store_true", + help="dev mode: 128-token perf, 6 short real-text cases, no long inputs", + ) + p.add_argument("--label", default="") + p.add_argument("--output", type=Path, required=True) + parsed = p.parse_args() + if parsed.quick: + parsed.lengths = [128] + parsed.long_lengths = [] + parsed.context_16k = False + main(parsed) diff --git a/benchmarks/deepseek_v41/trace.py b/benchmarks/deepseek_v41/trace.py new file mode 100644 index 0000000000..d1befb517f --- /dev/null +++ b/benchmarks/deepseek_v41/trace.py @@ -0,0 +1,246 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Opt-in qualification hooks; never installed by the serving model. + +CUDA events measure each rank's stream intervals, including dependency waits. +Nested expert/Engram intervals are part of target time, not additive stages. +No event is synchronized until the caller drains a completed request. +""" + +import time +from collections import defaultdict + +import torch + + +def install(worker): + runner = worker.model_runner + if hasattr(worker, "_v41_trace"): + raise RuntimeError("V4.1 qualification hooks are already installed") + trace = worker._v41_trace = dict( + events=[], commits=[], target_forwards=[], sampling=[] + ) + torch.accelerator.synchronize(worker.device) + + def wrap(owner, name, stage): + original = getattr(owner, name) + + def measured(*args, **kwargs): + if torch.cuda.is_current_stream_capturing(): + # Lazy draft/context/head graphs can capture this hook. + # Internal events express graph dependencies and do not + # have standalone timestamps for elapsed_time on drain. + # The enclosing graph call is timed outside capture. + return original(*args, **kwargs) + if name == "_model_forward": + from vllm.forward_context import get_forward_context + + ids = kwargs.get("input_ids") + if ids is None: + ids = args[0] + trace["target_forwards"].append( + dict( + rows=ids.numel(), + mode=get_forward_context().cudagraph_runtime_mode.name, + ) + ) + begin = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + begin.record() + cpu_start = time.perf_counter() + result = original(*args, **kwargs) + if name == "_sample": + metadata = kwargs.get("spec_decode_metadata") + if metadata is None and len(args) >= 3: + metadata = args[2] + if metadata is not None: + # State commit clips this same tensor in-place for EOS + # and the output budget. Preserve the verifier's actual + # result before clipping so final request tails do not + # masquerade as draft rejection. + remember_sampling(trace, metadata.num_draft_tokens, result) + cpu_ms = (time.perf_counter() - cpu_start) * 1000 + end.record() + trace["events"].append((stage, begin, end, cpu_ms)) + return result + + setattr(owner, name, measured) + + # Instrument above CUDAGraphWrapper: the model's Python forward is not + # invoked by replay. The same boundary also covers ordinary forwards. + wrap(runner, "_model_forward", "target") + wrap(runner.model, "compute_logits", "target_head") + wrap(runner, "_sample", "target_sampling") + target = runner.model.model + cache = target.layers[0].ffn.experts.quant_method.expert_cache + cache.runtime.statistics.zero_() + cache.tm.cold_transcodes.zero_() + target.engram_io.runtime.statistics.zero_() + wrap(cache, "resolve", "expert_transfer_and_group_wait") + wrap(cache.tm, "prepare", "prefill_expert_transfer_and_transcode") + wrap(target.engram_io.runtime, "lookup", "engram_wait_and_dequant") + original_commit = runner.model.commit_speculative_state + + def commit(counts): + original_commit(counts) + trace["commits"].append( + ( + counts.clone(), + target.pending_position.clone(), + target.pending_length.clone(), + target.pending_prefill.clone(), + target.hash_state.history.clone(), + target.hash_state.metadata.clone(), + [ + target.layers[i].attn.compressor.state.position.clone() + for i in (2, 8, 14) + ], + ) + ) + + runner.model.commit_speculative_state = commit + wrap(runner.model, "commit_speculative_state", "commit") + drafter = getattr(runner, "drafter", None) + if drafter is not None: + wrap(drafter.model, "forward", "draft") + wrap(drafter.model, "compute_logits", "draft_head") + wrap(drafter.model, "combine_hidden_states", "draft_main_projection") + wrap(drafter.model, "precompute_and_store_context_kv", "draft_context") + wrap(drafter, "_sample_draft_tokens", "draft_markov_sampling") + torch.accelerator.synchronize(worker.device) + + +def drain(worker): + from vllm.distributed import get_tp_group + + torch.accelerator.synchronize(worker.device) + trace = worker._v41_trace + stages = defaultdict(lambda: dict(calls=0, gpu_ms=0.0, cpu_enqueue_ms=0.0)) + for stage, begin, end, cpu_ms in trace["events"]: + record = stages[stage] + record["calls"] += 1 + record["gpu_ms"] += begin.elapsed_time(end) + record["cpu_enqueue_ms"] += cpu_ms + commits = [] + for counts, start, length, prefill, history, metadata, c2 in trace["commits"]: + commits.append( + dict( + accepted_inputs=int(counts.item()), + start=int(start.item()), + rows=int(length.item()), + prefill=bool(prefill.item()), + hash_history=history.tolist(), + hash_metadata=metadata.tolist(), + c2_positions=[int(t.item()) for t in c2], + ) + ) + trace["events"].clear() + trace["commits"].clear() + sampling = [] + for proposed, tensor in trace["sampling"]: + for count, ids in zip(proposed, tensor.tolist(), strict=True): + prefix = next((i for i, value in enumerate(ids) if value < 0), len(ids)) + sampling.append( + dict( + proposed=count, + accepted_drafts=max(0, min(count, prefix - 1)), + raw_sampled_token_ids=ids, + ) + ) + trace["sampling"].clear() + forwards = trace["target_forwards"][:] + trace["target_forwards"].clear() + drafter = getattr(worker.model_runner, "drafter", None) + graph = getattr(getattr(drafter, "model", None), "query_graph", None) + greedy = getattr(getattr(drafter, "model", None), "greedy_graph", None) + contexts = getattr(getattr(drafter, "model", None), "context_graphs", None) + communicator = get_tp_group().device_communicator + custom = getattr(communicator, "ca_comm", None) + return dict( + rank=worker.local_rank, + stages=dict(stages), + commits=commits, + sampling_before_stop_clipping=sampling, + custom_allreduce_enabled=bool( + communicator.use_custom_allreduce + and custom is not None + and not custom.disabled + ), + target_forwards=forwards, + draft_graph=( + dict(captures=graph.captures, replays=graph.replays) + if graph is not None + else None + ), + greedy_graph=( + dict(captures=greedy.captures, replays=greedy.replays) + if greedy is not None + else None + ), + context_graphs=( + dict( + captures=contexts.captures, + replays=contexts.replays, + rows=sorted(contexts.entries), + ) + if contexts is not None + else None + ), + ) + + +def remember_sampling(trace, proposed, result): + trace["sampling"].append((list(proposed), result.sampled_token_ids.clone())) + + +def set_draft_cap(worker, cap): + runner = worker.model_runner + drafter = getattr(runner, "drafter", None) + if drafter is None: + if cap is not None: + raise ValueError("A draft cap requires DSpark") + return + if cap is not None and cap not in range(6): + raise ValueError("V4.1 draft cap must be in 0..5") + enabled = cap is not None and cap < 5 + # The runner normally initializes these from the launch configuration. + # Changing only the proposer leaves its returned lengths unused, so a + # claimed M1 comparison would still verify the complete M6 proposal. + runner.dspark_confidence_scheduling = enabled + if enabled and runner._dspark_verification_lengths_cpu is None: + runner._dspark_verification_lengths_cpu = torch.empty( + runner.max_num_reqs, + dtype=torch.int32, + device="cpu", + pin_memory=runner.pin_memory, + ) + drafter._max_verification_tokens = 5 if cap is None else cap + drafter.confidence_scheduling_enabled = enabled + + +def force_draft(worker, continuation, reject_at): + """Keep the real draft forward, replace only greedy proposals for testing. + + A previously recorded M1 continuation supplies a known accepted prefix. + The deliberate mismatch at reject_at exercises all six selectors against + the actual target M6 and runner. This is not a quality or speed workload. + """ + drafter = worker.model_runner.drafter + if not hasattr(drafter, "_v41_original_sample"): + drafter._v41_original_sample = drafter._sample_draft_tokens + original = drafter._v41_original_sample + if continuation is None: + drafter._sample_draft_tokens = original + return + if reject_at not in range(6) or len(continuation) < 6: + raise ValueError("Forced draft requires six continuation tokens and 0..5") + values = torch.tensor(continuation[1:6], device=worker.device, dtype=torch.int64) + if reject_at < 5: + values[reject_at] = (values[reject_at] + 1) % 129280 + + def sample(*args, **kwargs): + # Execute native Markov sampling too, then override its proposal IDs. + original(*args, **kwargs) + return values.clone(), None + + drafter._sample_draft_tokens = sample diff --git a/benchmarks/deepseek_v41/v41k/.gitignore b/benchmarks/deepseek_v41/v41k/.gitignore new file mode 100644 index 0000000000..567609b123 --- /dev/null +++ b/benchmarks/deepseek_v41/v41k/.gitignore @@ -0,0 +1 @@ +build/ diff --git a/benchmarks/deepseek_v41/v41k/build.py b/benchmarks/deepseek_v41/v41k/build.py new file mode 100644 index 0000000000..51cdc313e3 --- /dev/null +++ b/benchmarks/deepseek_v41/v41k/build.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Build the V4.1 v41k CUDA extension (SM70) into ./build/v41k.so. + +The runtime loads it from VLLM_V41K_LIB, defaulting to this build directory. +""" + +import os + +from torch.utils import cpp_extension as ce + +here = os.path.dirname(os.path.abspath(__file__)) +src = os.path.abspath(os.path.join(here, "..", "..", "..")) +ce.load( + name="v41k", + sources=[os.path.join(here, "v41k.cu")], + extra_include_paths=[src + "/csrc", src + "/csrc/sm70_turbomind/lmdeploy"], + extra_cuda_cflags=[ + "-O3", + "-gencode", + "arch=compute_70,code=sm_70", + "--expt-relaxed-constexpr", + "-std=c++20", + "-DENABLE_SM70_TURBOMIND=1", + "-U__CUDA_NO_HALF_OPERATORS__", + "-U__CUDA_NO_HALF_CONVERSIONS__", + "-U__CUDA_NO_HALF2_OPERATORS__", + "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", + ], + extra_cflags=["-O3", "-std=c++20"], + is_python_module=False, + build_directory=os.path.join(here, "build"), + verbose=False, +) +print("built", os.path.join(here, "build", "v41k.so")) diff --git a/benchmarks/deepseek_v41/v41k/v41k.cu b/benchmarks/deepseek_v41/v41k/v41k.cu new file mode 100644 index 0000000000..63810d69d2 --- /dev/null +++ b/benchmarks/deepseek_v41/v41k/v41k.cu @@ -0,0 +1,3029 @@ +#include +#include +// V4.1 routed experts, expert-major (each distinct expert's weights are read +// and decoded once per layer for all tokens that selected it). Original GGUF +// IQ1_M / IQ2_XXS / IQ2_XS blocks, Q8_1 activations, dp4a integer dots. +// Codebooks live in a pre-transformed global table (L1-cached) instead of a +// per-CTA shared copy, so small CTAs pay no book initialization. +#include +#include +#include +#include +#include + +#include "sm70_turbomind/ops/gguf_dp4a.cuh" + +namespace v41k { +using vllm::sm70_gguf::Q8_1; + +template +constexpr int block_bytes() { + return Type == 16 ? 66 : Type == 17 ? 74 : 56; +} + +template +__global__ void build_table(uint32_t* out, int words) { + using Book = turbomind::gemm::LatticeCodebook; + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= words) return; + uint32_t word = Book::word(i) ^ 0x80808080U; + if constexpr (Type == 29) { + word = __vadd4(word, word); + word = __vadd4(word, word); + word = __vadd4(word, word); + } + out[i] = word; +} + +__device__ __forceinline__ uint32_t sign_mask(int s) { + return vllm::sm70_gguf::LatticeDot<22>::sign_mask(s); +} + +// One 32-weight group of an original block, decoded to 8 signed int8x4 words. +template +struct Group { + uint32_t w[8]; + float d; + int m0, m1; + __device__ __forceinline__ Group(const uint8_t* row, int group, + const uint32_t* __restrict__ book, + const uint32_t* masks) { + const uint8_t* b = row + (group >> 3) * block_bytes(); + const int sub = group & 7; + uint32_t meta = 0; + if constexpr (Type == 16) meta = vllm::sm70_gguf::load_u32_2(b + 6 + sub * 8); +#pragma unroll + for (int octet = 0; octet < 4; ++octet) { + int index, signs = 0, delta = 0; + if constexpr (Type == 16) { + index = b[2 + sub * 8 + octet]; + signs = (meta >> (7 * octet)) & 127; + } else if constexpr (Type == 17) { + const uint16_t code = + *reinterpret_cast(b + 2 + (sub * 4 + octet) * 2); + index = code & 511; + signs = code >> 9; + } else { + const int t = sub * 4 + octet; + const int high = (b[32 + t / 2] >> (4 * (t % 2))) & 15; + index = b[t] | ((high & 7) << 8); + delta = high >> 3; + } + uint32_t first = __ldg(book + index * 2), second = __ldg(book + index * 2 + 1); + if constexpr (Type == 29) { + const uint32_t bias = delta ? 0xffffffffU : 0x01010101U; + first = __vadd4(first, bias); + second = __vadd4(second, bias); + } else { + signs |= (__popc(signs) & 1) << 7; + const uint32_t lo = masks[signs & 15], hi = masks[signs >> 4]; + first = __vsub4(first ^ lo, lo); + second = __vsub4(second ^ hi, hi); + } + w[2 * octet] = first; + w[2 * octet + 1] = second; + } + if constexpr (Type == 29) { + const auto* scales = reinterpret_cast(b + 48); + const uint16_t bits = (scales[0] >> 12) | ((scales[1] >> 8) & 0xf0) | + ((scales[2] >> 4) & 0xf00) | (scales[3] & 0xf000); + d = __half2float(__ushort_as_half(bits)); + const int f = sub * 2, s = f + 1; + m0 = 2 * ((scales[f / 4] >> (3 * (f % 4))) & 7) + 1; + m1 = 2 * ((scales[s / 4] >> (3 * (s % 4))) & 7) + 1; + } else { + d = __half2float(*reinterpret_cast(b)); + if constexpr (Type == 16) { + m0 = m1 = 1 + 2 * (meta >> 28); + } else { + m0 = 1 + 2 * (b[66 + sub] & 15); + m1 = 1 + 2 * (b[66 + sub] >> 4); + } + } + } + // qs: 8 activation words; xd: activation Q8_1 scale. + __device__ __forceinline__ float apply(const int* qs, float xd) const { + int s0 = 0, s1 = 0; +#pragma unroll + for (int i = 0; i < 4; ++i) s0 = __dp4a(int(w[i]), qs[i], s0); +#pragma unroll + for (int i = 4; i < 8; ++i) s1 = __dp4a(int(w[i]), qs[i], s1); + return d * xd * float(s0 * m0 + s1 * m1) * 0.125f; + } +}; + +// Distinct expert `u` (first-occurrence order over routes) and its routes. +struct Members { + int expert, count; + int route[6]; +}; + +__device__ __forceinline__ void find_members(const int* ids, int routes, int u, + Members* out) { + // Executed by one warp; routes <= 36. + const int lane = threadIdx.x & 31; + __shared__ int first[64]; + int distinct = 0; + for (int r = 0; r < routes; ++r) { + const int e = ids[r]; + bool seen = false; + for (int q = lane; q < r; q += 32) seen |= ids[q] == e; + seen = __any_sync(0xffffffffU, seen); + if (!seen) { + if (lane == 0) first[distinct] = e; + ++distinct; + } + } + __syncwarp(); + if (lane == 0) { + out->expert = u < distinct ? first[u] : -1; + out->count = 0; + if (out->expert >= 0) + for (int r = 0; r < routes; ++r) + if (ids[r] == out->expert && out->count < 6) out->route[out->count++] = r; + } + __syncwarp(); +} + +template +__global__ void ids_to_int(const Index* ids, int* out, int n) { + const int i = threadIdx.x; + if (i < n) out[i] = int(ids[i]); +} + +// Gate/up for one distinct expert x 32-row tile; fused SwiGLU (clamp 10, +// route probability) and Q8_1 quantization of the down input. +// Block: 256 threads = 8 warps; warp w -> matrix w/4, rows (w%4)*8 .. +8; +// LANES lanes per row-dot. +template +__global__ void __launch_bounds__(256) gate_up_em( + Q8_1* __restrict__ down_in, const Q8_1* __restrict__ x, const int* ids, + const float* __restrict__ prob, const uint8_t* __restrict__ gate, + const uint8_t* __restrict__ up, int row_stride, + const uint32_t* __restrict__ book, int routes, int top_k) { + constexpr int GROUPS = 160, ROWS_PER_WARP = 32 / LANES; + extern __shared__ int smem_i[]; + __shared__ Members mem; + __shared__ uint32_t masks[16]; + __shared__ float red[6][2][32]; + const int t = threadIdx.x, warp = t / 32, lane = t % 32; + if (warp == 0) find_members(ids, routes, blockIdx.y, &mem); + if (t < 16) masks[t] = sign_mask(t); + __syncthreads(); + const int e = mem.expert, cnt = mem.count; + if (e < 0) return; + // Stage member activations: qs[c][g][8] ints, then d[c][g] floats. + int* xs = smem_i; + float* xd = reinterpret_cast(smem_i + 6 * GROUPS * 8); + for (int i = t; i < cnt * GROUPS * 9; i += blockDim.x) { + const int c = i / (GROUPS * 9), rest = i % (GROUPS * 9); + const int g = rest / 9, word = rest % 9; + const int token = mem.route[c] / top_k; + const uint32_t* src = + reinterpret_cast(x + token * GROUPS + g); + const uint32_t v = src[word]; + if (word == 0) + xd[c * GROUPS + g] = __low2float(*reinterpret_cast(&v)); + else + xs[(c * GROUPS + g) * 8 + word - 1] = int(v); + } + __syncthreads(); + const int mat = warp / 4; + const int local = (warp % 4) * 8 + lane / LANES * (8 / ROWS_PER_WARP); + // ROWS_PER_WARP row-dots per warp, each LANES lanes; with LANES=4 there + // are 8 row-dots per warp (= 8 rows), so `local` indexes them directly. + const int row = blockIdx.x * 32 + (warp % 4) * 8 + lane / LANES; + const uint8_t* base = (mat ? up : gate) + (int64_t(e) * 576 + row) * row_stride; + float acc[6] = {}; + for (int g = lane % LANES; g < GROUPS; g += LANES) { + const Group w(base, g, book, masks); +#pragma unroll + for (int c = 0; c < 6; ++c) + if (c < cnt) + acc[c] += w.apply(xs + (c * GROUPS + g) * 8, xd[c * GROUPS + g]); + } +#pragma unroll + for (int c = 0; c < 6; ++c) { +#pragma unroll + for (int o = LANES / 2; o; o >>= 1) + acc[c] += __shfl_xor_sync(0xffffffffU, acc[c], o); + if (c < cnt && lane % LANES == 0) + red[c][mat][(warp % 4) * 8 + lane / LANES] = acc[c]; + } + (void)local; + __syncthreads(); + if (warp < cnt) { + const int route = mem.route[warp]; + float g16 = __half2float(__float2half_rn(red[warp][0][lane])); + float u16 = __half2float(__float2half_rn(red[warp][1][lane])); + g16 = g16 > 10.f ? 10.f : g16; + u16 = u16 > 10.f ? 10.f : (u16 < -10.f ? -10.f : u16); + const float act = __fdiv_rn(g16, __fadd_rn(1.f, expf(-g16))); + const float value = __half2float( + __float2half_rn(__fmul_rn(__fmul_rn(act, u16), prob[route]))); + vllm::sm70_gguf::quantize_q8_1_warp(down_in + route * 18 + blockIdx.x, value); + } +} + +// Down for a 32*WARPS-row output tile over a chunk of distinct experts. +// Lane = output row; groups serial (18 = 576/32). Chunk partials go to `ws`; +// the last CTA of a tile sums chunks in order (deterministic). +template +__global__ void __launch_bounds__(32 * WARPS) down_em( + half* __restrict__ out, const Q8_1* __restrict__ h, const int* ids, + const uint8_t* __restrict__ down, int row_stride, int left_groups, + const uint32_t* __restrict__ book, int routes, int top_k, int tokens, + float* __restrict__ ws, int* counters) { + constexpr int GROUPS = 18, N = 5120; + __shared__ int distinct_e[36], s_nu, member_route[36][6], member_cnt[36]; + __shared__ uint32_t masks[16]; + __shared__ int xs[6][GROUPS][8]; + __shared__ float xd[6][GROUPS]; + __shared__ bool s_last; + const int t = threadIdx.x, lane = t % 32, warp = t / 32; + if (t < 16) masks[t] = sign_mask(t); + if (t == 0) { + int nu = 0; + for (int r = 0; r < routes; ++r) { + const int e = ids[r]; + int u = 0; + while (u < nu && distinct_e[u] != e) ++u; + if (u == nu) { + distinct_e[nu] = e; + member_cnt[nu] = 0; + ++nu; + } + member_route[u][member_cnt[u]++] = r; + } + s_nu = nu; + } + __syncthreads(); + const int nu = s_nu, chunks = gridDim.y, chunk = blockIdx.y; + const int row = blockIdx.x * 32 * WARPS + warp * 32 + lane; + float acc[6] = {}; // per token (tokens <= 6) + for (int u = chunk; u < nu; u += chunks) { + const int e = distinct_e[u], cnt = member_cnt[u]; + __syncthreads(); + for (int i = t; i < cnt * GROUPS * 9; i += blockDim.x) { + const int c = i / (GROUPS * 9), rest = i % (GROUPS * 9); + const int g = rest / 9, word = rest % 9; + const uint32_t v = reinterpret_cast( + h + member_route[u][c] * GROUPS + g)[word]; + if (word == 0) + xd[c][g] = __low2float(*reinterpret_cast(&v)); + else + xs[c][g][word - 1] = int(v); + } + __syncthreads(); + const uint8_t* base = down + (int64_t(e) * N + row) * row_stride; + float dot[6] = {}; + for (int g = 0; g < GROUPS; ++g) { + const Group w(base, g + left_groups, book, masks); +#pragma unroll + for (int c = 0; c < 6; ++c) + if (c < cnt) dot[c] += w.apply(xs[c][g], xd[c][g]); + } +#pragma unroll + for (int c = 0; c < 6; ++c) + if (c < cnt) { + const int token = member_route[u][c] / top_k; + const float v = __half2float(__float2half_rn(dot[c])); +#pragma unroll + for (int k = 0; k < 6; ++k) + if (k == token) acc[k] += v; + } + } + if (chunks == 1) { + for (int k = 0; k < tokens; ++k) out[k * N + row] = __float2half_rn(acc[k]); + return; + } + for (int k = 0; k < tokens; ++k) + __stcg(ws + (int64_t(chunk) * tokens + k) * N + row, acc[k]); + __threadfence(); + __syncthreads(); + if (t == 0) s_last = atomicAdd(counters + blockIdx.x, 1) == chunks - 1; + __syncthreads(); + if (!s_last) return; + __threadfence(); + for (int k = 0; k < tokens; ++k) { + float sum = 0.f; + for (int c = 0; c < chunks; ++c) + sum += __ldcg(ws + (int64_t(c) * tokens + k) * N + row); + out[k * N + row] = __float2half_rn(sum); + } + if (t == 0) counters[blockIdx.x] = 0; +} + + +// ---- v2: Codex reuse structure (16 lanes per row), 32-row tiles for the +// fused SwiGLU+Q8 epilogue, global L1 codebook, compact member lists. +template +__global__ void __launch_bounds__(512) gate_up_em2( + Q8_1* __restrict__ down_in, const Q8_1* __restrict__ x, const int* ids, + const float* __restrict__ prob, const uint8_t* __restrict__ gate, + const uint8_t* __restrict__ up, int row_stride, + const uint32_t* __restrict__ book, int routes, int top_k) { + constexpr int GROUPS = 160; + __shared__ int s_expert, s_cnt, s_route[6], s_tok[6]; + __shared__ uint32_t masks[16]; + __shared__ float red[6][2][32]; + const int t = threadIdx.x, warp = t / 32, lane = t % 32; + if (t < 16) masks[t] = sign_mask(t); + if (t == 0) { + // u-th distinct expert in first-occurrence order. + int u = blockIdx.y, e = -1, distinct = 0; + for (int r = 0; r < routes && e < 0; ++r) { + bool seen = false; + for (int q = 0; q < r; ++q) seen |= ids[q] == ids[r]; + if (!seen) { + if (distinct == u) e = ids[r]; + ++distinct; + } + } + int cnt = 0; + if (e >= 0) + for (int r = 0; r < routes; ++r) + if (ids[r] == e && cnt < 6) { + s_route[cnt] = r; + s_tok[cnt] = r / top_k; + ++cnt; + } + s_expert = e; + s_cnt = cnt; + } + __syncthreads(); + const int e = s_expert, cnt = s_cnt; + if (e < 0) return; + // 16 warps x 2 rows: lanes 0-15 row 2w, lanes 16-31 row 2w+1. + const int lrow = warp * 2 + lane / 16; + const int row = blockIdx.x * 32 + lrow; + const uint8_t* g = gate + (int64_t(e) * 576 + row) * row_stride; + const uint8_t* u = up + (int64_t(e) * 576 + row) * row_stride; + int tok[6]; +#pragma unroll + for (int c = 0; c < 6; ++c) tok[c] = c < cnt ? s_tok[c] : 0; + float gs[6] = {}, us[6] = {}; + for (int group = lane % 16; group < GROUPS; group += 16) { + const Group dg(g, group, book, masks); + const Group du(u, group, book, masks); +#pragma unroll + for (int c = 0; c < 6; ++c) + if (c < cnt) { + const Q8_1& a = x[tok[c] * GROUPS + group]; + const int* qs = reinterpret_cast(a.qs); + const float xd = __low2float(a.ds); + gs[c] += dg.apply(qs, xd); + us[c] += du.apply(qs, xd); + } + } +#pragma unroll + for (int c = 0; c < 6; ++c) { +#pragma unroll + for (int o = 8; o; o >>= 1) { + gs[c] += __shfl_xor_sync(0xffffffffU, gs[c], o); + us[c] += __shfl_xor_sync(0xffffffffU, us[c], o); + } + if (c < cnt && lane % 16 == 0) { + red[c][0][lrow] = gs[c]; + red[c][1][lrow] = us[c]; + } + } + __syncthreads(); + if (warp < cnt) { + const int route = s_route[warp]; + float g16 = __half2float(__float2half_rn(red[warp][0][lane])); + float u16 = __half2float(__float2half_rn(red[warp][1][lane])); + g16 = g16 > 10.f ? 10.f : g16; + u16 = u16 > 10.f ? 10.f : (u16 < -10.f ? -10.f : u16); + const float act = __fdiv_rn(g16, __fadd_rn(1.f, expf(-g16))); + const float value = __half2float( + __float2half_rn(__fmul_rn(__fmul_rn(act, u16), prob[route]))); + vllm::sm70_gguf::quantize_q8_1_warp(down_in + route * 18 + blockIdx.x, value); + } +} + +// Down v2: CTA = 8 output rows (one warp-row set), 4 warps split the +// distinct experts; 4 lanes per row over the 18 groups; per-token sums are +// reduced across warps in shared memory (fixed order: deterministic). +template +__global__ void __launch_bounds__(128) down_em2( + half* __restrict__ out, const Q8_1* __restrict__ h, const int* ids, + const uint8_t* __restrict__ down, int row_stride, int left_groups, + const uint32_t* __restrict__ book, int routes, int top_k, int tokens) { + constexpr int GROUPS = 18, N = 5120; + __shared__ int distinct_e[36], s_nu, member_route[36][6], member_cnt[36]; + __shared__ uint32_t masks[16]; + __shared__ float part[4][8][6]; + const int t = threadIdx.x, lane = t % 32, warp = t / 32; + if (t < 16) masks[t] = sign_mask(t); + if (t == 0) { + int nu = 0; + for (int r = 0; r < routes; ++r) { + const int e = ids[r]; + int q = 0; + while (q < nu && distinct_e[q] != e) ++q; + if (q == nu) { + distinct_e[nu] = e; + member_cnt[nu] = 0; + ++nu; + } + member_route[q][member_cnt[q]++] = r; + } + s_nu = nu; + } + __syncthreads(); + const int nu = s_nu; + const int lrow = lane / 4, sub = lane % 4; + const int row = blockIdx.x * 8 + lrow; + float acc[6] = {}; + for (int q = warp; q < nu; q += 4) { + const int e = distinct_e[q], cnt = member_cnt[q]; + const uint8_t* base = down + (int64_t(e) * N + row) * row_stride; + int rt[6]; +#pragma unroll + for (int c = 0; c < 6; ++c) rt[c] = c < cnt ? member_route[q][c] : 0; + float dot[6] = {}; + for (int g = sub; g < GROUPS; g += 4) { + const Group w(base, g + left_groups, book, masks); +#pragma unroll + for (int c = 0; c < 6; ++c) + if (c < cnt) { + const Q8_1& a = h[rt[c] * GROUPS + g]; + dot[c] += w.apply(reinterpret_cast(a.qs), __low2float(a.ds)); + } + } +#pragma unroll + for (int c = 0; c < 6; ++c) { + dot[c] += __shfl_xor_sync(0xffffffffU, dot[c], 1); + dot[c] += __shfl_xor_sync(0xffffffffU, dot[c], 2); + } +#pragma unroll + for (int c = 0; c < 6; ++c) + if (c < cnt) { + const int token = rt[c] / top_k; + const float v = __half2float(__float2half_rn(dot[c])); +#pragma unroll + for (int k = 0; k < 6; ++k) + if (k == token) acc[k] += v; + } + } + if (sub == 0) +#pragma unroll + for (int k = 0; k < 6; ++k) part[warp][lrow][k] = acc[k]; + __syncthreads(); + if (t < 8 * tokens) { + const int r = t / tokens, k = t % tokens; + const float v = part[0][r][k] + part[1][r][k] + part[2][r][k] + part[3][r][k]; + out[k * N + blockIdx.x * 8 + r] = __float2half_rn(v); + } +} + + +// ---- diagnostics: bandwidth floor and decode-only cost (gate+up, 18 tiles) +__global__ void __launch_bounds__(512) gu_read_only( + float* sink, const int* ids, const uint8_t* __restrict__ gate, + const uint8_t* __restrict__ up, int row_stride, int routes) { + __shared__ int s_expert; + if (threadIdx.x == 0) { + int u = blockIdx.y, e = -1, distinct = 0; + for (int r = 0; r < routes && e < 0; ++r) { + bool seen = false; + for (int q = 0; q < r; ++q) seen |= ids[q] == ids[r]; + if (!seen) { if (distinct == u) e = ids[r]; ++distinct; } + } + s_expert = e; + } + __syncthreads(); + const int e = s_expert; + if (e < 0) return; + const int64_t bytes = int64_t(32) * row_stride; // 32 rows per tile + const uint4* g = reinterpret_cast(gate + (int64_t(e) * 576 + blockIdx.x * 32) * row_stride); + const uint4* u = reinterpret_cast(up + (int64_t(e) * 576 + blockIdx.x * 32) * row_stride); + uint32_t acc = 0; + for (int64_t i = threadIdx.x; i < bytes / 16; i += blockDim.x) { + const uint4 a = __ldg(g + i), b = __ldg(u + i); + acc ^= a.x ^ a.y ^ a.z ^ a.w ^ b.x ^ b.y ^ b.z ^ b.w; + } + if (acc == 0x12345678u) sink[0] = 1.f; +} + +template +__global__ void __launch_bounds__(512) gu_decode_only( + float* sink, const int* ids, const uint8_t* __restrict__ gate, + const uint8_t* __restrict__ up, int row_stride, + const uint32_t* __restrict__ book, int routes) { + __shared__ int s_expert; + __shared__ uint32_t masks[16]; + const int t = threadIdx.x, warp = t / 32, lane = t % 32; + if (t < 16) masks[t] = sign_mask(t); + if (t == 0) { + int u = blockIdx.y, e = -1, distinct = 0; + for (int r = 0; r < routes && e < 0; ++r) { + bool seen = false; + for (int q = 0; q < r; ++q) seen |= ids[q] == ids[r]; + if (!seen) { if (distinct == u) e = ids[r]; ++distinct; } + } + s_expert = e; + } + __syncthreads(); + const int e = s_expert; + if (e < 0) return; + const int row = blockIdx.x * 32 + warp * 2 + lane / 16; + const uint8_t* g = gate + (int64_t(e) * 576 + row) * row_stride; + const uint8_t* u = up + (int64_t(e) * 576 + row) * row_stride; + uint32_t acc = 0; + float f = 0.f; + for (int group = lane % 16; group < 160; group += 16) { + const Group dg(g, group, book, masks); + const Group du(u, group, book, masks); +#pragma unroll + for (int i = 0; i < 8; ++i) acc ^= dg.w[i] ^ du.w[i]; + f += dg.d * dg.m0 + du.d * du.m1; + } + if (acc == 0x12345678u || f == 1234.5f) sink[0] = 1.f; +} + + +// ---- v3: vector loads + shared-memory codebook ---------------------------- +__device__ __forceinline__ void ld3(const uint8_t* p, uint32_t& w0, uint32_t& w1) { + // 8 bytes at any 2-byte-aligned address: three aligned words + funnel shifts. + const uintptr_t a = reinterpret_cast(p); + const uint32_t* q = reinterpret_cast(a & ~uintptr_t(3)); + const uint32_t sh = (a & 3) * 8; + const uint32_t x0 = __ldg(q), x1 = __ldg(q + 1), x2 = __ldg(q + 2); + w0 = __funnelshift_r(x0, x1, sh); + w1 = __funnelshift_r(x1, x2, sh); +} +__device__ __forceinline__ uint32_t ld16(const uint8_t* p) { + return *reinterpret_cast(p); +} + +#ifndef V41K_GU_UNROLL +#define V41K_GU_UNROLL 2 +#endif +constexpr int kGuUnroll = V41K_GU_UNROLL; +template +// Codebook entries as one 64-bit LDS; sign masks computed (PRMT), not looked up. +#ifndef V41K_SMEM_MASKS +#define V41K_MASK(n) sign_mask(n) +#else +#define V41K_MASK(n) masks[n] +#endif +struct Group3 { + uint32_t w[8]; + float d; + int m0, m1; + __device__ __forceinline__ Group3(const uint8_t* row, int group, + const uint32_t* book, const uint32_t* masks) { + const uint8_t* b = row + (group >> 3) * block_bytes(); + const int sub = group & 7; + if constexpr (Type == 29) { + const uint32_t qs = __ldg(reinterpret_cast(b + 4 * sub)); + const uint32_t qh = ld16(b + 32 + 2 * sub); + const uint2 sc = __ldg(reinterpret_cast(b + 48)); +#pragma unroll + for (int octet = 0; octet < 4; ++octet) { + const int high = (qh >> (4 * octet)) & 15; + const int index = ((qs >> (8 * octet)) & 255) | ((high & 7) << 8); + const uint32_t bias = (high >> 3) ? 0xffffffffU : 0x01010101U; + const uint2 e = reinterpret_cast(book)[index]; + w[2 * octet] = __vadd4(e.x, bias); + w[2 * octet + 1] = __vadd4(e.y, bias); + } + const uint32_t s01 = sc.x, s23 = sc.y; + const uint16_t bits = ((s01 & 0xffff) >> 12) | (((s01 >> 16) >> 8) & 0xf0) | + (((s23 & 0xffff) >> 4) & 0xf00) | ((s23 >> 16) & 0xf000); + d = __half2float(__ushort_as_half(bits)); + const uint64_t all = (uint64_t(s23) << 32) | s01; + const int f = sub * 2; + // scales[f/4] >> 3*(f%4) on 16-bit lanes + m0 = 2 * ((all >> (16 * (f / 4) + 3 * (f % 4))) & 7) + 1; + m1 = 2 * ((all >> (16 * ((f + 1) / 4) + 3 * ((f + 1) % 4))) & 7) + 1; + } else { + uint32_t w0, w1; + ld3(b + 2 + 8 * sub, w0, w1); + d = __half2float(__ushort_as_half(ld16(b))); + if constexpr (Type == 16) { + const uint32_t meta = w1; +#pragma unroll + for (int octet = 0; octet < 4; ++octet) { + const int index = (w0 >> (8 * octet)) & 255; + int signs = (meta >> (7 * octet)) & 127; + signs |= (__popc(signs) & 1) << 7; + const uint32_t lo = V41K_MASK(signs & 15), hi = V41K_MASK(signs >> 4); + const uint2 e = reinterpret_cast(book)[index]; + w[2 * octet] = __vsub4(e.x ^ lo, lo); + w[2 * octet + 1] = __vsub4(e.y ^ hi, hi); + } + m0 = m1 = 1 + 2 * (meta >> 28); + } else { +#pragma unroll + for (int octet = 0; octet < 4; ++octet) { + const uint32_t code = ((octet < 2 ? w0 : w1) >> (16 * (octet & 1))) & 0xffff; + const int index = code & 511; + int signs = code >> 9; + signs |= (__popc(signs) & 1) << 7; + const uint32_t lo = V41K_MASK(signs & 15), hi = V41K_MASK(signs >> 4); + const uint2 e = reinterpret_cast(book)[index]; + w[2 * octet] = __vsub4(e.x ^ lo, lo); + w[2 * octet + 1] = __vsub4(e.y ^ hi, hi); + } + const int sc = b[66 + sub]; + m0 = 1 + 2 * (sc & 15); + m1 = 1 + 2 * (sc >> 4); + } + } + } + __device__ __forceinline__ float apply(const int* qs, float xd) const { + int s0 = 0, s1 = 0; +#pragma unroll + for (int i = 0; i < 4; ++i) s0 = __dp4a(int(w[i]), qs[i], s0); +#pragma unroll + for (int i = 4; i < 8; ++i) s1 = __dp4a(int(w[i]), qs[i], s1); + return d * xd * float(s0 * m0 + s1 * m1) * 0.125f; + } +}; + +template +constexpr int book_words() { + return turbomind::gemm::LatticeCodebook::kBytes / 4; +} + +__device__ __forceinline__ int find_distinct(const int* ids, int routes, int u) { + int e = -1, distinct = 0; + for (int r = 0; r < routes && e < 0; ++r) { + bool seen = false; + for (int q = 0; q < r; ++q) seen |= ids[q] == ids[r]; + if (!seen) { + if (distinct == u) e = ids[r]; + ++distinct; + } + } + return e; +} + +template +__global__ void __launch_bounds__(512) gate_up_em3( + Q8_1* __restrict__ down_in, const Q8_1* __restrict__ x, const int* ids, + const float* __restrict__ prob, const uint8_t* __restrict__ gate, + const uint8_t* __restrict__ up, int row_stride, + const uint32_t* __restrict__ gbook, int routes, int top_k) { + constexpr int GROUPS = 160, BW = book_words(); + __shared__ __align__(16) uint32_t book[BW]; + __shared__ int s_expert, s_cnt, s_route[6], s_tok[6]; + __shared__ uint32_t masks[16]; + __shared__ float red[6][2][32]; + const int t = threadIdx.x, warp = t / 32, lane = t % 32; + if (t == 0) { + const int e = find_distinct(ids, routes, blockIdx.y); + int cnt = 0; + if (e >= 0) + for (int r = 0; r < routes; ++r) + if (ids[r] == e && cnt < 6) { + s_route[cnt] = r; + s_tok[cnt] = r / top_k; + ++cnt; + } + s_expert = e; + s_cnt = cnt; + } + __syncthreads(); + const int e = s_expert, cnt = s_cnt; + if (e < 0) return; + for (int i = t; i < BW; i += blockDim.x) book[i] = __ldg(gbook + i); + if (t < 16) masks[t] = sign_mask(t); + __syncthreads(); + const int lrow = warp * 2 + lane / 16; + const int row = blockIdx.x * 32 + lrow; + const uint8_t* g = gate + (int64_t(e) * 576 + row) * row_stride; + const uint8_t* u = up + (int64_t(e) * 576 + row) * row_stride; + int tok[6]; +#pragma unroll + for (int c = 0; c < 6; ++c) tok[c] = c < cnt ? s_tok[c] : 0; + float gs[6] = {}, us[6] = {}; +#pragma unroll 2 + for (int group = lane % 16; group < GROUPS; group += 16) { + const Group3 dg(g, group, book, masks); + const Group3 du(u, group, book, masks); +#pragma unroll + for (int c = 0; c < 6; ++c) + if (c < cnt) { + const Q8_1& a = x[tok[c] * GROUPS + group]; + const int* qs = reinterpret_cast(a.qs); + const float xd = __low2float(a.ds); + gs[c] += dg.apply(qs, xd); + us[c] += du.apply(qs, xd); + } + } +#pragma unroll + for (int c = 0; c < 6; ++c) { +#pragma unroll + for (int o = 8; o; o >>= 1) { + gs[c] += __shfl_xor_sync(0xffffffffU, gs[c], o); + us[c] += __shfl_xor_sync(0xffffffffU, us[c], o); + } + if (c < cnt && lane % 16 == 0) { + red[c][0][lrow] = gs[c]; + red[c][1][lrow] = us[c]; + } + } + __syncthreads(); + if (warp < cnt) { + const int route = s_route[warp]; + float g16 = __half2float(__float2half_rn(red[warp][0][lane])); + float u16 = __half2float(__float2half_rn(red[warp][1][lane])); + g16 = g16 > 10.f ? 10.f : g16; + u16 = u16 > 10.f ? 10.f : (u16 < -10.f ? -10.f : u16); + const float act = __fdiv_rn(g16, __fadd_rn(1.f, expf(-g16))); + const float value = __half2float( + __float2half_rn(__fmul_rn(__fmul_rn(act, u16), prob[route]))); + vllm::sm70_gguf::quantize_q8_1_warp(down_in + route * 18 + blockIdx.x, value); + } +} + +template +__global__ void __launch_bounds__(128) down_em3( + half* __restrict__ out, const Q8_1* __restrict__ h, const int* ids, + const uint8_t* __restrict__ down, int row_stride, int left_groups, + const uint32_t* __restrict__ gbook, int routes, int top_k, int tokens, + int rows_per_cta) { + constexpr int GROUPS = 18, N = 5120, BW = book_words(); + __shared__ __align__(16) uint32_t book[BW]; + __shared__ int distinct_e[36], s_nu, member_route[36][6], member_cnt[36]; + __shared__ uint32_t masks[16]; + __shared__ float part[4][8][6]; + const int t = threadIdx.x, lane = t % 32, warp = t / 32; + for (int i = t; i < BW; i += blockDim.x) book[i] = __ldg(gbook + i); + if (t < 16) masks[t] = sign_mask(t); + if (t == 0) { + int nu = 0; + for (int r = 0; r < routes; ++r) { + const int e = ids[r]; + int q = 0; + while (q < nu && distinct_e[q] != e) ++q; + if (q == nu) { + distinct_e[nu] = e; + member_cnt[nu] = 0; + ++nu; + } + member_route[q][member_cnt[q]++] = r; + } + s_nu = nu; + } + __syncthreads(); + const int nu = s_nu; + const int lrow = lane / 4, sub = lane % 4; + for (int tile = 0; tile < rows_per_cta / 8; ++tile) { + const int row0 = (blockIdx.x * rows_per_cta / 8 + tile) * 8; + const int row = row0 + lrow; + float acc[6] = {}; + for (int q = warp; q < nu; q += 4) { + const int e = distinct_e[q], cnt = member_cnt[q]; + const uint8_t* base = down + (int64_t(e) * N + row) * row_stride; + int rt[6]; +#pragma unroll + for (int c = 0; c < 6; ++c) rt[c] = c < cnt ? member_route[q][c] : 0; + float dot[6] = {}; + for (int g = sub; g < GROUPS; g += 4) { + const Group3 w(base, g + left_groups, book, masks); +#pragma unroll + for (int c = 0; c < 6; ++c) + if (c < cnt) { + const Q8_1& a = h[rt[c] * GROUPS + g]; + dot[c] += w.apply(reinterpret_cast(a.qs), __low2float(a.ds)); + } + } +#pragma unroll + for (int c = 0; c < 6; ++c) { + dot[c] += __shfl_xor_sync(0xffffffffU, dot[c], 1); + dot[c] += __shfl_xor_sync(0xffffffffU, dot[c], 2); + } +#pragma unroll + for (int c = 0; c < 6; ++c) + if (c < cnt) { + const int token = rt[c] / top_k; + const float v = __half2float(__float2half_rn(dot[c])); +#pragma unroll + for (int k = 0; k < 6; ++k) + if (k == token) acc[k] += v; + } + } + __syncthreads(); + if (sub == 0) +#pragma unroll + for (int k = 0; k < 6; ++k) part[warp][lrow][k] = acc[k]; + __syncthreads(); + if (t < 8 * tokens) { + const int r = t / tokens, k = t % tokens; + const float v = part[0][r][k] + part[1][r][k] + part[2][r][k] + part[3][r][k]; + out[k * N + row0 + r] = __float2half_rn(v); + } + } +} + +template +__global__ void __launch_bounds__(512) gu_decode3( + float* sink, const int* ids, const uint8_t* __restrict__ gate, + const uint8_t* __restrict__ up, int row_stride, + const uint32_t* __restrict__ gbook, int routes) { + constexpr int BW = book_words(); + __shared__ __align__(16) uint32_t book[BW]; + __shared__ int s_expert; + __shared__ uint32_t masks[16]; + const int t = threadIdx.x, warp = t / 32, lane = t % 32; + if (t == 0) s_expert = find_distinct(ids, routes, blockIdx.y); + __syncthreads(); + const int e = s_expert; + if (e < 0) return; + for (int i = t; i < BW; i += blockDim.x) book[i] = __ldg(gbook + i); + if (t < 16) masks[t] = sign_mask(t); + __syncthreads(); + const int row = blockIdx.x * 32 + warp * 2 + lane / 16; + const uint8_t* g = gate + (int64_t(e) * 576 + row) * row_stride; + const uint8_t* u = up + (int64_t(e) * 576 + row) * row_stride; + uint32_t acc = 0; + float f = 0.f; + for (int group = lane % 16; group < 160; group += 16) { + const Group3 dg(g, group, book, masks); + const Group3 du(u, group, book, masks); +#pragma unroll + for (int i = 0; i < 8; ++i) acc ^= dg.w[i] ^ du.w[i]; + f += dg.d * dg.m0 + du.d * du.m1; + } + if (acc == 0x12345678u || f == 1234.5f) sink[0] = 1.f; +} + + +// ---- v4: v3 decode + activations staged once in aligned shared memory ---- +// x layout in smem: qs as int4 pairs [tok][group][2] (32 B), d as float. +template +__global__ void __launch_bounds__(512) gate_up_em4( + Q8_1* __restrict__ down_in, const Q8_1* __restrict__ x, const int* ids, + const float* __restrict__ prob, const uint8_t* __restrict__ gate, + const uint8_t* __restrict__ up, int row_stride, + const uint32_t* __restrict__ gbook, int routes, int top_k, int tokens) { + constexpr int GROUPS = 160, BW = book_words(); + extern __shared__ int4 dyn[]; + int4* xq = dyn; // [6][160][2] + float* xd = reinterpret_cast(dyn + 6 * GROUPS * 2); // [6][160] + uint32_t* book = reinterpret_cast(xd + 6 * GROUPS); + __shared__ int s_expert, s_cnt, s_route[6], s_tok[6]; + __shared__ uint32_t masks[16]; + __shared__ float red[6][2][32]; + const int t = threadIdx.x, warp = t / 32, lane = t % 32; + if (t == 0) { + const int e = find_distinct(ids, routes, blockIdx.y); + int cnt = 0; + if (e >= 0) + for (int r = 0; r < routes; ++r) + if (ids[r] == e && cnt < 6) { + s_route[cnt] = r; + s_tok[cnt] = r / top_k; + ++cnt; + } + s_expert = e; + s_cnt = cnt; + } + __syncthreads(); + const int e = s_expert, cnt = s_cnt; + if (e < 0) return; + for (int i = t; i < BW; i += blockDim.x) book[i] = __ldg(gbook + i); + if (t < 16) masks[t] = sign_mask(t); + // stage only member tokens, compacted: slot c <- token s_tok[c] + for (int i = t; i < cnt * GROUPS * 9; i += blockDim.x) { + const int c = i / (GROUPS * 9), rest = i % (GROUPS * 9); + const int g = rest / 9, word = rest % 9; + const uint32_t v = reinterpret_cast(x + s_tok[c] * GROUPS + g)[word]; + if (word == 0) + xd[c * GROUPS + g] = __low2float(*reinterpret_cast(&v)); + else + reinterpret_cast(xq + (c * GROUPS + g) * 2)[word - 1] = v; + } + __syncthreads(); + const int lrow = warp * 2 + lane / 16; + const int row = blockIdx.x * 32 + lrow; + const uint8_t* g = gate + (int64_t(e) * 576 + row) * row_stride; + const uint8_t* u = up + (int64_t(e) * 576 + row) * row_stride; + float gs[6] = {}, us[6] = {}; +#pragma unroll 2 + for (int group = lane % 16; group < GROUPS; group += 16) { + const Group3 dg(g, group, book, masks); + const Group3 du(u, group, book, masks); +#pragma unroll + for (int c = 0; c < 6; ++c) + if (c < cnt) { + const int4 a0 = xq[(c * GROUPS + group) * 2], a1 = xq[(c * GROUPS + group) * 2 + 1]; + const int qs[8] = {a0.x, a0.y, a0.z, a0.w, a1.x, a1.y, a1.z, a1.w}; + const float d = xd[c * GROUPS + group]; + gs[c] += dg.apply(qs, d); + us[c] += du.apply(qs, d); + } + } +#pragma unroll + for (int c = 0; c < 6; ++c) { +#pragma unroll + for (int o = 8; o; o >>= 1) { + gs[c] += __shfl_xor_sync(0xffffffffU, gs[c], o); + us[c] += __shfl_xor_sync(0xffffffffU, us[c], o); + } + if (c < cnt && lane % 16 == 0) { + red[c][0][lrow] = gs[c]; + red[c][1][lrow] = us[c]; + } + } + __syncthreads(); + if (warp < cnt) { + const int route = s_route[warp]; + float g16 = __half2float(__float2half_rn(red[warp][0][lane])); + float u16 = __half2float(__float2half_rn(red[warp][1][lane])); + g16 = g16 > 10.f ? 10.f : g16; + u16 = u16 > 10.f ? 10.f : (u16 < -10.f ? -10.f : u16); + const float act = __fdiv_rn(g16, __fadd_rn(1.f, expf(-g16))); + const float value = __half2float( + __float2half_rn(__fmul_rn(__fmul_rn(act, u16), prob[route]))); + vllm::sm70_gguf::quantize_q8_1_warp(down_in + route * 18 + blockIdx.x, value); + } + (void)tokens; +} + +template +__global__ void __launch_bounds__(128) down_em4( + half* __restrict__ out, const Q8_1* __restrict__ h, const int* ids, + const uint8_t* __restrict__ down, int row_stride, int left_groups, + const uint32_t* __restrict__ gbook, int routes, int top_k, int tokens, + int rows_per_cta) { + constexpr int GROUPS = 18, N = 5120, BW = book_words(); + __shared__ __align__(16) uint32_t book[BW]; + __shared__ int4 hq[36 * GROUPS * 2]; + __shared__ float hd[36 * GROUPS]; + __shared__ int distinct_e[36], s_nu, member_route[36][6], member_cnt[36]; + __shared__ uint32_t masks[16]; + __shared__ float part[4][8][6]; + const int t = threadIdx.x, lane = t % 32, warp = t / 32; + for (int i = t; i < BW; i += blockDim.x) book[i] = __ldg(gbook + i); + if (t < 16) masks[t] = sign_mask(t); + for (int i = t; i < routes * GROUPS * 9; i += blockDim.x) { + const int rg = i / 9, word = i % 9; + const uint32_t v = reinterpret_cast(h + rg)[word]; + if (word == 0) + hd[rg] = __low2float(*reinterpret_cast(&v)); + else + reinterpret_cast(hq + rg * 2)[word - 1] = v; + } + if (t == 0) { + int nu = 0; + for (int r = 0; r < routes; ++r) { + const int e = ids[r]; + int q = 0; + while (q < nu && distinct_e[q] != e) ++q; + if (q == nu) { + distinct_e[nu] = e; + member_cnt[nu] = 0; + ++nu; + } + member_route[q][member_cnt[q]++] = r; + } + s_nu = nu; + } + __syncthreads(); + const int nu = s_nu; + const int lrow = lane / 4, sub = lane % 4; + for (int tile = 0; tile < rows_per_cta / 8; ++tile) { + const int row0 = (blockIdx.x * rows_per_cta / 8 + tile) * 8; + const int row = row0 + lrow; + float acc[6] = {}; + for (int q = warp; q < nu; q += 4) { + const int e = distinct_e[q], cnt = member_cnt[q]; + const uint8_t* base = down + (int64_t(e) * N + row) * row_stride; + int rt[6]; +#pragma unroll + for (int c = 0; c < 6; ++c) rt[c] = c < cnt ? member_route[q][c] : 0; + float dot[6] = {}; + for (int g = sub; g < GROUPS; g += 4) { + const Group3 w(base, g + left_groups, book, masks); +#pragma unroll + for (int c = 0; c < 6; ++c) + if (c < cnt) { + const int rg = rt[c] * GROUPS + g; + const int4 a0 = hq[rg * 2], a1 = hq[rg * 2 + 1]; + const int qs[8] = {a0.x, a0.y, a0.z, a0.w, a1.x, a1.y, a1.z, a1.w}; + dot[c] += w.apply(qs, hd[rg]); + } + } +#pragma unroll + for (int c = 0; c < 6; ++c) { + dot[c] += __shfl_xor_sync(0xffffffffU, dot[c], 1); + dot[c] += __shfl_xor_sync(0xffffffffU, dot[c], 2); + } +#pragma unroll + for (int c = 0; c < 6; ++c) + if (c < cnt) { + const int token = rt[c] / top_k; + const float v = __half2float(__float2half_rn(dot[c])); +#pragma unroll + for (int k = 0; k < 6; ++k) + if (k == token) acc[k] += v; + } + } + __syncthreads(); + if (sub == 0) +#pragma unroll + for (int k = 0; k < 6; ++k) part[warp][lrow][k] = acc[k]; + __syncthreads(); + if (t < 8 * tokens) { + const int r = t / tokens, k = t % tokens; + const float v = part[0][r][k] + part[1][r][k] + part[2][r][k] + part[3][r][k]; + out[k * N + row0 + r] = __float2half_rn(v); + } + } +} + + +// ---- FP8 (e4m3, UE8M0 32x32 block scales) small-M GEMV for the DSpark draft. +// CTA = 32 output rows (one scale row block), 8 warps x 4 rows; lanes stride +// K in 16-byte chunks. e4m3 -> fp16 by bit placement, x 2^8 folded into the +// block scale (exact, subnormals included). Split-K partials are reduced in a +// fixed order by fp8_reduce (deterministic). +template +__global__ void __launch_bounds__(256) fp8_gemv_kernel( + float* __restrict__ partials, half* __restrict__ out, + const half* __restrict__ x, const uint8_t* __restrict__ w, + const uint8_t* __restrict__ scales, int N, int K, int direct, int ksplit) { + const int warp = threadIdx.x / 32, lane = threadIdx.x % 32; + const int n0 = blockIdx.x * 32 + warp * 4; + const int s = blockIdx.y; + const int kbeg = s * ksplit, kend = min(K, kbeg + ksplit); + const int sgroups = K / 32; + float acc[4][M]; +#pragma unroll + for (int r = 0; r < 4; ++r) +#pragma unroll + for (int m = 0; m < M; ++m) acc[r][m] = 0.f; + for (int k0 = kbeg + lane * 16; k0 < kend; k0 += 512) { + float xv[M][16]; +#pragma unroll + for (int m = 0; m < M; ++m) { + const uint4 a = __ldg(reinterpret_cast(x + int64_t(m) * K + k0)); + const uint4 b = __ldg(reinterpret_cast(x + int64_t(m) * K + k0 + 8)); + const half2* ha = reinterpret_cast(&a); + const half2* hb = reinterpret_cast(&b); +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float2 fa = __half22float2(ha[i]), fb = __half22float2(hb[i]); + xv[m][2 * i] = fa.x; xv[m][2 * i + 1] = fa.y; + xv[m][8 + 2 * i] = fb.x; xv[m][8 + 2 * i + 1] = fb.y; + } + } +#pragma unroll + for (int r = 0; r < 4; ++r) { + const int n = n0 + r; + if (n >= N) break; + const uint4 q = __ldg(reinterpret_cast(w + int64_t(n) * K + k0)); + const int e = scales[(n / 32) * sgroups + k0 / 32]; + const float scale = e + 8 < 255 ? __int_as_float((e + 8) << 23) + : __int_as_float(e << 23) * 256.f; + const uint32_t words[4] = {q.x, q.y, q.z, q.w}; + float dot[M]; +#pragma unroll + for (int m = 0; m < M; ++m) dot[m] = 0.f; +#pragma unroll + for (int i = 0; i < 16; ++i) { + const uint32_t b = (words[i / 4] >> (8 * (i % 4))) & 0xffu; + const uint16_t bits = uint16_t(((b & 0x7fu) << 7) | ((b & 0x80u) << 8)); + const float v = __half2float(__ushort_as_half(bits)); +#pragma unroll + for (int m = 0; m < M; ++m) dot[m] = fmaf(v, xv[m][i], dot[m]); + } +#pragma unroll + for (int m = 0; m < M; ++m) acc[r][m] = fmaf(dot[m], scale, acc[r][m]); + } + } +#pragma unroll + for (int r = 0; r < 4; ++r) +#pragma unroll + for (int m = 0; m < M; ++m) { + float v = acc[r][m]; +#pragma unroll + for (int o = 16; o; o >>= 1) v += __shfl_xor_sync(0xffffffffU, v, o); + acc[r][m] = v; + } + if (lane == 0) { +#pragma unroll + for (int r = 0; r < 4; ++r) { + const int n = n0 + r; + if (n >= N) break; +#pragma unroll + for (int m = 0; m < M; ++m) { + if (direct) + out[int64_t(m) * N + n] = __float2half_rn(acc[r][m]); + else + partials[(int64_t(s) * M + m) * N + n] = acc[r][m]; + } + } + } +} + +__global__ void fp8_reduce(half* out, const float* partials, const half* bias, + int M, int N, int split) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= M * N) return; + float v = 0.f; + for (int s = 0; s < split; ++s) v += partials[int64_t(s) * M * N + i]; + if (bias) v += __half2float(bias[i % N]); + out[i] = __float2half_rn(v); +} + + +// ---- v5: aligned activations (int8 qs as two int4 + float scale arrays) ---- +__global__ void q8_aligned_kernel(int8_t* __restrict__ qs, float* __restrict__ ds, + const half* __restrict__ x, int K) { + // one warp per 32-element group + const int g = blockIdx.x * (blockDim.x / 32) + threadIdx.x / 32; + const int lane = threadIdx.x % 32, row = blockIdx.y, groups = K / 32; + if (g >= groups) return; + const float v = __half2float(x[int64_t(row) * K + g * 32 + lane]); + float m = fabsf(v); +#pragma unroll + for (int o = 16; o; o >>= 1) m = fmaxf(m, __shfl_xor_sync(0xffffffffU, m, o)); + const float d = m / 127.f; + qs[(int64_t(row) * groups + g) * 32 + lane] = m == 0.f ? 0 : int8_t(roundf(v / d)); + if (!lane) ds[int64_t(row) * groups + g] = __half2float(__float2half_rn(d)); +} + +__device__ __forceinline__ void load_q8(const int8_t* qs, const float* ds, int64_t idx, + int (&q)[8], float& d) { + const int4* p = reinterpret_cast(qs + idx * 32); + const int4 a = __ldg(p), b = __ldg(p + 1); + q[0] = a.x; q[1] = a.y; q[2] = a.z; q[3] = a.w; + q[4] = b.x; q[5] = b.y; q[6] = b.z; q[7] = b.w; + d = __ldg(ds + idx); +} + +__device__ __forceinline__ void q8_warp_aligned(int8_t* qs, float* ds, int64_t idx, float value) { + const int lane = threadIdx.x % 32; + float m = fabsf(value); +#pragma unroll + for (int o = 16; o; o >>= 1) m = fmaxf(m, __shfl_xor_sync(0xffffffffU, m, o)); + const float d = m / 127.f; + qs[idx * 32 + lane] = m == 0.f ? 0 : int8_t(roundf(value / d)); + if (!lane) ds[idx] = __half2float(__float2half_rn(d)); +} + +#ifndef V41K_GU_MINB +#define V41K_GU_MINB 2 +#endif +template +__global__ void __launch_bounds__(512, V41K_GU_MINB) gate_up_em5( + int8_t* __restrict__ hq, float* __restrict__ hd, const int8_t* __restrict__ xq, + const float* __restrict__ xd, const int* ids, const float* __restrict__ prob, + const uint8_t* __restrict__ gate, const uint8_t* __restrict__ up, int row_stride, + const uint32_t* __restrict__ gbook, int routes, int top_k, const int* __restrict__ plan) { + constexpr int GROUPS = 160, BW = book_words(); + __shared__ __align__(16) uint32_t book[BW]; + __shared__ int s_expert, s_cnt, s_route[6], s_tok[6]; + __shared__ uint32_t masks[16]; + __shared__ float red[6][2][32]; + const int t = threadIdx.x, warp = t / 32, lane = t % 32; + if (plan != nullptr) { + // per-layer plan: [0] distinct count, [1+u] expert, [37+u] routes, [73+6u+c] route + const int u = blockIdx.y; + if (u >= __ldg(plan)) return; + for (int i = t; i < BW; i += blockDim.x) book[i] = __ldg(gbook + i); + if (t < 16) masks[t] = sign_mask(t); + if (t < 6) { + const int cnt = __ldg(plan + 37 + u); + const int r = t < cnt ? __ldg(plan + 73 + 6 * u + t) : 0; + s_route[t] = r; + s_tok[t] = r / top_k; + if (t == 0) { s_expert = __ldg(plan + 1 + u); s_cnt = cnt; } + } + } else if (t == 0) { + const int e = find_distinct(ids, routes, blockIdx.y); + int cnt = 0; + if (e >= 0) + for (int r = 0; r < routes; ++r) + if (ids[r] == e && cnt < 6) { s_route[cnt] = r; s_tok[cnt] = r / top_k; ++cnt; } + s_expert = e; + s_cnt = cnt; + } + __syncthreads(); + const int e = s_expert, cnt = s_cnt; + if (plan == nullptr) { + if (e < 0) return; + for (int i = t; i < BW; i += blockDim.x) book[i] = __ldg(gbook + i); + if (t < 16) masks[t] = sign_mask(t); + __syncthreads(); + } + const int lrow = warp * 2 + lane / 16; + const int row = blockIdx.x * 32 + lrow; + const uint8_t* g = gate + (int64_t(e) * 576 + row) * row_stride; + const uint8_t* u = up + (int64_t(e) * 576 + row) * row_stride; + int tok[6]; +#pragma unroll + for (int c = 0; c < 6; ++c) tok[c] = c < cnt ? s_tok[c] : 0; + float gs[6] = {}, us[6] = {}; +#pragma unroll kGuUnroll + for (int group = lane % 16; group < GROUPS; group += 16) { + const Group3 dg(g, group, book, masks); + const Group3 du(u, group, book, masks); +#pragma unroll + for (int c = 0; c < 6; ++c) + if (c < cnt) { + int q[8]; + float d; + load_q8(xq, xd, int64_t(tok[c]) * GROUPS + group, q, d); + gs[c] += dg.apply(q, d); + us[c] += du.apply(q, d); + } + } +#pragma unroll + for (int c = 0; c < 6; ++c) { +#pragma unroll + for (int o = 8; o; o >>= 1) { + gs[c] += __shfl_xor_sync(0xffffffffU, gs[c], o); + us[c] += __shfl_xor_sync(0xffffffffU, us[c], o); + } + if (c < cnt && lane % 16 == 0) { red[c][0][lrow] = gs[c]; red[c][1][lrow] = us[c]; } + } + __syncthreads(); + if (warp < cnt) { + const int route = s_route[warp]; + float g16 = __half2float(__float2half_rn(red[warp][0][lane])); + float u16 = __half2float(__float2half_rn(red[warp][1][lane])); + g16 = g16 > 10.f ? 10.f : g16; + u16 = u16 > 10.f ? 10.f : (u16 < -10.f ? -10.f : u16); + const float act = __fdiv_rn(g16, __fadd_rn(1.f, expf(-g16))); + const float value = __half2float(__float2half_rn(__fmul_rn(__fmul_rn(act, u16), prob[route]))); + q8_warp_aligned(hq, hd, int64_t(route) * 18 + blockIdx.x, value); + } +} + +template +__global__ void __launch_bounds__(512) gate_up_em6( + int8_t* __restrict__ hq, float* __restrict__ hd, const int8_t* __restrict__ xq, + const float* __restrict__ xd, const int* ids, const float* __restrict__ prob, + const uint8_t* __restrict__ gate, const uint8_t* __restrict__ up, int row_stride, + const uint32_t* __restrict__ gbook, int routes, int top_k) { + constexpr int GROUPS = 160, BW = book_words(); + __shared__ __align__(16) uint32_t book[BW]; + __shared__ int s_expert, s_cnt, s_route[6], s_tok[6]; + __shared__ uint32_t masks[16]; + __shared__ float red[6][2][32]; + const int t = threadIdx.x, warp = t / 32, lane = t % 32; + if (t == 0) { + const int e = find_distinct(ids, routes, blockIdx.y); + int cnt = 0; + if (e >= 0) + for (int r = 0; r < routes; ++r) + if (ids[r] == e && cnt < 6) { s_route[cnt] = r; s_tok[cnt] = r / top_k; ++cnt; } + s_expert = e; + s_cnt = cnt; + } + __syncthreads(); + const int e = s_expert, cnt = s_cnt; + if (e < 0) return; + for (int i = t; i < BW; i += blockDim.x) book[i] = __ldg(gbook + i); + if (t < 16) masks[t] = sign_mask(t); + __syncthreads(); + const int lrow = warp * 2 + lane / 16; + for (int rb = 0; rb < RB; ++rb) { + const int block = blockIdx.x * RB + rb; + const int row = block * 32 + lrow; + const uint8_t* g = gate + (int64_t(e) * 576 + row) * row_stride; + const uint8_t* u = up + (int64_t(e) * 576 + row) * row_stride; + int tok[6]; +#pragma unroll + for (int c = 0; c < 6; ++c) tok[c] = c < cnt ? s_tok[c] : 0; + float gs[6] = {}, us[6] = {}; +#pragma unroll 2 + for (int group = lane % 16; group < GROUPS; group += 16) { + const Group3 dg(g, group, book, masks); + const Group3 du(u, group, book, masks); +#pragma unroll + for (int c = 0; c < 6; ++c) + if (c < cnt) { + int q[8]; + float d; + load_q8(xq, xd, int64_t(tok[c]) * GROUPS + group, q, d); + gs[c] += dg.apply(q, d); + us[c] += du.apply(q, d); + } + } +#pragma unroll + for (int c = 0; c < 6; ++c) { +#pragma unroll + for (int o = 8; o; o >>= 1) { + gs[c] += __shfl_xor_sync(0xffffffffU, gs[c], o); + us[c] += __shfl_xor_sync(0xffffffffU, us[c], o); + } + if (c < cnt && lane % 16 == 0) { red[c][0][lrow] = gs[c]; red[c][1][lrow] = us[c]; } + } + __syncthreads(); + if (warp < cnt) { + const int route = s_route[warp]; + float g16 = __half2float(__float2half_rn(red[warp][0][lane])); + float u16 = __half2float(__float2half_rn(red[warp][1][lane])); + g16 = g16 > 10.f ? 10.f : g16; + u16 = u16 > 10.f ? 10.f : (u16 < -10.f ? -10.f : u16); + const float act = __fdiv_rn(g16, __fadd_rn(1.f, expf(-g16))); + const float value = __half2float(__float2half_rn(__fmul_rn(__fmul_rn(act, u16), prob[route]))); + q8_warp_aligned(hq, hd, int64_t(route) * 18 + block, value); + } + __syncthreads(); + } +} + +template +__global__ void __launch_bounds__(128) down_em5( + half* __restrict__ out, const int8_t* __restrict__ hq, const float* __restrict__ hd, + const int* ids, const uint8_t* __restrict__ down, int row_stride, int left_groups, + const uint32_t* __restrict__ gbook, int routes, int top_k, int tokens, int rows_per_cta) { + constexpr int GROUPS = 18, N = 5120, BW = book_words(); + __shared__ __align__(16) uint32_t book[BW]; + __shared__ int distinct_e[36], s_nu, member_route[36][6], member_cnt[36]; + __shared__ uint32_t masks[16]; + __shared__ float part[4][8][6]; + const int t = threadIdx.x, lane = t % 32, warp = t / 32; + for (int i = t; i < BW; i += blockDim.x) book[i] = __ldg(gbook + i); + if (t < 16) masks[t] = sign_mask(t); + if (t == 0) { + int nu = 0; + for (int r = 0; r < routes; ++r) { + const int e = ids[r]; + int q = 0; + while (q < nu && distinct_e[q] != e) ++q; + if (q == nu) { distinct_e[nu] = e; member_cnt[nu] = 0; ++nu; } + member_route[q][member_cnt[q]++] = r; + } + s_nu = nu; + } + __syncthreads(); + const int nu = s_nu; + const int lrow = lane / 4, sub = lane % 4; + for (int tile = 0; tile < rows_per_cta / 8; ++tile) { + const int row0 = (blockIdx.x * rows_per_cta / 8 + tile) * 8; + const int row = row0 + lrow; + float acc[6] = {}; + for (int q = warp; q < nu; q += 4) { + const int e = distinct_e[q], cnt = member_cnt[q]; + const uint8_t* base = down + (int64_t(e) * N + row) * row_stride; + int rt[6]; +#pragma unroll + for (int c = 0; c < 6; ++c) rt[c] = c < cnt ? member_route[q][c] : 0; + float dot[6] = {}; + for (int g = sub; g < GROUPS; g += 4) { + const Group3 w(base, g + left_groups, book, masks); +#pragma unroll + for (int c = 0; c < 6; ++c) + if (c < cnt) { + int qv[8]; + float d; + load_q8(hq, hd, int64_t(rt[c]) * GROUPS + g, qv, d); + dot[c] += w.apply(qv, d); + } + } +#pragma unroll + for (int c = 0; c < 6; ++c) { + dot[c] += __shfl_xor_sync(0xffffffffU, dot[c], 1); + dot[c] += __shfl_xor_sync(0xffffffffU, dot[c], 2); + } +#pragma unroll + for (int c = 0; c < 6; ++c) + if (c < cnt) { + const int token = rt[c] / top_k; + const float v = __half2float(__float2half_rn(dot[c])); +#pragma unroll + for (int k = 0; k < 6; ++k) + if (k == token) acc[k] += v; + } + } + __syncthreads(); + if (sub == 0) +#pragma unroll + for (int k = 0; k < 6; ++k) part[warp][lrow][k] = acc[k]; + __syncthreads(); + if (t < 8 * tokens) { + const int r = t / tokens, k = t % tokens; + out[k * N + row0 + r] = __float2half_rn(part[0][r][k] + part[1][r][k] + part[2][r][k] + part[3][r][k]); + } + } +} + +template +__global__ void __launch_bounds__(32 * NW, 8 / NW * 4) down_em6( + half* __restrict__ out, const int8_t* __restrict__ hq, const float* __restrict__ hd, + const int* ids, const uint8_t* __restrict__ down, int row_stride, int left_groups, + const uint32_t* __restrict__ gbook, int routes, int top_k, int tokens, int rows_per_cta, + const int* __restrict__ plan) { + constexpr int GROUPS = 18, N = 5120, BW = book_words(); + __shared__ __align__(16) uint32_t book[BW]; + __shared__ int distinct_e[36], s_nu, member_route[36][6], member_cnt[36]; + __shared__ uint32_t masks[16]; + __shared__ float part[NW][8][6]; + const int t = threadIdx.x, lane = t % 32, warp = t / 32; + for (int i = t; i < BW; i += blockDim.x) book[i] = __ldg(gbook + i); + if (t < 16) masks[t] = sign_mask(t); + if (plan != nullptr) { + for (int i = t; i < 36 + 36 + 216; i += blockDim.x) { + if (i < 36) distinct_e[i] = __ldg(plan + 1 + i); + else if (i < 72) member_cnt[i - 36] = __ldg(plan + 1 + i); + else member_route[(i - 72) / 6][(i - 72) % 6] = __ldg(plan + 1 + i); + } + if (t == 0) s_nu = __ldg(plan); + } else if (t == 0) { + int nu = 0; + for (int r = 0; r < routes; ++r) { + const int e = ids[r]; + int q = 0; + while (q < nu && distinct_e[q] != e) ++q; + if (q == nu) { distinct_e[nu] = e; member_cnt[nu] = 0; ++nu; } + member_route[q][member_cnt[q]++] = r; + } + s_nu = nu; + } + __syncthreads(); + const int nu = s_nu; + const int lrow = lane / 4, sub = lane % 4; + for (int tile = 0; tile < rows_per_cta / 8; ++tile) { + const int row0 = (blockIdx.x * rows_per_cta / 8 + tile) * 8; + const int row = row0 + lrow; + float acc[6] = {}; + for (int q = warp; q < nu; q += NW) { + const int e = distinct_e[q], cnt = member_cnt[q]; + const uint8_t* base = down + (int64_t(e) * N + row) * row_stride; + int rt[6]; +#pragma unroll + for (int c = 0; c < 6; ++c) rt[c] = c < cnt ? member_route[q][c] : 0; + float dot[6] = {}; + for (int g = sub; g < GROUPS; g += 4) { + const Group3 w(base, g + left_groups, book, masks); +#pragma unroll + for (int c = 0; c < 6; ++c) + if (c < cnt) { + int qv[8]; + float d; + load_q8(hq, hd, int64_t(rt[c]) * GROUPS + g, qv, d); + dot[c] += w.apply(qv, d); + } + } +#pragma unroll + for (int c = 0; c < 6; ++c) { + dot[c] += __shfl_xor_sync(0xffffffffU, dot[c], 1); + dot[c] += __shfl_xor_sync(0xffffffffU, dot[c], 2); + } +#pragma unroll + for (int c = 0; c < 6; ++c) + if (c < cnt) { + const int token = rt[c] / top_k; + const float v = __half2float(__float2half_rn(dot[c])); +#pragma unroll + for (int k = 0; k < 6; ++k) + if (k == token) acc[k] += v; + } + } + __syncthreads(); + if (sub == 0) +#pragma unroll + for (int k = 0; k < 6; ++k) part[warp][lrow][k] = acc[k]; + __syncthreads(); + if (t < 8 * tokens) { + const int r = t / tokens, k = t % tokens; + float v = 0.f; +#pragma unroll + for (int w8 = 0; w8 < NW; ++w8) v += part[w8][r][k]; + out[k * N + row0 + r] = __float2half_rn(v); + } + } +} + + +// Row-split down (G3): each rank owns 1280 output rows of every expert with the +// full K = 2304 (72 Q8 groups). The routed Q8 inputs of all four ranks are +// gathered rank-major: rank r's block holds hq [routes, 18, 32] then hd +// [routes, 18] at `rank_stride` bytes. Only rows [row_base, row_base + rows) +// of out are written; the caller zeroes the rest before the TP all-reduce. +template +__global__ void __launch_bounds__(32 * NW, 8 / NW * 4) down_rows( + half* __restrict__ out, const uint8_t* __restrict__ gathered, int64_t rank_stride, + const uint8_t* __restrict__ down, int row_stride, int rows_local, int row_base, + const uint32_t* __restrict__ gbook, int routes, int top_k, int tokens, int rows_per_cta, + const int* __restrict__ plan) { + constexpr int LG = 18, GROUPS = 72, N = 5120, BW = book_words(), RT = 32 / SUB; + __shared__ __align__(16) uint32_t book[BW]; + __shared__ int distinct_e[36], s_nu, member_route[36][6], member_cnt[36]; + __shared__ uint32_t masks[16]; + __shared__ float part[NW][8][6]; + const int t = threadIdx.x, lane = t % 32, warp = t / 32; + for (int i = t; i < BW; i += blockDim.x) book[i] = __ldg(gbook + i); + if (t < 16) masks[t] = sign_mask(t); + for (int i = t; i < 36 + 36 + 216; i += blockDim.x) { + if (i < 36) distinct_e[i] = __ldg(plan + 1 + i); + else if (i < 72) member_cnt[i - 36] = __ldg(plan + 1 + i); + else member_route[(i - 72) / 6][(i - 72) % 6] = __ldg(plan + 1 + i); + } + if (t == 0) s_nu = __ldg(plan); + __syncthreads(); + const int nu = s_nu; + const int lrow = lane / SUB, sub = lane % SUB; + const int64_t q_bytes = int64_t(routes) * LG * 32; + for (int tile = 0; tile < rows_per_cta / RT; ++tile) { + const int row0 = (blockIdx.x * rows_per_cta / RT + tile) * RT; + const int row = row0 + lrow; + float acc[6] = {}; + for (int q = warp; q < nu; q += NW) { + const int e = distinct_e[q], cnt = member_cnt[q]; + const uint8_t* base = down + (int64_t(e) * rows_local + row) * row_stride; + int rt[6]; +#pragma unroll + for (int c = 0; c < 6; ++c) rt[c] = c < cnt ? member_route[q][c] : 0; + float dot[6] = {}; + for (int g = sub; g < GROUPS; g += SUB) { + const Group3 w(base, g, book, masks); + const uint8_t* rank_base = gathered + (g / LG) * rank_stride; + const int8_t* hq = reinterpret_cast(rank_base); + const float* hd = reinterpret_cast(rank_base + q_bytes); + const int gl = g % LG; +#pragma unroll + for (int c = 0; c < 6; ++c) + if (c < cnt) { + int qv[8]; + float d; + load_q8(hq, hd, int64_t(rt[c]) * LG + gl, qv, d); + dot[c] += w.apply(qv, d); + } + } +#pragma unroll + for (int c = 0; c < 6; ++c) { +#pragma unroll + for (int o = 1; o < SUB; o <<= 1) dot[c] += __shfl_xor_sync(0xffffffffU, dot[c], o); + } +#pragma unroll + for (int c = 0; c < 6; ++c) + if (c < cnt) { + const int token = rt[c] / top_k; + const float v = __half2float(__float2half_rn(dot[c])); +#pragma unroll + for (int k = 0; k < 6; ++k) + if (k == token) acc[k] += v; + } + } + __syncthreads(); + if (sub == 0) +#pragma unroll + for (int k = 0; k < 6; ++k) part[warp][lrow][k] = acc[k]; + __syncthreads(); + if (t < RT * tokens) { + const int r = t / tokens, k = t % tokens; + float v = 0.f; +#pragma unroll + for (int w8 = 0; w8 < NW; ++w8) v += part[w8][r][k]; + out[k * N + row_base + row0 + r] = __float2half_rn(v); +#pragma unroll + for (int o = 0; o < 4; ++o) + if (o * rows_local != row_base) out[k * N + o * rows_local + row0 + r] = __float2half_rn(0.f); + } + } +} + + +// One-shot TP4 gather over CUDA IPC (G3 decode). Every rank's slot r of every +// rank's buffer receives rank r's `bytes`; flags at `flag_offset` carry a +// per-device epoch (release/acquire at system scope, as in custom all-reduce). +// The MoE output all-reduce that follows bounds rank skew, so one buffer is +// enough: no rank can rewrite a slot before its owner left down_rows. +__device__ __forceinline__ void v41k_st_release(unsigned* p, unsigned v) { + asm volatile("st.release.sys.global.u32 [%0], %1;" ::"l"(p), "r"(v) : "memory"); +} +__device__ __forceinline__ unsigned v41k_ld_acquire(const unsigned* p) { + unsigned v; + asm volatile("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} + +template +__global__ void __launch_bounds__(128 * WPD) peer_gather( + uint8_t* p0, uint8_t* p1, uint8_t* p2, uint8_t* p3, int rank, int64_t stride, + int64_t flag_offset, unsigned* __restrict__ epoch_counter) { + __shared__ unsigned s_epoch; + if (threadIdx.x == 0) s_epoch = *epoch_counter + 1; + __syncthreads(); + const unsigned epoch = s_epoch; + uint8_t* peers[4] = {p0, p1, p2, p3}; + const int dest = threadIdx.x / (32 * WPD), lane = threadIdx.x % (32 * WPD); + if (dest != rank) { + const int4* src = reinterpret_cast(peers[rank] + rank * stride); + int4* dst = reinterpret_cast(peers[dest] + rank * stride); + for (int64_t i = lane; i < stride / 16; i += 32 * WPD) dst[i] = src[i]; + } + __syncthreads(); + if (threadIdx.x < 4) + v41k_st_release(reinterpret_cast(peers[threadIdx.x] + flag_offset) + rank, epoch); + if (threadIdx.x == 0) { + const unsigned* own = reinterpret_cast(peers[rank] + flag_offset); + for (int r = 0; r < 4; ++r) + while (int(v41k_ld_acquire(own + r) - epoch) < 0) { + } + *epoch_counter = epoch; + } +} + + +// ---- G4b decode indexer: scores over packed FP4 index keys ---------------- +// Packed key row (68 B): 64 B E2M1 nibbles (even element in the low nibble) +// then 4 power-of-two scale exponents (biased) for groups of 32. Rows q/w are +// (token, head) pairs padded to 64. Output fp16 out[t, n] = +// sum_h fp16(fp16(relu(fp16(q[t,h] . k[n]))) * w[t,h]), for n < live end. +__device__ __forceinline__ float v41k_e2m1(int code) { + const int m = code & 7; + const float mag = m < 2 ? 0.5f * m : __uint_as_float(uint32_t(((m >> 1) + 126) << 23) | + uint32_t((m & 1) << 22)); + return (code & 8) ? -mag : mag; +} + +template +__global__ void __launch_bounds__(256) index_scores_decode( + const half* __restrict__ q, const half* __restrict__ w, const uint8_t* __restrict__ keys, + const int64_t* __restrict__ positions, half* __restrict__ out, int tokens, int heads, + int64_t n_out, int ratio) { + __shared__ half qs[64][136]; + __shared__ float ks[KB][129]; + __shared__ float ws[64]; + const int t = threadIdx.x; + const int64_t end = (positions[tokens - 1] + 1) / ratio; + const int64_t n0 = int64_t(blockIdx.x) * KB; + if (n0 >= end) return; + const int rows = tokens * heads; + for (int i = t; i < 64 * 128; i += 256) { + const int r = i / 128, d = i % 128; + qs[r][d] = r < rows ? q[r * 128 + d] : __float2half(0.f); + } + if (t < 64) ws[t] = t < rows ? __half2float(w[t]) : 0.f; + for (int i = t; i < KB * 64; i += 256) { + const int k = i / 64, b = i % 64; + const int64_t n = n0 + k; + float lo = 0.f, hi = 0.f; + if (n < end) { + const uint8_t* row = keys + n * 68; + const int byte = row[b]; + const float scale = __uint_as_float(uint32_t(row[64 + (2 * b) / 32]) << 23); + lo = v41k_e2m1(byte & 15) * scale; + hi = v41k_e2m1((byte >> 4) & 15) * scale; + } + // Round like the FP16 cache values (exact for these products). + ks[k][2 * b] = __half2float(__float2half_rn(lo)); + ks[k][2 * b + 1] = __half2float(__float2half_rn(hi)); + } + __syncthreads(); + // Each thread: one key column (k), a stripe of (token) outputs. + for (int item = t; item < KB * tokens; item += 256) { + const int k = item % KB, tok = item / KB; + const int64_t n = n0 + k; + if (n >= end) continue; + float total = 0.f; + for (int h = 0; h < heads; ++h) { + const int r = tok * heads + h; + float dot = 0.f; +#pragma unroll 16 + for (int d = 0; d < 128; ++d) dot = fmaf(__half2float(qs[r][d]), ks[k][d], dot); + float s = __half2float(__float2half_rn(dot)); + s = s > 0.f ? s : 0.f; + total += __half2float(__float2half_rn(s * ws[r])); + } + out[int64_t(tok) * n_out + n] = __float2half_rn(total); + } +} + + +// ---- HC transition (V4.1 mHC, 4 streams, H=5120): parts + finish ---- +// parts: one CTA per 32-column chunk serves every token. thread (o = warp, c = lane) +// owns output stream o at column c; fn is read once per CTA. +template +__global__ void __launch_bounds__(128) hc_parts( + float* __restrict__ residual, const half* __restrict__ y, const float* __restrict__ post, + const float* __restrict__ comb, const float* __restrict__ fn, const float* __restrict__ cpre, + float* __restrict__ cbuf, float* __restrict__ dots, float* __restrict__ squares, + float* __restrict__ csquares) { + constexpr int H = 5120, C = H / 32; + const int chunk = blockIdx.x, o = threadIdx.x / 32, lane = threadIdx.x % 32; + const int col = chunk * 32 + lane; + __shared__ float red[TOK][4][32]; + __shared__ float part[TOK][4][26]; + float r[TOK][4], yv[TOK]; + // issue every load first: all tokens' four sources, y and (once) fn +#pragma unroll + for (int t = 0; t < TOK; ++t) { +#pragma unroll + for (int s = 0; s < 4; ++s) r[t][s] = residual[(int64_t(t) * 4 + s) * H + col]; + yv[t] = __half2float(y[int64_t(t) * H + col]); + } + float w[MIX ? 24 : 1]; + if (MIX) { +#pragma unroll + for (int k = 0; k < 24; ++k) w[k] = __ldg(fn + (int64_t(k) * 4 + o) * H + col); + } + __syncthreads(); // every source row read before any stream is overwritten +#pragma unroll + for (int t = 0; t < TOK; ++t) { + const float* cb = comb + t * 16; + const float n = cb[o] * r[t][0] + cb[4 + o] * r[t][1] + cb[8 + o] * r[t][2] + + cb[12 + o] * r[t][3] + post[t * 4 + o] * yv[t]; + residual[(int64_t(t) * 4 + o) * H + col] = n; + red[t][o][lane] = n * cpre[t * 4 + o]; + if (MIX) { + float acc[25]; +#pragma unroll + for (int k = 0; k < 24; ++k) acc[k] = w[k] * n; + acc[24] = n * n; +#pragma unroll + for (int k = 0; k < 25; ++k) +#pragma unroll + for (int m = 16; m; m >>= 1) acc[k] += __shfl_xor_sync(0xffffffffU, acc[k], m); + if (lane == 0) +#pragma unroll + for (int k = 0; k < 25; ++k) part[t][o][k] = acc[k]; + } + } + __syncthreads(); + for (int t = o; t < TOK; t += 4) { + const float coll = red[t][0][lane] + red[t][1][lane] + red[t][2][lane] + red[t][3][lane]; + cbuf[int64_t(t) * H + col] = coll; + float q = coll * coll; +#pragma unroll + for (int m = 16; m; m >>= 1) q += __shfl_xor_sync(0xffffffffU, q, m); + if (lane == 0) csquares[t * C + chunk] = q; + } + if (MIX) + for (int i = threadIdx.x; i < TOK * 25; i += 128) { + const int t = i / 25, k = i % 25; + const float v = part[t][0][k] + part[t][1][k] + part[t][2][k] + part[t][3][k]; + if (k < 24) dots[(int64_t(t) * C + chunk) * 24 + k] = v; + else squares[t * C + chunk] = v; + } +} + +template +__global__ void __launch_bounds__(256) hc_finish( + const float* __restrict__ dots, const float* __restrict__ squares, + const float* __restrict__ csquares, const float* __restrict__ cbuf, + const float* __restrict__ norm, half* __restrict__ out, const float* __restrict__ scale, + const float* __restrict__ basev, float* __restrict__ pre, float* __restrict__ post, + float* __restrict__ comb, float norm_eps, float hc_eps, int iters) { + constexpr int H = 5120, C = H / 32; + const int t = blockIdx.x, part = blockIdx.y, tid = threadIdx.x, lane = tid % 32, warp = tid / 32; + __shared__ float wsum[8][26]; + __shared__ float s_inv; + // fixed-order reductions over the C chunk partials (independent of M) + float cs = 0.f, sq = 0.f, d[24]; +#pragma unroll + for (int r = 0; r < 24; ++r) d[r] = 0.f; + const bool mixer = MIX && part == 0; + for (int c = tid; c < C; c += 256) { + cs += csquares[t * C + c]; + if (mixer) { + sq += squares[t * C + c]; +#pragma unroll + for (int r = 0; r < 24; ++r) d[r] += dots[(int64_t(t) * C + c) * 24 + r]; + } + } +#pragma unroll + for (int k = 16; k; k >>= 1) { + cs += __shfl_xor_sync(0xffffffffU, cs, k); + if (mixer) { + sq += __shfl_xor_sync(0xffffffffU, sq, k); +#pragma unroll + for (int r = 0; r < 24; ++r) d[r] += __shfl_xor_sync(0xffffffffU, d[r], k); + } + } + if (lane == 0) { + wsum[warp][24] = cs; + if (mixer) { + wsum[warp][25] = sq; +#pragma unroll + for (int r = 0; r < 24; ++r) wsum[warp][r] = d[r]; + } + } + __syncthreads(); + if (tid == 0) { + float total = 0.f; + for (int w8 = 0; w8 < 8; ++w8) total += wsum[w8][24]; + s_inv = rsqrtf(total / H + norm_eps); + } + __syncthreads(); + const float inv = s_inv; + for (int c = part * 1024 + tid; c < (part + 1) * 1024; c += 256) + out[int64_t(t) * H + c] = __float2half_rn(cbuf[int64_t(t) * H + c] * inv * norm[c]); + if (mixer && tid < 32) { + // lane < 16 holds comb[i][j], i = lane / 4, j = lane % 4; lanes 16..23 hold pre/post + float tot = 0.f; + for (int w8 = 0; w8 < 8; ++w8) tot += wsum[w8][25]; + const float rs = rsqrtf(tot / (4 * H) + norm_eps); + float acc = 0.f; + const int r = lane < 16 ? 8 + lane : (lane < 24 ? lane - 16 : 0); + for (int w8 = 0; w8 < 8; ++w8) acc += wsum[w8][r]; + if (lane >= 16 && lane < 20) + pre[t * 4 + lane - 16] = 1.f / (1.f + expf(-(acc * rs * scale[0] + basev[lane - 16]))) + hc_eps; + if (lane >= 20 && lane < 24) + post[t * 4 + lane - 20] = 2.f / (1.f + expf(-(acc * rs * scale[1] + basev[lane - 16]))); + float v = acc * rs * scale[2] + basev[r]; + const unsigned m16 = 0x0000ffffU; + if (lane < 16) { + // softmax over j within row i (4 lanes), then column normalize + float mx = v; + mx = fmaxf(mx, __shfl_xor_sync(m16, mx, 1)); + mx = fmaxf(mx, __shfl_xor_sync(m16, mx, 2)); + v = expf(v - mx); + float rsum = v + __shfl_xor_sync(m16, v, 1); + rsum += __shfl_xor_sync(m16, rsum, 2); + v = v / rsum + hc_eps; + float csum = v + __shfl_xor_sync(m16, v, 4); + csum += __shfl_xor_sync(m16, csum, 8); + v = v / (csum + hc_eps); + for (int it = 1; it < iters; ++it) { + rsum = v + __shfl_xor_sync(m16, v, 1); + rsum += __shfl_xor_sync(m16, rsum, 2); + v = v / (rsum + hc_eps); + csum = v + __shfl_xor_sync(m16, v, 4); + csum += __shfl_xor_sync(m16, csum, 8); + v = v / (csum + hc_eps); + } + comb[t * 16 + lane] = v; + } + } +} + + +// ---- fp16 small-M GEMV: out[m, n] = sum_k x[m, k] * w[n, k] (M <= 8, K % 8 == 0) ---- +template +__global__ void __launch_bounds__(256) hgemv_kernel(half* __restrict__ out, const half* __restrict__ x, + const half* __restrict__ w, int N, int K) { + // 8 warps x 4 rows per CTA; each lane issues all its weight loads first. + constexpr int R = 4, MAXC = 3; // K / 8 <= 96 chunks -> <= 3 per lane + extern __shared__ half xs[]; + for (int i = threadIdx.x; i < M * K / 8; i += blockDim.x) + reinterpret_cast(xs)[i] = reinterpret_cast(x)[i]; + const int lane = threadIdx.x % 32; + const int n0 = (blockIdx.x * 8 + threadIdx.x / 32) * R; + const int chunks = K / 8; + uint4 wv[R][MAXC]; +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int i = 0; i < MAXC; ++i) { + const int c = lane + 32 * i; + wv[r][i] = (n0 + r < N && c < chunks) + ? __ldg(reinterpret_cast(w + int64_t(n0 + r) * K) + c) + : make_uint4(0, 0, 0, 0); + } + __syncthreads(); + float acc[R][M] = {}; +#pragma unroll + for (int i = 0; i < MAXC; ++i) { + const int c = lane + 32 * i; + if (c >= chunks) break; +#pragma unroll + for (int m = 0; m < M; ++m) { + const uint4 xv = reinterpret_cast(xs + m * K)[c]; + const half2* xh = reinterpret_cast(&xv); +#pragma unroll + for (int r = 0; r < R; ++r) { + const half2* wh = reinterpret_cast(&wv[r][i]); +#pragma unroll + for (int j = 0; j < 4; ++j) { + const float2 a = __half22float2(wh[j]), b = __half22float2(xh[j]); + acc[r][m] += a.x * b.x + a.y * b.y; + } + } + } + } +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int m = 0; m < M; ++m) { +#pragma unroll + for (int o = 16; o; o >>= 1) acc[r][m] += __shfl_xor_sync(0xffffffffU, acc[r][m], o); + if (lane == 0 && n0 + r < N) out[int64_t(m) * N + n0 + r] = __float2half_rn(acc[r][m]); + } +} + + +// ---- fused sparse attention over gathered keys (V = K), 16 heads x 512 dims ---- +// grid (T, splits): each CTA takes 64 keys, computes S = Q K^T (fp32), a masked +// partial softmax, and the unnormalized P V partial; attn_combine folds splits + sink. +namespace attnw = nvcuda::wmma; +constexpr int kAH = 16, kAD = 512, kACH = 64; + +__global__ void __launch_bounds__(128) attn_split( + const half* __restrict__ q, const half* __restrict__ keys, const bool* __restrict__ unused, + half* __restrict__ ows, float* __restrict__ mws, float* __restrict__ lws, int width, + float scale) { + constexpr int CH = 32; + __shared__ __align__(32) half qs[kAH * kAD]; // 16 KB + __shared__ __align__(32) float ss[2][kAH * CH]; // partial S per d-half + __shared__ __align__(32) half ps[kAH * CH]; + __shared__ __align__(32) half kz[16 * kAD]; // zero rows for a short last chunk + const int t = blockIdx.x, sp = blockIdx.y, splits = gridDim.y; + const int tid = threadIdx.x, warp = tid / 32; + const int k0 = sp * CH; + const int nk = min(CH, width - k0); + const uint4* qg = reinterpret_cast(q + int64_t(t) * kAH * kAD); + for (int i = tid; i < kAH * kAD / 8; i += 128) reinterpret_cast(qs)[i] = qg[i]; + if (nk < CH) + for (int i = tid; i < 16 * kAD / 8; i += 128) reinterpret_cast(kz)[i] = make_uint4(0, 0, 0, 0); + __syncthreads(); + const half* kg = keys + (int64_t(t) * width + k0) * kAD; + // key tile kt (16 keys): global rows, or zeros when the tile runs past width + auto ktile = [&](int kt) -> const half* { return kt * 16 < nk ? kg + kt * 16 * kAD : kz; }; + { // S: warp -> key tile (warp & 1), d half (warp >> 1) + const int kt = warp & 1, dh = warp >> 1; + attnw::fragment acc; + attnw::fill_fragment(acc, 0.f); + attnw::fragment a; + attnw::fragment b; + const half* kb = ktile(kt); +#pragma unroll 4 + for (int d = dh * 256; d < dh * 256 + 256; d += 16) { + attnw::load_matrix_sync(a, qs + d, kAD); + attnw::load_matrix_sync(b, kb + d, kAD); + attnw::mma_sync(acc, a, b, acc); + } + attnw::store_matrix_sync(ss[dh] + kt * 16, acc, CH, attnw::mem_row_major); + } + __syncthreads(); + { // masked partial softmax: 8 threads per head, 4 keys each + const int h = tid / 8, part = tid % 8; + float v[4], m = -INFINITY; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int key = part * 4 + j; + const bool dead = key >= nk || unused[int64_t(t) * width + k0 + key]; + v[j] = dead ? -INFINITY : (ss[0][h * CH + key] + ss[1][h * CH + key]) * scale; + m = fmaxf(m, v[j]); + } +#pragma unroll + for (int o = 4; o; o >>= 1) m = fmaxf(m, __shfl_xor_sync(0xffffffffU, m, o)); + float l = 0.f; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const float e = m == -INFINITY ? 0.f : expf(v[j] - m); + l += e; + ps[h * CH + part * 4 + j] = __float2half_rn(e); + } +#pragma unroll + for (int o = 4; o; o >>= 1) l += __shfl_xor_sync(0xffffffffU, l, o); + if (part == 0) { + mws[(int64_t(t) * splits + sp) * kAH + h] = m; + lws[(int64_t(t) * splits + sp) * kAH + h] = l; + } + } + __syncthreads(); + { // O partial (fp16): 32 dim tiles, 8 per warp, k = 32 keys + attnw::fragment a0, a1; + attnw::load_matrix_sync(a0, ps, CH); + attnw::load_matrix_sync(a1, ps + 16, CH); + attnw::fragment b; + half* og = ows + (int64_t(t) * splits + sp) * kAH * kAD; + const half* kb0 = ktile(0); + const half* kb1 = ktile(1); +#pragma unroll 2 + for (int tile = warp; tile < kAD / 16; tile += 4) { + attnw::fragment acc; + attnw::fill_fragment(acc, 0.f); + attnw::load_matrix_sync(b, kb0 + tile * 16, kAD); + attnw::mma_sync(acc, a0, b, acc); + attnw::load_matrix_sync(b, kb1 + tile * 16, kAD); + attnw::mma_sync(acc, a1, b, acc); + attnw::fragment hacc; +#pragma unroll + for (int e = 0; e < acc.num_elements; ++e) hacc.x[e] = __float2half_rn(acc.x[e]); + attnw::store_matrix_sync(og + tile * 16, hacc, kAD, attnw::mem_row_major); + } + } +} + +__global__ void __launch_bounds__(256) attn_combine( + const half* __restrict__ ows, const float* __restrict__ mws, const float* __restrict__ lws, + const float* __restrict__ sink, half* __restrict__ out, int splits) { + const int t = blockIdx.x, h = blockIdx.y, tid = threadIdx.x; + float mx = sink[h]; + for (int s = 0; s < splits; ++s) mx = fmaxf(mx, mws[(int64_t(t) * splits + s) * kAH + h]); + float total = expf(sink[h] - mx); + for (int s = 0; s < splits; ++s) { + const float m = mws[(int64_t(t) * splits + s) * kAH + h]; + total += m == -INFINITY ? 0.f : lws[(int64_t(t) * splits + s) * kAH + h] * expf(m - mx); + } + const float inv = 1.f / total; + for (int d = tid; d < kAD; d += 256) { + float acc = 0.f; + for (int s = 0; s < splits; ++s) { + const float m = mws[(int64_t(t) * splits + s) * kAH + h]; + if (m != -INFINITY) + acc += __half2float(ows[((int64_t(t) * splits + s) * kAH + h) * kAD + d]) * expf(m - mx); + } + out[(int64_t(t) * kAH + h) * kAD + d] = __float2half_rn(acc * inv); + } +} + + +// One warp builds the per-layer distinct-expert plan (first-occurrence order), +// replacing the serial thread-0 prologue of every gate/up and down CTA. +__global__ void __launch_bounds__(32) moe_plan(int* __restrict__ plan, const int* __restrict__ ids, + int routes) { + __shared__ int s_ids[64], s_first_u[64]; + const int lane = threadIdx.x; + for (int r = lane; r < 64; r += 32) s_ids[r] = r < routes ? ids[r] : -1; + __syncwarp(); + bool first[2]; + int e[2]; +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int r = lane + 32 * h; + e[h] = s_ids[r]; + bool f = r < routes && e[h] >= 0; // id < 0: route dropped (truncated verify position) + for (int q = 0; q < r && f; ++q) f = s_ids[q] != e[h]; + first[h] = f; + } + const unsigned b0 = __ballot_sync(0xffffffffU, first[0]); + const unsigned b1 = __ballot_sync(0xffffffffU, first[1]); + const unsigned below = (1u << lane) - 1u; +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int r = lane + 32 * h; + const int u = h == 0 ? __popc(b0 & below) : __popc(b0) + __popc(b1 & below); + if (first[h]) s_first_u[r] = u; + } + __syncwarp(); +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int r = lane + 32 * h; + if (r >= routes || e[h] < 0) continue; + int q0 = 0, c = 0; + while (s_ids[q0] != e[h]) ++q0; // first occurrence of this expert + for (int q = q0; q < r; ++q) c += s_ids[q] == e[h]; + const int u = s_first_u[q0]; + if (c < 6) plan[73 + 6 * u + c] = r; + if (first[h]) { + int cnt = 0; + for (int q = r; q < routes; ++q) cnt += s_ids[q] == e[h]; + plan[1 + u] = e[h]; + plan[37 + u] = cnt < 6 ? cnt : 6; + } + } + if (lane == 0) plan[0] = __popc(b0) + __popc(b1); +} + + +// down with all row tiles of the CTA accumulated in registers and one reduction +// barrier at the end (down_em6 synchronized twice per 8-row tile). Same FP order. +template +__global__ void __launch_bounds__(32 * NW, 8 / NW * 4) down_em8( + half* __restrict__ out, const int8_t* __restrict__ hq, const float* __restrict__ hd, + const uint8_t* __restrict__ down, int row_stride, int left_groups, + const uint32_t* __restrict__ gbook, int top_k, int tokens, const int* __restrict__ plan) { + constexpr int GROUPS = 18, N = 5120, BW = book_words(); + __shared__ __align__(16) uint32_t book[BW]; + __shared__ int distinct_e[36], s_nu, member_route[36][6], member_cnt[36]; + __shared__ uint32_t masks[16]; + __shared__ float part[NW][TL][8][6]; + const int t = threadIdx.x, lane = t % 32, warp = t / 32; + for (int i = t; i < BW; i += blockDim.x) book[i] = __ldg(gbook + i); + if (t < 16) masks[t] = sign_mask(t); + for (int i = t; i < 36 + 36 + 216; i += blockDim.x) { + if (i < 36) distinct_e[i] = __ldg(plan + 1 + i); + else if (i < 72) member_cnt[i - 36] = __ldg(plan + 1 + i); + else member_route[(i - 72) / 6][(i - 72) % 6] = __ldg(plan + 1 + i); + } + if (t == 0) s_nu = __ldg(plan); + __syncthreads(); + const int nu = s_nu; + const int lrow = lane / 4, sub = lane % 4; + float acc[TL][6] = {}; + for (int q = warp; q < nu; q += NW) { + const int e = distinct_e[q], cnt = member_cnt[q]; + int rt[6]; +#pragma unroll + for (int c = 0; c < 6; ++c) rt[c] = c < cnt ? member_route[q][c] : 0; +#pragma unroll + for (int tile = 0; tile < TL; ++tile) { + const int row = (blockIdx.x * TL + tile) * 8 + lrow; + const uint8_t* base = down + (int64_t(e) * N + row) * row_stride; + float dot[6] = {}; + for (int g = sub; g < GROUPS; g += 4) { + const Group3 w(base, g + left_groups, book, masks); +#pragma unroll + for (int c = 0; c < 6; ++c) + if (c < cnt) { + int qv[8]; + float d; + load_q8(hq, hd, int64_t(rt[c]) * GROUPS + g, qv, d); + dot[c] += w.apply(qv, d); + } + } +#pragma unroll + for (int c = 0; c < 6; ++c) { + dot[c] += __shfl_xor_sync(0xffffffffU, dot[c], 1); + dot[c] += __shfl_xor_sync(0xffffffffU, dot[c], 2); + } +#pragma unroll + for (int c = 0; c < 6; ++c) + if (c < cnt) { + const int token = rt[c] / top_k; + const float v = __half2float(__float2half_rn(dot[c])); +#pragma unroll + for (int k = 0; k < 6; ++k) + if (k == token) acc[tile][k] += v; + } + } + } + if (sub == 0) +#pragma unroll + for (int tile = 0; tile < TL; ++tile) +#pragma unroll + for (int k = 0; k < 6; ++k) part[warp][tile][lrow][k] = acc[tile][k]; + __syncthreads(); + for (int i = t; i < TL * 8 * tokens; i += 32 * NW) { + const int tile = i / (8 * tokens), r = (i / tokens) % 8, k = i % tokens; + float v = 0.f; +#pragma unroll + for (int w8 = 0; w8 < NW; ++w8) v += part[w8][tile][r][k]; + out[k * N + (blockIdx.x * TL + tile) * 8 + r] = __float2half_rn(v); + } +} + + +// ---- H1 in-place expert swap (shared host arena, TP rank carrier banks) ---- +// Plan contract (from publish): slots[op], loads[op] = incoming's host offset, +// which also receives the victim (in-place exchange). Per shard layout: +// S source row bytes, R source row start, F first byte, D device row bytes, +// W exclusive write bytes, N rows, P projection offset. +struct SwapShard { int64_t S, R, F, D, W, N, P, scratch; }; +struct SwapLayout { SwapShard sh[3]; int64_t rows_total, scratch_per_op; }; + +__device__ __forceinline__ void warp_copy(uint8_t* dst, const uint8_t* src, int64_t len, int lane) { + const uintptr_t a = reinterpret_cast(dst) | reinterpret_cast(src) | uintptr_t(len); + if ((a & 15) == 0) { + for (int64_t i = lane; i < len / 16; i += 32) + reinterpret_cast(dst)[i] = reinterpret_cast(src)[i]; + } else if ((a & 7) == 0) { + for (int64_t i = lane; i < len / 8; i += 32) + reinterpret_cast(dst)[i] = reinterpret_cast(src)[i]; + } else if ((a & 3) == 0) { + for (int64_t i = lane; i < len / 4; i += 32) + reinterpret_cast(dst)[i] = reinterpret_cast(src)[i]; + } else if ((a & 1) == 0) { + for (int64_t i = lane; i < len / 2; i += 32) + reinterpret_cast(dst)[i] = reinterpret_cast(src)[i]; + } else { + for (int64_t i = lane; i < len; i += 32) dst[i] = src[i]; + } +} + +__device__ __forceinline__ void unit_of(const SwapLayout& L, int64_t unit, int64_t& op, int& shard, int64_t& row) { + op = unit / L.rows_total; + row = unit % L.rows_total; + shard = 0; + while (shard < 2 && row >= L.sh[shard].N) { row -= L.sh[shard].N; ++shard; } +} + +__global__ void __launch_bounds__(256) swap_stage_kernel( + uint8_t* b0, uint8_t* b1, uint8_t* b2, const uint8_t* __restrict__ host, uint8_t* scratch, + const int* __restrict__ slots, const int64_t* __restrict__ loads, const int* count, + const int* fast, SwapLayout L) { + if (*fast != 0) return; + const int64_t units = int64_t(*count) * L.rows_total; + const int lane = threadIdx.x % 32; + for (int64_t unit = (int64_t(blockIdx.x) * blockDim.x + threadIdx.x) / 32; unit < units; + unit += int64_t(gridDim.x) * blockDim.x / 32) { + int64_t op, row; int shard; + unit_of(L, unit, op, shard, row); + const SwapShard& g = L.sh[shard]; + uint8_t* bank = shard == 0 ? b0 : (shard == 1 ? b1 : b2); + uint8_t* dev = bank + (int64_t(slots[op]) * g.N + row) * g.D; + // victim's exclusive bytes -> scratch, then incoming host row -> slot + warp_copy(scratch + op * L.scratch_per_op + g.scratch + row * g.W, dev, g.W, lane); + __syncwarp(); + warp_copy(dev, host + loads[op] + (row + g.R) * g.S + g.F + g.P, g.D, lane); + } +} + +__global__ void swap_signal_kernel(unsigned long long* control, const long long* counter, + long long* wb_seq, const int* count, const int* fast, + int flag_word) { + if (*fast != 0 || *count == 0) return; + const long long seq = *counter; + wb_seq[0] = seq; + wb_seq[1] = *count; // private snapshot: the next layer may rewrite count/fast + __threadfence_system(); + atomicExch(control + flag_word, (unsigned long long)seq); + __threadfence_system(); +} + +__global__ void __launch_bounds__(256) swap_writeback_kernel( + uint8_t* host, const uint8_t* __restrict__ scratch, const int64_t* __restrict__ loads, + const int* count, const int* fast, const long long* wb_seq, + unsigned long long* control, int flag0, int flag_stride, int ranks, int ack_word, + unsigned* done, SwapLayout L) { + (void)count; (void)fast; + const long long pending = reinterpret_cast(wb_seq)[1]; + if (pending == 0) return; + const unsigned long long seq = (unsigned long long)wb_seq[0]; + if (threadIdx.x == 0) { + // every rank must finish reading the incoming bytes before any rank + // overwrites them (ffn_down boundary blocks are read by two ranks). + // Only CTA 0 polls host memory; the others wait on a device word. + const long long begin = clock64(); + volatile unsigned* gate = done + 1; + if (blockIdx.x == 0) { + bool all = false; + while (!all) { + all = true; + for (int r = 0; r < ranks; ++r) + all &= *reinterpret_cast(control + flag0 + r * flag_stride) >= seq; + if (!all) __nanosleep(1000); + if (clock64() - begin > (long long)60e9) asm volatile("trap;"); + } + __threadfence(); + *gate = 1; + } else { + while (*gate == 0) { + __nanosleep(500); + if (clock64() - begin > (long long)60e9) asm volatile("trap;"); + } + } + } + __syncthreads(); + const int64_t units = int64_t(pending) * L.rows_total; + const int lane = threadIdx.x % 32; + for (int64_t unit = (int64_t(blockIdx.x) * blockDim.x + threadIdx.x) / 32; unit < units; + unit += int64_t(gridDim.x) * blockDim.x / 32) { + int64_t op, row; int shard; + unit_of(L, unit, op, shard, row); + const SwapShard& g = L.sh[shard]; + warp_copy(host + loads[op] + (row + g.R) * g.S + g.F + g.P, + scratch + op * L.scratch_per_op + g.scratch + row * g.W, g.W, lane); + } + __threadfence_system(); + __syncthreads(); + if (threadIdx.x == 0) { + if (atomicAdd(done, 1u) == gridDim.x - 1) { + *done = 0; + done[1] = 0; + reinterpret_cast(const_cast(wb_seq))[1] = 0; // consumed + __threadfence_system(); + atomicExch(control + ack_word, seq); + __threadfence_system(); + } + } +} + + + +// PCIe-friendly warp copies for arbitrary alignment (zero-copy host memory): +// host reads are whole aligned 16 B words (a superset of the range is read, +// harmless); host writes are aligned 16 B words in the middle and exact bytes +// only at the edges (neighbouring bytes belong to other ranks). +constexpr int kStage = 512; // bytes staged per warp per step + +__device__ __forceinline__ void warp_h2d(uint8_t* dst, const uint8_t* src, int64_t len, + int lane, uint8_t* stage) { + const uintptr_t s = reinterpret_cast(src); + const uintptr_t a0 = s & ~uintptr_t(15), a1 = (s + len + 15) & ~uintptr_t(15); + for (uintptr_t base = a0; base < a1; base += kStage) { + const uintptr_t w = base + uintptr_t(lane) * 16; + if (w < a1) reinterpret_cast(stage)[lane] = *reinterpret_cast(w); + __syncwarp(); + const int64_t lo = int64_t(base > s ? base - s : 0); // dst index of stage[ (s+lo)-base ] + const int64_t hi = min(len, int64_t(base + kStage - s)); + for (int64_t i = lo + lane; i < hi; i += 32) dst[i] = stage[s + i - base]; + __syncwarp(); + } +} + +__device__ __forceinline__ void warp_d2h(uint8_t* dst, const uint8_t* src, int64_t len, + int lane, uint8_t* stage) { + const uintptr_t d = reinterpret_cast(dst); + const uintptr_t a0 = (d + 15) & ~uintptr_t(15), a1 = (d + len) & ~uintptr_t(15); + if (a0 >= a1) { // short: exact bytes + for (int64_t i = lane; i < len; i += 32) dst[i] = src[i]; + return; + } + for (int64_t i = lane; i < int64_t(a0 - d); i += 32) dst[i] = src[i]; // head + for (int64_t i = int64_t(a1 - d) + lane; i < len; i += 32) dst[i] = src[i]; // tail + for (uintptr_t base = a0; base < a1; base += kStage) { + const int64_t off = int64_t(base - d), n = min(int64_t(kStage), int64_t(a1 - base)); + for (int64_t i = lane; i < n; i += 32) stage[i] = src[off + i]; + __syncwarp(); + if (lane * 16 < n) + *reinterpret_cast(base + uintptr_t(lane) * 16) = reinterpret_cast(stage)[lane]; + __syncwarp(); + } +} + +// Free-slot exchange: per row, save the victim's exclusive bytes into its new +// host allocation, then read the incoming row (sources and destinations never +// alias, so no cross-rank ordering is needed). Full-duplex PCIe; ends with a +// system fence so the following commit may publish completion. +__global__ void __launch_bounds__(256) swap_combined_kernel( + uint8_t* b0, uint8_t* b1, uint8_t* b2, uint8_t* host, const int* __restrict__ slots, + const int64_t* __restrict__ loads, const int64_t* __restrict__ saves, const int* count, + const int* fast, SwapLayout L) { + if (*fast != 0) return; + __shared__ __align__(16) uint8_t swap_stage_smem[8 * kStage]; + const int64_t units = int64_t(*count) * L.rows_total; + const int lane = threadIdx.x % 32; + for (int64_t unit = (int64_t(blockIdx.x) * blockDim.x + threadIdx.x) / 32; unit < units; + unit += int64_t(gridDim.x) * blockDim.x / 32) { + int64_t op, row; int shard; + unit_of(L, unit, op, shard, row); + if (loads[op] < 0) continue; // rescue: the expert never left its slot + const SwapShard& g = L.sh[shard]; + uint8_t* bank = shard == 0 ? b0 : (shard == 1 ? b1 : b2); + uint8_t* dev = bank + (int64_t(slots[op]) * g.N + row) * g.D; + const int64_t rowoff = (row + g.R) * g.S + g.F + g.P; + uint8_t* stage = swap_stage_smem + (threadIdx.x / 32) * kStage; + if (saves[op] >= 0) warp_d2h(host + saves[op] + rowoff, dev, g.W, lane, stage); + __syncwarp(); + warp_h2d(dev, host + loads[op] + rowoff, g.D, lane, stage); + } + __threadfence_system(); +} + +} // namespace v41k + +torch::Tensor v41k_table(int64_t type, torch::Device device) { + int words = 0; + if (type == 16) words = turbomind::gemm::LatticeCodebook<16>::kBytes / 4; + else if (type == 17) words = turbomind::gemm::LatticeCodebook<17>::kBytes / 4; + else if (type == 29) words = turbomind::gemm::LatticeCodebook<29>::kBytes / 4; + else TORCH_CHECK(false, "type"); + auto out = torch::empty({words}, torch::dtype(torch::kInt32).device(device)); + auto stream = at::cuda::getCurrentCUDAStream(); + auto* p = reinterpret_cast(out.data_ptr()); + if (type == 16) v41k::build_table<16><<<(words + 255) / 256, 256, 0, stream>>>(p, words); + if (type == 17) v41k::build_table<17><<<(words + 255) / 256, 256, 0, stream>>>(p, words); + if (type == 29) v41k::build_table<29><<<(words + 255) / 256, 256, 0, stream>>>(p, words); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + +void v41k_gate_up(torch::Tensor down_in, torch::Tensor x, torch::Tensor ids32, + torch::Tensor prob, torch::Tensor gate, torch::Tensor up, + int64_t type, torch::Tensor book, int64_t lanes) { + const at::cuda::OptionalCUDAGuard guard(x.device()); + const int routes = ids32.numel(), top_k = 6; + TORCH_CHECK(ids32.scalar_type() == torch::kInt32 && routes <= 36); + const dim3 grid(18, routes); + const size_t smem = 6 * 160 * 9 * 4; + auto stream = at::cuda::getCurrentCUDAStream(); + auto* o = reinterpret_cast(down_in.data_ptr()); + auto* xi = reinterpret_cast(x.data_ptr()); +#define V41K_GU(T, L) \ + { \ + auto k = v41k::gate_up_em; \ + cudaFuncSetAttribute(k, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); \ + k<<>>(o, xi, ids32.data_ptr(), \ + prob.data_ptr(), \ + gate.data_ptr(), \ + up.data_ptr(), gate.size(2), \ + reinterpret_cast( \ + book.data_ptr()), \ + routes, top_k); \ + } + if (lanes == 4) { + if (type == 16) V41K_GU(16, 4) else if (type == 17) V41K_GU(17, 4) else V41K_GU(29, 4) + } else { + TORCH_CHECK(false, "lanes"); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void v41k_down(torch::Tensor out, torch::Tensor h, torch::Tensor ids32, + torch::Tensor down, int64_t type, int64_t left, + torch::Tensor book, torch::Tensor ws, torch::Tensor counters, + int64_t chunks, int64_t warps) { + const at::cuda::OptionalCUDAGuard guard(out.device()); + const int routes = ids32.numel(), tokens = out.size(0); + auto stream = at::cuda::getCurrentCUDAStream(); + auto* hi = reinterpret_cast(h.data_ptr()); +#define V41K_DN(T, W) \ + v41k::down_em<<>>( \ + reinterpret_cast(out.data_ptr()), hi, ids32.data_ptr(), \ + down.data_ptr(), down.size(2), left / 32, \ + reinterpret_cast(book.data_ptr()), routes, 6, tokens, \ + ws.data_ptr(), counters.data_ptr()) + if (warps == 4) { + if (type == 16) V41K_DN(16, 4); else if (type == 17) V41K_DN(17, 4); else V41K_DN(29, 4); + } else if (warps == 2) { + if (type == 16) V41K_DN(16, 2); else if (type == 17) V41K_DN(17, 2); else V41K_DN(29, 2); + } else { + if (type == 16) V41K_DN(16, 1); else if (type == 17) V41K_DN(17, 1); else V41K_DN(29, 1); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + + +void v41k_gate_up2(torch::Tensor down_in, torch::Tensor x, torch::Tensor ids32, + torch::Tensor prob, torch::Tensor gate, torch::Tensor up, + int64_t type, torch::Tensor book) { + const at::cuda::OptionalCUDAGuard guard(x.device()); + const int routes = ids32.numel(); + auto stream = at::cuda::getCurrentCUDAStream(); + auto* o = reinterpret_cast(down_in.data_ptr()); + auto* xi = reinterpret_cast(x.data_ptr()); +#define V41K_GU2(T) \ + v41k::gate_up_em2<<>>( \ + o, xi, ids32.data_ptr(), prob.data_ptr(), \ + gate.data_ptr(), up.data_ptr(), gate.size(2), \ + reinterpret_cast(book.data_ptr()), routes, 6) + if (type == 16) V41K_GU2(16); else if (type == 17) V41K_GU2(17); else V41K_GU2(29); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void v41k_down2(torch::Tensor out, torch::Tensor h, torch::Tensor ids32, + torch::Tensor down, int64_t type, int64_t left, torch::Tensor book) { + const at::cuda::OptionalCUDAGuard guard(out.device()); + const int routes = ids32.numel(), tokens = out.size(0); + auto stream = at::cuda::getCurrentCUDAStream(); + auto* hi = reinterpret_cast(h.data_ptr()); +#define V41K_DN2(T) \ + v41k::down_em2<<<5120 / 8, 128, 0, stream>>>( \ + reinterpret_cast(out.data_ptr()), hi, ids32.data_ptr(), \ + down.data_ptr(), down.size(2), left / 32, \ + reinterpret_cast(book.data_ptr()), routes, 6, tokens) + if (type == 16) V41K_DN2(16); else if (type == 17) V41K_DN2(17); else V41K_DN2(29); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + + +void v41k_diag(torch::Tensor sink, torch::Tensor ids32, torch::Tensor gate, + torch::Tensor up, int64_t type, torch::Tensor book, int64_t mode) { + const at::cuda::OptionalCUDAGuard guard(gate.device()); + const int routes = ids32.numel(); + auto stream = at::cuda::getCurrentCUDAStream(); + const dim3 grid(18, routes); + auto* bk = reinterpret_cast(book.data_ptr()); + if (mode == 0) { + v41k::gu_read_only<<>>(sink.data_ptr(), ids32.data_ptr(), + gate.data_ptr(), up.data_ptr(), gate.size(2), routes); + } else { +#define V41K_DG(T) v41k::gu_decode_only<<>>(sink.data_ptr(), \ + ids32.data_ptr(), gate.data_ptr(), up.data_ptr(), gate.size(2), bk, routes) + if (type == 16) V41K_DG(16); else if (type == 17) V41K_DG(17); else V41K_DG(29); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + + +void v41k_gate_up3(torch::Tensor down_in, torch::Tensor x, torch::Tensor ids32, + torch::Tensor prob, torch::Tensor gate, torch::Tensor up, + int64_t type, torch::Tensor book) { + const at::cuda::OptionalCUDAGuard guard(x.device()); + const int routes = ids32.numel(); + auto stream = at::cuda::getCurrentCUDAStream(); + auto* o = reinterpret_cast(down_in.data_ptr()); + auto* xi = reinterpret_cast(x.data_ptr()); +#define V41K_GU3(T) \ + v41k::gate_up_em3<<>>( \ + o, xi, ids32.data_ptr(), prob.data_ptr(), \ + gate.data_ptr(), up.data_ptr(), gate.size(2), \ + reinterpret_cast(book.data_ptr()), routes, 6) + if (type == 16) V41K_GU3(16); else if (type == 17) V41K_GU3(17); else V41K_GU3(29); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void v41k_down3(torch::Tensor out, torch::Tensor h, torch::Tensor ids32, + torch::Tensor down, int64_t type, int64_t left, torch::Tensor book, + int64_t rows_per_cta) { + const at::cuda::OptionalCUDAGuard guard(out.device()); + const int routes = ids32.numel(), tokens = out.size(0); + auto stream = at::cuda::getCurrentCUDAStream(); + auto* hi = reinterpret_cast(h.data_ptr()); + TORCH_CHECK(rows_per_cta % 8 == 0 && 5120 % rows_per_cta == 0); +#define V41K_DN3(T) \ + v41k::down_em3<<<5120 / rows_per_cta, 128, 0, stream>>>( \ + reinterpret_cast(out.data_ptr()), hi, ids32.data_ptr(), \ + down.data_ptr(), down.size(2), left / 32, \ + reinterpret_cast(book.data_ptr()), routes, 6, tokens, \ + int(rows_per_cta)) + if (type == 16) V41K_DN3(16); else if (type == 17) V41K_DN3(17); else V41K_DN3(29); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void v41k_decode3(torch::Tensor sink, torch::Tensor ids32, torch::Tensor gate, + torch::Tensor up, int64_t type, torch::Tensor book) { + const at::cuda::OptionalCUDAGuard guard(gate.device()); + const int routes = ids32.numel(); + auto stream = at::cuda::getCurrentCUDAStream(); + auto* bk = reinterpret_cast(book.data_ptr()); +#define V41K_D3(T) v41k::gu_decode3<<>>( \ + sink.data_ptr(), ids32.data_ptr(), gate.data_ptr(), \ + up.data_ptr(), gate.size(2), bk, routes) + if (type == 16) V41K_D3(16); else if (type == 17) V41K_D3(17); else V41K_D3(29); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + + +void v41k_gate_up4(torch::Tensor down_in, torch::Tensor x, torch::Tensor ids32, + torch::Tensor prob, torch::Tensor gate, torch::Tensor up, + int64_t type, torch::Tensor book) { + const at::cuda::OptionalCUDAGuard guard(x.device()); + const int routes = ids32.numel(), tokens = x.size(0); + auto stream = at::cuda::getCurrentCUDAStream(); + auto* o = reinterpret_cast(down_in.data_ptr()); + auto* xi = reinterpret_cast(x.data_ptr()); +#define V41K_GU4(T) \ + { \ + const size_t smem = 6 * 160 * 36 + v41k::book_words() * 4; \ + auto k = v41k::gate_up_em4; \ + static bool configured = false; \ + if (!configured) { \ + cudaFuncSetAttribute(k, cudaFuncAttributeMaxDynamicSharedMemorySize, smem);\ + configured = true; \ + } \ + k<<>>( \ + o, xi, ids32.data_ptr(), prob.data_ptr(), \ + gate.data_ptr(), up.data_ptr(), gate.size(2), \ + reinterpret_cast(book.data_ptr()), routes, 6, tokens);\ + } + if (type == 16) V41K_GU4(16) else if (type == 17) V41K_GU4(17) else V41K_GU4(29) + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void v41k_down4(torch::Tensor out, torch::Tensor h, torch::Tensor ids32, + torch::Tensor down, int64_t type, int64_t left, torch::Tensor book, + int64_t rows_per_cta) { + const at::cuda::OptionalCUDAGuard guard(out.device()); + const int routes = ids32.numel(), tokens = out.size(0); + auto stream = at::cuda::getCurrentCUDAStream(); + auto* hi = reinterpret_cast(h.data_ptr()); + TORCH_CHECK(rows_per_cta % 8 == 0 && 5120 % rows_per_cta == 0 && routes <= 36); +#define V41K_DN4(T) \ + v41k::down_em4<<<5120 / rows_per_cta, 128, 0, stream>>>( \ + reinterpret_cast(out.data_ptr()), hi, ids32.data_ptr(), \ + down.data_ptr(), down.size(2), left / 32, \ + reinterpret_cast(book.data_ptr()), routes, 6, tokens, \ + int(rows_per_cta)) + if (type == 16) V41K_DN4(16); else if (type == 17) V41K_DN4(17); else V41K_DN4(29); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + + +void v41k_fp8_gemv(torch::Tensor out, torch::Tensor x, torch::Tensor w, + torch::Tensor scales, torch::Tensor partials, int64_t split, + c10::optional bias) { + const at::cuda::OptionalCUDAGuard guard(x.device()); + const int M = x.size(0), K = x.size(1), N = w.size(0); + TORCH_CHECK(M >= 1 && M <= 8 && K % 32 == 0 && split >= 1); + int ksplit = (K + split - 1) / split; + ksplit = (ksplit + 15) / 16 * 16; + auto stream = at::cuda::getCurrentCUDAStream(); + const dim3 grid((N + 31) / 32, split); + const bool direct = split == 1 && !bias.has_value(); + float* part = direct ? nullptr : partials.data_ptr(); + half* o = reinterpret_cast(out.data_ptr()); + auto* xi = reinterpret_cast(x.data_ptr()); + auto* wi = reinterpret_cast(w.data_ptr()); + auto* si = reinterpret_cast(scales.data_ptr()); +#define V41K_FP8(MM) v41k::fp8_gemv_kernel<<>>( \ + part, o, xi, wi, si, N, K, direct ? 1 : 0, ksplit) + switch (M) { + case 1: V41K_FP8(1); break; case 2: V41K_FP8(2); break; case 3: V41K_FP8(3); break; + case 4: V41K_FP8(4); break; case 5: V41K_FP8(5); break; case 6: V41K_FP8(6); break; + case 7: V41K_FP8(7); break; default: V41K_FP8(8); break; + } + if (!direct) { + const int total = M * N; + v41k::fp8_reduce<<<(total + 255) / 256, 256, 0, stream>>>( + o, part, bias.has_value() ? reinterpret_cast(bias->data_ptr()) : nullptr, + M, N, split); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + + +void v41k_q8_aligned(torch::Tensor qs, torch::Tensor ds, torch::Tensor x) { + const at::cuda::OptionalCUDAGuard guard(x.device()); + const int M = x.size(0), K = x.size(1); + auto stream = at::cuda::getCurrentCUDAStream(); + v41k::q8_aligned_kernel<<>>( + reinterpret_cast(qs.data_ptr()), ds.data_ptr(), + reinterpret_cast(x.data_ptr()), K); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void v41k_moe5(torch::Tensor out, torch::Tensor xq, torch::Tensor xd, torch::Tensor hq, + torch::Tensor hd, torch::Tensor ids32, torch::Tensor prob, torch::Tensor gate, + torch::Tensor up, torch::Tensor down, int64_t gtype, int64_t dtype, + int64_t left, torch::Tensor gbook, torch::Tensor dbook, int64_t rows) { + const at::cuda::OptionalCUDAGuard guard(out.device()); + const int routes = ids32.numel(), tokens = out.size(0); + auto stream = at::cuda::getCurrentCUDAStream(); + auto i8 = [](torch::Tensor& t) { return reinterpret_cast(t.data_ptr()); }; +#define V41K_GU5(T) v41k::gate_up_em5<<>>( \ + i8(hq), hd.data_ptr(), i8(xq), xd.data_ptr(), ids32.data_ptr(), \ + prob.data_ptr(), gate.data_ptr(), up.data_ptr(), gate.size(2), \ + reinterpret_cast(gbook.data_ptr()), routes, 6, nullptr) + if (gtype == 16) V41K_GU5(16); else if (gtype == 17) V41K_GU5(17); else V41K_GU5(29); +#define V41K_DN5(T) v41k::down_em5<<<5120 / rows, 128, 0, stream>>>( \ + reinterpret_cast(out.data_ptr()), i8(hq), hd.data_ptr(), \ + ids32.data_ptr(), down.data_ptr(), down.size(2), left / 32, \ + reinterpret_cast(dbook.data_ptr()), routes, 6, tokens, int(rows)) + if (dtype == 16) V41K_DN5(16); else if (dtype == 17) V41K_DN5(17); else V41K_DN5(29); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + + +void v41k_moe6(torch::Tensor out, torch::Tensor xq, torch::Tensor xd, torch::Tensor hq, + torch::Tensor hd, torch::Tensor ids32, torch::Tensor prob, torch::Tensor gate, + torch::Tensor up, torch::Tensor down, int64_t gtype, int64_t dtype, + int64_t left, torch::Tensor gbook, torch::Tensor dbook, int64_t rows, + int64_t rb) { + const at::cuda::OptionalCUDAGuard guard(out.device()); + const int routes = ids32.numel(), tokens = out.size(0); + auto stream = at::cuda::getCurrentCUDAStream(); + auto i8 = [](torch::Tensor& t) { return reinterpret_cast(t.data_ptr()); }; +#define V41K_GU6(T, R) v41k::gate_up_em6<<>>( \ + i8(hq), hd.data_ptr(), i8(xq), xd.data_ptr(), ids32.data_ptr(), \ + prob.data_ptr(), gate.data_ptr(), up.data_ptr(), gate.size(2), \ + reinterpret_cast(gbook.data_ptr()), routes, 6) +#define V41K_GU6T(T) if (rb == 1) V41K_GU6(T, 1); else if (rb == 2) V41K_GU6(T, 2); \ + else if (rb == 3) V41K_GU6(T, 3); else if (rb == 6) V41K_GU6(T, 6); else V41K_GU6(T, 9) + if (gtype == 16) { V41K_GU6T(16); } else if (gtype == 17) { V41K_GU6T(17); } else { V41K_GU6T(29); } + if (dtype == 16) V41K_DN5(16); else if (dtype == 17) V41K_DN5(17); else V41K_DN5(29); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + + +void v41k_hc_transition(torch::Tensor residual, torch::Tensor y, torch::Tensor post, + torch::Tensor comb, torch::Tensor cpre, torch::Tensor norm, + torch::Tensor out, torch::Tensor cbuf, torch::Tensor csq, + std::optional fn, std::optional scale, + std::optional base, std::optional dots, + std::optional sq, std::optional pre_n, + std::optional post_n, std::optional comb_n, + double norm_eps, double hc_eps, int64_t iters) { + const at::cuda::OptionalCUDAGuard guard(residual.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + const int tokens = residual.size(0); + const bool mix = fn.has_value(); + auto f = [](const std::optional& t) { return t ? t->data_ptr() : nullptr; }; + const half* yp = reinterpret_cast(y.data_ptr()); + if (mix) { +#define V41K_HCP(M, T) v41k::hc_parts<<<160, 128, 0, stream>>>( \ + residual.data_ptr(), yp, post.data_ptr(), comb.data_ptr(), \ + M ? f(fn) : nullptr, cpre.data_ptr(), cbuf.data_ptr(), \ + M ? f(dots) : nullptr, M ? f(sq) : nullptr, csq.data_ptr()) +#define V41K_HCPT(M) switch (tokens) { case 1: V41K_HCP(M, 1); break; case 2: V41K_HCP(M, 2); break; \ + case 3: V41K_HCP(M, 3); break; case 4: V41K_HCP(M, 4); break; case 5: V41K_HCP(M, 5); break; \ + case 6: V41K_HCP(M, 6); break; case 7: V41K_HCP(M, 7); break; default: V41K_HCP(M, 8); } + V41K_HCPT(true) + v41k::hc_finish<<>>( + f(dots), f(sq), csq.data_ptr(), cbuf.data_ptr(), norm.data_ptr(), + reinterpret_cast(out.data_ptr()), f(scale), f(base), f(pre_n), f(post_n), + f(comb_n), float(norm_eps), float(hc_eps), int(iters)); + } else { + V41K_HCPT(false) + v41k::hc_finish<<>>( + nullptr, nullptr, csq.data_ptr(), cbuf.data_ptr(), norm.data_ptr(), + reinterpret_cast(out.data_ptr()), nullptr, nullptr, nullptr, nullptr, nullptr, + float(norm_eps), float(hc_eps), int(iters)); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + + +void v41k_moe5w(torch::Tensor out, torch::Tensor xq, torch::Tensor xd, torch::Tensor hq, + torch::Tensor hd, torch::Tensor ids32, torch::Tensor prob, torch::Tensor gate, + torch::Tensor up, torch::Tensor down, int64_t gtype, int64_t dtype, + int64_t left, torch::Tensor gbook, torch::Tensor dbook, int64_t rows, + int64_t dn_warps) { + if (dn_warps == 4) { + v41k_moe5(out, xq, xd, hq, hd, ids32, prob, gate, up, down, gtype, dtype, left, gbook, + dbook, rows); + return; + } + const at::cuda::OptionalCUDAGuard guard(out.device()); + const int routes = ids32.numel(), tokens = out.size(0); + auto stream = at::cuda::getCurrentCUDAStream(); + auto i8 = [](torch::Tensor& t) { return reinterpret_cast(t.data_ptr()); }; + if (gtype == 16) V41K_GU5(16); else if (gtype == 17) V41K_GU5(17); else V41K_GU5(29); +#define V41K_DN6(T, W) v41k::down_em6<<<5120 / rows, 32 * W, 0, stream>>>( \ + reinterpret_cast(out.data_ptr()), i8(hq), hd.data_ptr(), \ + ids32.data_ptr(), down.data_ptr(), down.size(2), left / 32, \ + reinterpret_cast(dbook.data_ptr()), routes, 6, tokens, int(rows), nullptr) + if (dn_warps == 8) { + if (dtype == 16) V41K_DN6(16, 8); else if (dtype == 17) V41K_DN6(17, 8); else V41K_DN6(29, 8); + } else { + if (dtype == 16) V41K_DN6(16, 2); else if (dtype == 17) V41K_DN6(17, 2); else V41K_DN6(29, 2); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + + +void v41k_hgemv(torch::Tensor out, torch::Tensor x, torch::Tensor w) { + const at::cuda::OptionalCUDAGuard guard(x.device()); + const int M = x.size(0), K = x.size(1), N = w.size(0); + TORCH_CHECK(M >= 1 && M <= 8 && K % 8 == 0 && K <= 768 && x.is_contiguous() && w.is_contiguous()); + auto stream = at::cuda::getCurrentCUDAStream(); + const size_t smem = size_t(M) * K * 2; + auto* o = reinterpret_cast(out.data_ptr()); + auto* xp = reinterpret_cast(x.data_ptr()); + auto* wp = reinterpret_cast(w.data_ptr()); + const dim3 grid((N + 31) / 32); + switch (M) { +#define V41K_HG(MM) case MM: v41k::hgemv_kernel<<>>(o, xp, wp, N, K); break; + V41K_HG(1) V41K_HG(2) V41K_HG(3) V41K_HG(4) V41K_HG(5) V41K_HG(6) V41K_HG(7) V41K_HG(8) + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + + +void v41k_attn_gathered(torch::Tensor out, torch::Tensor q, torch::Tensor keys, + torch::Tensor unused, torch::Tensor sink, torch::Tensor ws, + double scale) { + const at::cuda::OptionalCUDAGuard guard(q.device()); + const int T = q.size(0), width = keys.size(1); + TORCH_CHECK(q.size(1) == 16 && q.size(2) == 512 && keys.size(2) == 512); + TORCH_CHECK(q.is_contiguous() && keys.is_contiguous() && unused.is_contiguous()); + const int splits = (width + 31) / 32; + // ws (fp32 elements): O partials as fp16 [T, splits, 16, 512], then m and l + const int64_t o_floats = int64_t(T) * splits * 16 * 512 / 2; + TORCH_CHECK(ws.numel() >= o_floats + int64_t(T) * splits * 16 * 2); + half* ows = reinterpret_cast(ws.data_ptr()); + float* mws = ws.data_ptr() + o_floats; + float* lws = mws + int64_t(T) * splits * 16; + auto stream = at::cuda::getCurrentCUDAStream(); + v41k::attn_split<<>>( + reinterpret_cast(q.data_ptr()), reinterpret_cast(keys.data_ptr()), + unused.data_ptr(), ows, mws, lws, width, float(scale)); + v41k::attn_combine<<>>(ows, mws, lws, sink.data_ptr(), + reinterpret_cast(out.data_ptr()), + splits); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + + +void v41k_moe7(torch::Tensor out, torch::Tensor xq, torch::Tensor xd, torch::Tensor hq, + torch::Tensor hd, torch::Tensor ids32, torch::Tensor prob, torch::Tensor gate, + torch::Tensor up, torch::Tensor down, int64_t gtype, int64_t dtype, + int64_t left, torch::Tensor gbook, torch::Tensor dbook, int64_t rows, + int64_t dn_warps, torch::Tensor plan) { + const at::cuda::OptionalCUDAGuard guard(out.device()); + const int routes = ids32.numel(), tokens = out.size(0); + TORCH_CHECK(routes <= 36 && plan.numel() >= 73 + 216); + auto stream = at::cuda::getCurrentCUDAStream(); + auto i8 = [](torch::Tensor& t) { return reinterpret_cast(t.data_ptr()); }; + int* pl = plan.data_ptr(); + v41k::moe_plan<<<1, 32, 0, stream>>>(pl, ids32.data_ptr(), routes); +#define V41K_GU7(T) v41k::gate_up_em5<<>>( \ + i8(hq), hd.data_ptr(), i8(xq), xd.data_ptr(), ids32.data_ptr(), \ + prob.data_ptr(), gate.data_ptr(), up.data_ptr(), gate.size(2), \ + reinterpret_cast(gbook.data_ptr()), routes, 6, pl) + if (gtype == 16) V41K_GU7(16); else if (gtype == 17) V41K_GU7(17); else V41K_GU7(29); +#define V41K_DN7(T, W) v41k::down_em6<<<5120 / rows, 32 * W, 0, stream>>>( \ + reinterpret_cast(out.data_ptr()), i8(hq), hd.data_ptr(), \ + ids32.data_ptr(), down.data_ptr(), down.size(2), left / 32, \ + reinterpret_cast(dbook.data_ptr()), routes, 6, tokens, int(rows), pl) + static const bool em8 = [] { const char* v = std::getenv("VLLM_V41_EM8"); return v == nullptr || v[0] != '0'; }(); + if (em8) { +#define V41K_DN8(T, W, L) v41k::down_em8<<<5120 / (8 * L), 32 * W, 0, stream>>>( \ + reinterpret_cast(out.data_ptr()), i8(hq), hd.data_ptr(), \ + down.data_ptr(), down.size(2), left / 32, \ + reinterpret_cast(dbook.data_ptr()), 6, tokens, pl) +#define V41K_DN8T(T) if (dn_warps == 8 && rows == 16) V41K_DN8(T, 8, 2); \ + else if (dn_warps == 8 && rows == 32) V41K_DN8(T, 8, 4); \ + else if (dn_warps == 4 && rows == 16) V41K_DN8(T, 4, 2); \ + else if (dn_warps == 8) V41K_DN8(T, 8, 1); else V41K_DN8(T, 4, 1) + if (dtype == 16) { V41K_DN8T(16); } else if (dtype == 17) { V41K_DN8T(17); } else { V41K_DN8T(29); } + } else if (dn_warps == 8) { + if (dtype == 16) V41K_DN7(16, 8); else if (dtype == 17) V41K_DN7(17, 8); else V41K_DN7(29, 8); + } else { + if (dtype == 16) V41K_DN7(16, 4); else if (dtype == 17) V41K_DN7(17, 4); else V41K_DN7(29, 4); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + + +static v41k::SwapLayout v41k_swap_layout(const std::vector& v) { + TORCH_CHECK(v.size() == 21, "swap layout: 3 x (S, R, F, D, W, N, P)"); + v41k::SwapLayout L{}; + int64_t scratch = 0, rows = 0; + for (int i = 0; i < 3; ++i) { + auto& g = L.sh[i]; + g.S = v[7 * i]; g.R = v[7 * i + 1]; g.F = v[7 * i + 2]; g.D = v[7 * i + 3]; + g.W = v[7 * i + 4]; g.N = v[7 * i + 5]; g.P = v[7 * i + 6]; + g.scratch = scratch; + scratch += g.N * g.W; + rows += g.N; + } + L.rows_total = rows; + L.scratch_per_op = (scratch + 255) / 256 * 256; + return L; +} + +void v41k_swap_stage(torch::Tensor b0, torch::Tensor b1, torch::Tensor b2, torch::Tensor host, + torch::Tensor scratch, torch::Tensor slots, torch::Tensor loads, + torch::Tensor count, torch::Tensor fast, std::vector layout) { + const at::cuda::OptionalCUDAGuard guard(b0.device()); + auto L = v41k_swap_layout(layout); + TORCH_CHECK(scratch.numel() >= L.scratch_per_op * slots.numel(), "swap scratch too small"); + auto stream = at::cuda::getCurrentCUDAStream(); + v41k::swap_stage_kernel<<<1024, 256, 0, stream>>>( + b0.data_ptr(), b1.data_ptr(), b2.data_ptr(), + host.data_ptr(), scratch.data_ptr(), slots.data_ptr(), + loads.data_ptr(), count.data_ptr(), fast.data_ptr(), L); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void v41k_swap_signal(torch::Tensor control, torch::Tensor counter, torch::Tensor wb_seq, + torch::Tensor count, torch::Tensor fast, int64_t flag_word) { + const at::cuda::OptionalCUDAGuard guard(counter.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + v41k::swap_signal_kernel<<<1, 1, 0, stream>>>( + reinterpret_cast(control.data_ptr()), + reinterpret_cast(counter.data_ptr()), + reinterpret_cast(wb_seq.data_ptr()), count.data_ptr(), + fast.data_ptr(), int(flag_word)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void v41k_swap_writeback(torch::Tensor host, torch::Tensor scratch, torch::Tensor loads, + torch::Tensor count, torch::Tensor fast, torch::Tensor wb_seq, + torch::Tensor control, int64_t flag0, int64_t flag_stride, + int64_t ranks, int64_t ack_word, torch::Tensor done, + std::vector layout, int64_t ctas) { + const at::cuda::OptionalCUDAGuard guard(scratch.device()); + auto L = v41k_swap_layout(layout); + auto stream = at::cuda::getCurrentCUDAStream(); + TORCH_CHECK(done.numel() >= 2, "swap_writeback: done needs 2 words"); + v41k::swap_writeback_kernel<<>>( + host.data_ptr(), scratch.data_ptr(), loads.data_ptr(), + count.data_ptr(), fast.data_ptr(), + reinterpret_cast(wb_seq.data_ptr()), + reinterpret_cast(control.data_ptr()), int(flag0), + int(flag_stride), int(ranks), int(ack_word), + reinterpret_cast(done.data_ptr()), L); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +int64_t v41k_swap_scratch_per_op(std::vector layout) { + return v41k_swap_layout(layout).scratch_per_op; +} + + +void v41k_swap_combined(torch::Tensor b0, torch::Tensor b1, torch::Tensor b2, torch::Tensor host, + torch::Tensor slots, torch::Tensor loads, torch::Tensor saves, + torch::Tensor count, torch::Tensor fast, std::vector layout, + int64_t ctas) { + const at::cuda::OptionalCUDAGuard guard(b0.device()); + auto L = v41k_swap_layout(layout); + auto stream = at::cuda::getCurrentCUDAStream(); + v41k::swap_combined_kernel<<>>( + b0.data_ptr(), b1.data_ptr(), b2.data_ptr(), + host.data_ptr(), slots.data_ptr(), loads.data_ptr(), + saves.data_ptr(), count.data_ptr(), fast.data_ptr(), L); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void v41k_gate_up7(torch::Tensor xq, torch::Tensor xd, torch::Tensor hq, torch::Tensor hd, + torch::Tensor ids32, torch::Tensor prob, torch::Tensor gate, torch::Tensor up, + int64_t gtype, torch::Tensor gbook, torch::Tensor plan) { + const at::cuda::OptionalCUDAGuard guard(xq.device()); + const int routes = ids32.numel(); + TORCH_CHECK(routes <= 36 && plan.numel() >= 73 + 216); + auto stream = at::cuda::getCurrentCUDAStream(); + auto i8 = [](torch::Tensor& t) { return reinterpret_cast(t.data_ptr()); }; + int* pl = plan.data_ptr(); + v41k::moe_plan<<<1, 32, 0, stream>>>(pl, ids32.data_ptr(), routes); +#define V41K_GU7R(T) v41k::gate_up_em5<<>>( \ + i8(hq), hd.data_ptr(), i8(xq), xd.data_ptr(), ids32.data_ptr(), \ + prob.data_ptr(), gate.data_ptr(), up.data_ptr(), gate.size(2), \ + reinterpret_cast(gbook.data_ptr()), routes, 6, pl) + if (gtype == 16) V41K_GU7R(16); else if (gtype == 17) V41K_GU7R(17); else V41K_GU7R(29); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void v41k_down_rows(torch::Tensor out, torch::Tensor gathered, int64_t rank_stride, + int64_t routes, torch::Tensor down, int64_t dtype, int64_t row_base, + torch::Tensor dbook, int64_t rows, int64_t dn_warps, torch::Tensor plan) { + const at::cuda::OptionalCUDAGuard guard(out.device()); + const int tokens = out.size(0), rows_local = down.size(1); + TORCH_CHECK(routes <= 36 && rows_local % rows == 0 && rows % 2 == 0 && rank_stride % 16 == 0); + TORCH_CHECK(gathered.numel() >= 4 * rank_stride && out.size(1) == 5120); + auto stream = at::cuda::getCurrentCUDAStream(); + const int sub = rows % 4 ? 16 : (rows % 8 ? 8 : 4); +#define V41K_DNR(T, W) do { if (sub == 16) V41K_DNR2(T, W, 16) else if (sub == 8) V41K_DNR2(T, W, 8) else V41K_DNR2(T, W, 4) } while (0) +#define V41K_DNR2(T, W, S) v41k::down_rows<<>>( \ + reinterpret_cast(out.data_ptr()), gathered.data_ptr(), rank_stride, \ + down.data_ptr(), down.size(2), rows_local, int(row_base), \ + reinterpret_cast(dbook.data_ptr()), int(routes), 6, tokens, int(rows), \ + plan.data_ptr()); + if (dn_warps == 8) { + if (dtype == 16) V41K_DNR(16, 8); else if (dtype == 17) V41K_DNR(17, 8); else V41K_DNR(29, 8); + } else { + if (dtype == 16) V41K_DNR(16, 4); else if (dtype == 17) V41K_DNR(17, 4); else V41K_DNR(29, 4); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + + +void v41k_peer_gather(int64_t p0, int64_t p1, int64_t p2, int64_t p3, int64_t rank, + int64_t stride, int64_t flag_offset, torch::Tensor epoch) { + const at::cuda::OptionalCUDAGuard guard(epoch.device()); + TORCH_CHECK(stride % 16 == 0 && flag_offset % 16 == 0 && rank >= 0 && rank < 4); + auto stream = at::cuda::getCurrentCUDAStream(); + auto ptr = [](int64_t v) { return reinterpret_cast(v); }; + static const int wpd = [] { const char* v = std::getenv("VLLM_V41_GATHER_WPD"); return v ? std::atoi(v) : 4; }(); + unsigned* ep = reinterpret_cast(epoch.data_ptr()); + if (wpd == 1) v41k::peer_gather<1><<<1, 128, 0, stream>>>(ptr(p0), ptr(p1), ptr(p2), ptr(p3), int(rank), stride, flag_offset, ep); + else if (wpd == 2) v41k::peer_gather<2><<<1, 256, 0, stream>>>(ptr(p0), ptr(p1), ptr(p2), ptr(p3), int(rank), stride, flag_offset, ep); + else v41k::peer_gather<4><<<1, 512, 0, stream>>>(ptr(p0), ptr(p1), ptr(p2), ptr(p3), int(rank), stride, flag_offset, ep); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + + +void v41k_index_scores_decode(torch::Tensor out, torch::Tensor q, torch::Tensor w, + torch::Tensor keys, torch::Tensor positions, int64_t heads, + int64_t ratio) { + const at::cuda::OptionalCUDAGuard guard(out.device()); + const int tokens = out.size(0); + TORCH_CHECK(tokens * heads <= 64 && q.size(-1) == 128 && keys.size(1) == 68 && + positions.scalar_type() == torch::kInt64); + auto stream = at::cuda::getCurrentCUDAStream(); + constexpr int KB = 32; + const int blocks = (keys.size(0) + KB - 1) / KB; + v41k::index_scores_decode<<>>( + reinterpret_cast(q.data_ptr()), reinterpret_cast(w.data_ptr()), + keys.data_ptr(), positions.data_ptr(), + reinterpret_cast(out.data_ptr()), tokens, int(heads), out.size(1), int(ratio)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + + +TORCH_LIBRARY(v41k, m) { + m.def("index_scores_decode(Tensor(a!) out, Tensor q, Tensor w, Tensor keys, Tensor positions, " + "int heads, int ratio) -> ()", &v41k_index_scores_decode); + m.def("peer_gather(int p0, int p1, int p2, int p3, int rank, int stride, int flag_offset, " + "Tensor(a!) epoch) -> ()", &v41k_peer_gather); + m.def("gate_up7(Tensor xq, Tensor xd, Tensor(a!) hq, Tensor(b!) hd, Tensor ids32, Tensor prob, " + "Tensor gate, Tensor up, int gtype, Tensor gbook, Tensor(c!) plan) -> ()", &v41k_gate_up7); + m.def("down_rows(Tensor(a!) out, Tensor gathered, int rank_stride, int routes, Tensor down, " + "int dtype, int row_base, Tensor dbook, int rows, int dn_warps, Tensor plan) -> ()", + &v41k_down_rows); + m.def("table(int type, Device device) -> Tensor", &v41k_table); + m.def("gate_up(Tensor(a!) down_in, Tensor x, Tensor ids32, Tensor prob, " + "Tensor gate, Tensor up, int type, Tensor book, int lanes) -> ()", + &v41k_gate_up); + m.def("down(Tensor(a!) out, Tensor h, Tensor ids32, Tensor down, int type, " + "int left, Tensor book, Tensor ws, Tensor counters, int chunks, " + "int warps) -> ()", + &v41k_down); + m.def("diag(Tensor(a!) sink, Tensor ids32, Tensor gate, Tensor up, int type, " + "Tensor book, int mode) -> ()", &v41k_diag); + m.def("q8_aligned(Tensor(a!) qs, Tensor(b!) ds, Tensor x) -> ()", &v41k_q8_aligned); + m.def("moe5(Tensor(a!) out, Tensor xq, Tensor xd, Tensor(b!) hq, Tensor(c!) hd, Tensor ids32, " + "Tensor prob, Tensor gate, Tensor up, Tensor down, int gtype, int dtype, int left, " + "Tensor gbook, Tensor dbook, int rows) -> ()", &v41k_moe5); + m.def("moe6(Tensor(a!) out, Tensor xq, Tensor xd, Tensor(b!) hq, Tensor(c!) hd, Tensor ids32, " + "Tensor prob, Tensor gate, Tensor up, Tensor down, int gtype, int dtype, int left, " + "Tensor gbook, Tensor dbook, int rows, int rb) -> ()", &v41k_moe6); + m.def("hc_transition(Tensor(a!) residual, Tensor y, Tensor post, Tensor comb, Tensor cpre, " + "Tensor norm, Tensor(b!) out, Tensor(c!) cbuf, Tensor(d!) csq, Tensor? fn, Tensor? scale, " + "Tensor? base, Tensor(e!)? dots, Tensor(f!)? sq, Tensor(g!)? pre_n, Tensor(h!)? post_n, " + "Tensor(i!)? comb_n, float norm_eps, float hc_eps, int iters) -> ()", &v41k_hc_transition); + m.def("moe5w(Tensor(a!) out, Tensor xq, Tensor xd, Tensor(b!) hq, Tensor(c!) hd, Tensor ids32, " + "Tensor prob, Tensor gate, Tensor up, Tensor down, int gtype, int dtype, int left, " + "Tensor gbook, Tensor dbook, int rows, int dn_warps) -> ()", &v41k_moe5w); + m.def("hgemv(Tensor(a!) out, Tensor x, Tensor w) -> ()", &v41k_hgemv); + m.def("attn_gathered(Tensor(a!) out, Tensor q, Tensor keys, Tensor unused, Tensor sink, " + "Tensor(b!) ws, float scale) -> ()", &v41k_attn_gathered); + m.def("moe7(Tensor(a!) out, Tensor xq, Tensor xd, Tensor(b!) hq, Tensor(c!) hd, Tensor ids32, " + "Tensor prob, Tensor gate, Tensor up, Tensor down, int gtype, int dtype, int left, " + "Tensor gbook, Tensor dbook, int rows, int dn_warps, Tensor(d!) plan) -> ()", &v41k_moe7); + m.def("swap_stage(Tensor(a!) b0, Tensor(b!) b1, Tensor(c!) b2, Tensor host, Tensor(d!) scratch, " + "Tensor slots, Tensor loads, Tensor count, Tensor fast, int[] layout) -> ()", &v41k_swap_stage); + m.def("swap_signal(Tensor(a!) control, Tensor counter, Tensor(b!) wb_seq, Tensor count, " + "Tensor fast, int flag_word) -> ()", &v41k_swap_signal); + m.def("swap_writeback(Tensor(a!) host, Tensor scratch, Tensor loads, Tensor count, Tensor fast, " + "Tensor wb_seq, Tensor(b!) control, int flag0, int flag_stride, int ranks, int ack_word, " + "Tensor(c!) done, int[] layout, int ctas) -> ()", &v41k_swap_writeback); + m.def("swap_scratch_per_op(int[] layout) -> int", &v41k_swap_scratch_per_op); + m.def("swap_combined(Tensor(a!) b0, Tensor(b!) b1, Tensor(c!) b2, Tensor(d!) host, Tensor slots, " + "Tensor loads, Tensor saves, Tensor count, Tensor fast, int[] layout, int ctas) -> ()", + &v41k_swap_combined); + m.def("fp8_gemv(Tensor(a!) out, Tensor x, Tensor w, Tensor scales, Tensor partials, " + "int split, Tensor? bias) -> ()", &v41k_fp8_gemv); + m.def("gate_up4(Tensor(a!) down_in, Tensor x, Tensor ids32, Tensor prob, " + "Tensor gate, Tensor up, int type, Tensor book) -> ()", &v41k_gate_up4); + m.def("down4(Tensor(a!) out, Tensor h, Tensor ids32, Tensor down, int type, " + "int left, Tensor book, int rows_per_cta) -> ()", &v41k_down4); + m.def("gate_up3(Tensor(a!) down_in, Tensor x, Tensor ids32, Tensor prob, " + "Tensor gate, Tensor up, int type, Tensor book) -> ()", &v41k_gate_up3); + m.def("down3(Tensor(a!) out, Tensor h, Tensor ids32, Tensor down, int type, " + "int left, Tensor book, int rows_per_cta) -> ()", &v41k_down3); + m.def("decode3(Tensor(a!) sink, Tensor ids32, Tensor gate, Tensor up, int type, " + "Tensor book) -> ()", &v41k_decode3); + m.def("gate_up2(Tensor(a!) down_in, Tensor x, Tensor ids32, Tensor prob, " + "Tensor gate, Tensor up, int type, Tensor book) -> ()", &v41k_gate_up2); + m.def("down2(Tensor(a!) out, Tensor h, Tensor ids32, Tensor down, int type, " + "int left, Tensor book) -> ()", &v41k_down2); +} diff --git a/csrc/gguf_cpu_iq_tables.h b/csrc/gguf_cpu_iq_tables.h new file mode 100644 index 0000000000..caf4c1da4c --- /dev/null +++ b/csrc/gguf_cpu_iq_tables.h @@ -0,0 +1,733 @@ +// SPDX-License-Identifier: MIT +// IQ grids and sign tables for the CPU cold-expert kernels, extracted from +// ggml-common.h of llama.cpp (ggml-org/llama.cpp @ 9b8446b9, MIT License, +// Copyright (c) 2023-2024 The ggml authors). +#pragma once +#include + +namespace vllm::gguf::cpu_iq { + +inline constexpr uint64_t iq2xxs_grid[256] = { + 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, + 0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x08080808082b0808, + 0x08080808082b082b, 0x08080808082b2b08, 0x08080808082b2b2b, 0x0808080819080819, + 0x0808080819081908, 0x0808080819190808, 0x0808080819192b08, 0x08080808192b0819, + 0x08080808192b1908, 0x080808082b080808, 0x080808082b08082b, 0x080808082b082b2b, + 0x080808082b2b082b, 0x0808081908080819, 0x0808081908081908, 0x0808081908190808, + 0x0808081908191919, 0x0808081919080808, 0x080808192b081908, 0x080808192b192b08, + 0x0808082b08080808, 0x0808082b0808082b, 0x0808082b082b082b, 0x0808082b2b08082b, + 0x0808190808080819, 0x0808190808081908, 0x0808190808190808, 0x08081908082b0819, + 0x08081908082b1908, 0x0808190819080808, 0x080819081908082b, 0x0808190819082b08, + 0x08081908192b0808, 0x080819082b080819, 0x080819082b081908, 0x080819082b190808, + 0x080819082b2b1908, 0x0808191908080808, 0x080819190808082b, 0x0808191908082b08, + 0x08081919082b0808, 0x080819191908192b, 0x08081919192b2b19, 0x080819192b080808, + 0x080819192b190819, 0x0808192b08082b19, 0x0808192b08190808, 0x0808192b19080808, + 0x0808192b2b081908, 0x0808192b2b2b1908, 0x08082b0808080808, 0x08082b0808081919, + 0x08082b0808082b08, 0x08082b0808191908, 0x08082b08082b2b08, 0x08082b0819080819, + 0x08082b0819081908, 0x08082b0819190808, 0x08082b081919082b, 0x08082b082b082b08, + 0x08082b1908081908, 0x08082b1919080808, 0x08082b2b0808082b, 0x08082b2b08191908, + 0x0819080808080819, 0x0819080808081908, 0x0819080808190808, 0x08190808082b0819, + 0x0819080819080808, 0x08190808192b0808, 0x081908082b081908, 0x081908082b190808, + 0x081908082b191919, 0x0819081908080808, 0x0819081908082b08, 0x08190819082b0808, + 0x0819081919190808, 0x0819081919192b2b, 0x081908192b080808, 0x0819082b082b1908, + 0x0819082b19081919, 0x0819190808080808, 0x0819190808082b08, 0x08191908082b0808, + 0x08191908082b1919, 0x0819190819082b19, 0x081919082b080808, 0x0819191908192b08, + 0x08191919192b082b, 0x0819192b08080808, 0x0819192b0819192b, 0x08192b0808080819, + 0x08192b0808081908, 0x08192b0808190808, 0x08192b0819080808, 0x08192b082b080819, + 0x08192b1908080808, 0x08192b1908081919, 0x08192b192b2b0808, 0x08192b2b19190819, + 0x082b080808080808, 0x082b08080808082b, 0x082b080808082b2b, 0x082b080819081908, + 0x082b0808192b0819, 0x082b08082b080808, 0x082b08082b08082b, 0x082b0819082b2b19, + 0x082b081919082b08, 0x082b082b08080808, 0x082b082b0808082b, 0x082b190808080819, + 0x082b190808081908, 0x082b190808190808, 0x082b190819080808, 0x082b19081919192b, + 0x082b191908080808, 0x082b191919080819, 0x082b1919192b1908, 0x082b192b2b190808, + 0x082b2b0808082b08, 0x082b2b08082b0808, 0x082b2b082b191908, 0x082b2b2b19081908, + 0x1908080808080819, 0x1908080808081908, 0x1908080808190808, 0x1908080808192b08, + 0x19080808082b0819, 0x19080808082b1908, 0x1908080819080808, 0x1908080819082b08, + 0x190808081919192b, 0x19080808192b0808, 0x190808082b080819, 0x190808082b081908, + 0x190808082b190808, 0x1908081908080808, 0x19080819082b0808, 0x19080819192b0819, + 0x190808192b080808, 0x190808192b081919, 0x1908082b08080819, 0x1908082b08190808, + 0x1908082b19082b08, 0x1908082b1919192b, 0x1908082b192b2b08, 0x1908190808080808, + 0x1908190808082b08, 0x19081908082b0808, 0x190819082b080808, 0x190819082b192b19, + 0x190819190819082b, 0x19081919082b1908, 0x1908192b08080808, 0x19082b0808080819, + 0x19082b0808081908, 0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919, + 0x19082b1908080808, 0x19082b1919192b08, 0x19082b19192b0819, 0x19082b192b08082b, + 0x19082b2b19081919, 0x19082b2b2b190808, 0x1919080808080808, 0x1919080808082b08, + 0x1919080808190819, 0x1919080808192b19, 0x19190808082b0808, 0x191908082b080808, + 0x191908082b082b08, 0x1919081908081908, 0x191908191908082b, 0x191908192b2b1908, + 0x1919082b2b190819, 0x191919082b190808, 0x191919082b19082b, 0x1919191908082b2b, + 0x1919192b08080819, 0x1919192b19191908, 0x19192b0808080808, 0x19192b0808190819, + 0x19192b0808192b19, 0x19192b08192b1908, 0x19192b1919080808, 0x19192b2b08082b08, + 0x192b080808081908, 0x192b080808190808, 0x192b080819080808, 0x192b0808192b2b08, + 0x192b081908080808, 0x192b081919191919, 0x192b082b08192b08, 0x192b082b192b0808, + 0x192b190808080808, 0x192b190808081919, 0x192b191908190808, 0x192b19190819082b, + 0x192b19192b081908, 0x192b2b081908082b, 0x2b08080808080808, 0x2b0808080808082b, + 0x2b08080808082b2b, 0x2b08080819080819, 0x2b0808082b08082b, 0x2b08081908081908, + 0x2b08081908192b08, 0x2b08081919080808, 0x2b08082b08190819, 0x2b08190808080819, + 0x2b08190808081908, 0x2b08190808190808, 0x2b08190808191919, 0x2b08190819080808, + 0x2b081908192b0808, 0x2b08191908080808, 0x2b0819191908192b, 0x2b0819192b191908, + 0x2b08192b08082b19, 0x2b08192b19080808, 0x2b08192b192b0808, 0x2b082b080808082b, + 0x2b082b1908081908, 0x2b082b2b08190819, 0x2b19080808081908, 0x2b19080808190808, + 0x2b190808082b1908, 0x2b19080819080808, 0x2b1908082b2b0819, 0x2b1908190819192b, + 0x2b1908192b080808, 0x2b19082b19081919, 0x2b19190808080808, 0x2b191908082b082b, + 0x2b19190819081908, 0x2b19191919190819, 0x2b192b082b080819, 0x2b192b19082b0808, + 0x2b2b08080808082b, 0x2b2b080819190808, 0x2b2b08082b081919, 0x2b2b081908082b19, + 0x2b2b082b08080808, 0x2b2b190808192b08, 0x2b2b2b0819190808, 0x2b2b2b1908081908}; + +inline constexpr uint64_t iq2xs_grid[512] = { + 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, + 0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x080808080819192b, + 0x0808080808192b19, 0x08080808082b0808, 0x08080808082b082b, 0x08080808082b1919, + 0x08080808082b2b08, 0x0808080819080819, 0x0808080819081908, 0x080808081908192b, + 0x0808080819082b19, 0x0808080819190808, 0x080808081919082b, 0x0808080819191919, + 0x0808080819192b08, 0x08080808192b0819, 0x08080808192b1908, 0x080808082b080808, + 0x080808082b08082b, 0x080808082b081919, 0x080808082b082b08, 0x080808082b190819, + 0x080808082b191908, 0x080808082b192b19, 0x080808082b2b0808, 0x0808081908080819, + 0x0808081908081908, 0x080808190808192b, 0x0808081908082b19, 0x0808081908190808, + 0x080808190819082b, 0x0808081908191919, 0x0808081908192b08, 0x0808081908192b2b, + 0x08080819082b0819, 0x08080819082b1908, 0x0808081919080808, 0x080808191908082b, + 0x0808081919081919, 0x0808081919082b08, 0x0808081919190819, 0x0808081919191908, + 0x08080819192b0808, 0x08080819192b2b08, 0x080808192b080819, 0x080808192b081908, + 0x080808192b190808, 0x0808082b08080808, 0x0808082b0808082b, 0x0808082b08081919, + 0x0808082b08082b08, 0x0808082b08190819, 0x0808082b08191908, 0x0808082b082b0808, + 0x0808082b19080819, 0x0808082b19081908, 0x0808082b19190808, 0x0808082b19191919, + 0x0808082b2b080808, 0x0808082b2b082b2b, 0x0808190808080819, 0x0808190808081908, + 0x080819080808192b, 0x0808190808082b19, 0x0808190808190808, 0x080819080819082b, + 0x0808190808191919, 0x0808190808192b08, 0x08081908082b0819, 0x08081908082b1908, + 0x0808190819080808, 0x080819081908082b, 0x0808190819081919, 0x0808190819082b08, + 0x0808190819190819, 0x0808190819191908, 0x080819081919192b, 0x08081908192b0808, + 0x080819082b080819, 0x080819082b081908, 0x080819082b190808, 0x0808191908080808, + 0x080819190808082b, 0x0808191908081919, 0x0808191908082b08, 0x0808191908190819, + 0x0808191908191908, 0x08081919082b0808, 0x0808191919080819, 0x0808191919081908, + 0x0808191919190808, 0x08081919192b0819, 0x080819192b080808, 0x0808192b08080819, + 0x0808192b08081908, 0x0808192b08190808, 0x0808192b082b192b, 0x0808192b19080808, + 0x0808192b1908082b, 0x0808192b2b081908, 0x08082b0808080808, 0x08082b080808082b, + 0x08082b0808081919, 0x08082b0808082b08, 0x08082b0808082b2b, 0x08082b0808190819, + 0x08082b0808191908, 0x08082b08082b0808, 0x08082b08082b1919, 0x08082b0819080819, + 0x08082b0819081908, 0x08082b0819190808, 0x08082b0819192b08, 0x08082b082b080808, + 0x08082b082b2b0808, 0x08082b082b2b2b2b, 0x08082b1908080819, 0x08082b1908081908, + 0x08082b1908190808, 0x08082b1919080808, 0x08082b192b080819, 0x08082b192b082b19, + 0x08082b2b08080808, 0x08082b2b082b0808, 0x08082b2b082b2b08, 0x08082b2b2b19192b, + 0x08082b2b2b2b0808, 0x0819080808080819, 0x0819080808081908, 0x081908080808192b, + 0x0819080808082b19, 0x0819080808190808, 0x081908080819082b, 0x0819080808191919, + 0x0819080808192b08, 0x08190808082b0819, 0x08190808082b1908, 0x0819080819080808, + 0x081908081908082b, 0x0819080819081919, 0x0819080819082b08, 0x0819080819190819, + 0x0819080819191908, 0x08190808192b0808, 0x08190808192b2b2b, 0x081908082b080819, + 0x081908082b081908, 0x081908082b190808, 0x0819081908080808, 0x081908190808082b, + 0x0819081908081919, 0x0819081908082b08, 0x0819081908190819, 0x0819081908191908, + 0x08190819082b0808, 0x0819081919080819, 0x0819081919081908, 0x0819081919190808, + 0x081908192b080808, 0x081908192b191908, 0x081908192b19192b, 0x0819082b08080819, + 0x0819082b08081908, 0x0819082b0808192b, 0x0819082b08190808, 0x0819082b19080808, + 0x0819082b192b0808, 0x0819190808080808, 0x081919080808082b, 0x0819190808081919, + 0x0819190808082b08, 0x0819190808190819, 0x0819190808191908, 0x08191908082b0808, + 0x0819190819080819, 0x0819190819081908, 0x0819190819082b19, 0x0819190819190808, + 0x08191908192b1908, 0x081919082b080808, 0x0819191908080819, 0x0819191908081908, + 0x0819191908190808, 0x0819191919080808, 0x0819192b08080808, 0x0819192b08191908, + 0x0819192b19082b19, 0x08192b0808080819, 0x08192b0808081908, 0x08192b0808190808, + 0x08192b080819082b, 0x08192b0819080808, 0x08192b0819191908, 0x08192b082b08192b, + 0x08192b1908080808, 0x08192b1908081919, 0x08192b19192b192b, 0x08192b2b19190819, + 0x08192b2b2b2b2b19, 0x082b080808080808, 0x082b08080808082b, 0x082b080808081919, + 0x082b080808082b08, 0x082b080808082b2b, 0x082b080808190819, 0x082b080808191908, + 0x082b0808082b0808, 0x082b080819080819, 0x082b080819081908, 0x082b080819190808, + 0x082b08082b080808, 0x082b08082b2b0808, 0x082b081908080819, 0x082b081908081908, + 0x082b081908190808, 0x082b081919080808, 0x082b081919082b08, 0x082b0819192b1919, + 0x082b082b08080808, 0x082b082b082b082b, 0x082b082b2b080808, 0x082b082b2b2b2b08, + 0x082b190808080819, 0x082b190808081908, 0x082b190808190808, 0x082b1908082b2b19, + 0x082b190819080808, 0x082b191908080808, 0x082b191919080819, 0x082b19191919082b, + 0x082b19192b192b19, 0x082b192b08080819, 0x082b192b08192b2b, 0x082b192b2b2b192b, + 0x082b2b0808080808, 0x082b2b0808082b08, 0x082b2b0808082b2b, 0x082b2b08082b0808, + 0x082b2b0819191919, 0x082b2b082b082b08, 0x082b2b082b2b082b, 0x082b2b19192b2b08, + 0x082b2b192b190808, 0x082b2b2b08082b08, 0x082b2b2b082b0808, 0x082b2b2b2b08082b, + 0x082b2b2b2b082b08, 0x082b2b2b2b082b2b, 0x1908080808080819, 0x1908080808081908, + 0x190808080808192b, 0x1908080808082b19, 0x1908080808190808, 0x190808080819082b, + 0x1908080808191919, 0x1908080808192b08, 0x19080808082b0819, 0x19080808082b1908, + 0x1908080819080808, 0x190808081908082b, 0x1908080819081919, 0x1908080819082b08, + 0x1908080819082b2b, 0x1908080819190819, 0x1908080819191908, 0x19080808192b0808, + 0x19080808192b1919, 0x190808082b080819, 0x190808082b081908, 0x190808082b190808, + 0x1908081908080808, 0x190808190808082b, 0x1908081908081919, 0x1908081908082b08, + 0x1908081908190819, 0x1908081908191908, 0x19080819082b0808, 0x1908081919080819, + 0x1908081919081908, 0x1908081919190808, 0x190808192b080808, 0x190808192b081919, + 0x190808192b2b082b, 0x1908082b08080819, 0x1908082b08081908, 0x1908082b08190808, + 0x1908082b0819082b, 0x1908082b082b2b19, 0x1908082b19080808, 0x1908190808080808, + 0x190819080808082b, 0x1908190808081919, 0x1908190808082b08, 0x1908190808190819, + 0x1908190808191908, 0x1908190808192b19, 0x19081908082b0808, 0x1908190819080819, + 0x1908190819081908, 0x1908190819190808, 0x190819082b080808, 0x190819082b191908, + 0x1908191908080819, 0x1908191908081908, 0x1908191908190808, 0x19081919082b1908, + 0x1908191919080808, 0x190819192b192b2b, 0x1908192b08080808, 0x1908192b08082b2b, + 0x1908192b19081908, 0x1908192b19190808, 0x19082b0808080819, 0x19082b0808081908, + 0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919, 0x19082b0819191908, + 0x19082b08192b082b, 0x19082b1908080808, 0x19082b1908190819, 0x19082b1919081908, + 0x19082b1919190808, 0x19082b19192b2b19, 0x19082b2b08081908, 0x1919080808080808, + 0x191908080808082b, 0x1919080808081919, 0x1919080808082b08, 0x1919080808190819, + 0x1919080808191908, 0x19190808082b0808, 0x19190808082b2b08, 0x1919080819080819, + 0x1919080819081908, 0x1919080819190808, 0x191908082b080808, 0x1919081908080819, + 0x1919081908081908, 0x1919081908190808, 0x1919081908191919, 0x1919081919080808, + 0x191908191908082b, 0x1919082b08080808, 0x1919082b19081908, 0x1919082b2b2b2b2b, + 0x1919190808080819, 0x1919190808081908, 0x1919190808190808, 0x19191908082b0819, + 0x1919190819080808, 0x19191908192b0808, 0x191919082b080819, 0x191919082b2b0819, + 0x1919191908080808, 0x1919191908082b08, 0x191919192b080808, 0x191919192b082b08, + 0x1919192b082b0819, 0x1919192b192b2b08, 0x1919192b2b2b0819, 0x19192b0808080808, + 0x19192b0808191908, 0x19192b0819080819, 0x19192b0819190808, 0x19192b082b192b19, + 0x19192b1908192b2b, 0x19192b1919080808, 0x19192b191908082b, 0x19192b2b2b081919, + 0x192b080808080819, 0x192b080808081908, 0x192b080808190808, 0x192b080819080808, + 0x192b080819191908, 0x192b0808192b082b, 0x192b08082b08192b, 0x192b08082b2b2b19, + 0x192b081908080808, 0x192b082b082b1908, 0x192b082b19082b2b, 0x192b082b2b19082b, + 0x192b190808080808, 0x192b19080819192b, 0x192b191908190808, 0x192b191919080808, + 0x192b191919081919, 0x192b19192b2b1908, 0x192b2b0808080819, 0x192b2b08192b2b2b, + 0x192b2b19082b1919, 0x192b2b2b0808192b, 0x192b2b2b19191908, 0x192b2b2b192b082b, + 0x2b08080808080808, 0x2b0808080808082b, 0x2b08080808081919, 0x2b08080808082b08, + 0x2b08080808190819, 0x2b08080808191908, 0x2b080808082b0808, 0x2b080808082b2b2b, + 0x2b08080819080819, 0x2b08080819081908, 0x2b08080819190808, 0x2b0808082b080808, + 0x2b0808082b08082b, 0x2b0808082b2b2b08, 0x2b0808082b2b2b2b, 0x2b08081908080819, + 0x2b08081908081908, 0x2b0808190808192b, 0x2b08081908190808, 0x2b08081919080808, + 0x2b08081919190819, 0x2b08081919192b19, 0x2b08082b08080808, 0x2b08082b082b0808, + 0x2b08082b2b080808, 0x2b08082b2b08082b, 0x2b08082b2b2b0808, 0x2b08082b2b2b2b08, + 0x2b08190808080819, 0x2b08190808081908, 0x2b08190808190808, 0x2b0819080819082b, + 0x2b08190808191919, 0x2b08190819080808, 0x2b081908192b0808, 0x2b0819082b082b19, + 0x2b08191908080808, 0x2b08191919081908, 0x2b0819192b2b1919, 0x2b08192b08192b08, + 0x2b08192b192b2b2b, 0x2b082b0808080808, 0x2b082b0808082b08, 0x2b082b08082b1919, + 0x2b082b0819192b2b, 0x2b082b082b080808, 0x2b082b082b08082b, 0x2b082b082b2b2b08, + 0x2b082b190808192b, 0x2b082b2b082b082b, 0x2b082b2b2b080808, 0x2b082b2b2b082b08, + 0x2b082b2b2b19192b, 0x2b082b2b2b2b2b08, 0x2b19080808080819, 0x2b19080808081908, + 0x2b19080808190808, 0x2b19080819080808, 0x2b1908081919192b, 0x2b1908082b081908, + 0x2b19081908080808, 0x2b190819082b082b, 0x2b190819192b1908, 0x2b19082b1919192b, + 0x2b19082b2b082b19, 0x2b19190808080808, 0x2b19190808081919, 0x2b19190819081908, + 0x2b19190819190808, 0x2b19190819192b08, 0x2b191919082b2b19, 0x2b1919192b190808, + 0x2b1919192b19082b, 0x2b19192b19080819, 0x2b192b0819190819, 0x2b192b082b2b192b, + 0x2b192b1919082b19, 0x2b192b2b08191919, 0x2b192b2b192b0808, 0x2b2b080808080808, + 0x2b2b08080808082b, 0x2b2b080808082b08, 0x2b2b080808082b2b, 0x2b2b0808082b0808, + 0x2b2b0808082b2b2b, 0x2b2b08082b2b0808, 0x2b2b081919190819, 0x2b2b081919192b19, + 0x2b2b08192b2b192b, 0x2b2b082b08080808, 0x2b2b082b0808082b, 0x2b2b082b08082b08, + 0x2b2b082b082b2b2b, 0x2b2b082b2b080808, 0x2b2b082b2b2b0808, 0x2b2b190819080808, + 0x2b2b19082b191919, 0x2b2b192b192b1919, 0x2b2b192b2b192b08, 0x2b2b2b0808082b2b, + 0x2b2b2b08082b0808, 0x2b2b2b08082b082b, 0x2b2b2b08082b2b08, 0x2b2b2b082b2b0808, + 0x2b2b2b082b2b2b08, 0x2b2b2b1908081908, 0x2b2b2b192b081908, 0x2b2b2b192b08192b, + 0x2b2b2b2b082b2b08, 0x2b2b2b2b082b2b2b, 0x2b2b2b2b2b190819, 0x2b2b2b2b2b2b2b2b}; + +inline constexpr uint64_t iq1s_grid[2048] = { + 0xffffffffffffffff, 0xffffffffffffff01, 0xffffffffffff0000, 0xffffffffffff01ff, + 0xffffffffffff0101, 0xffffffffff00ff00, 0xffffffffff000000, 0xffffffffff01ffff, + 0xffffffffff01ff01, 0xffffffffff0101ff, 0xffffffffff010101, 0xffffffff00ff0000, + 0xffffffff0000ff00, 0xffffffff000000ff, 0xffffffff00000001, 0xffffffff00010000, + 0xffffffff01ffffff, 0xffffffff01ffff01, 0xffffffff01ff01ff, 0xffffffff01ff0101, + 0xffffffff01000000, 0xffffffff0101ffff, 0xffffffff0101ff01, 0xffffffff010101ff, + 0xffffffff01010101, 0xffffff00ffff00ff, 0xffffff00ffff0000, 0xffffff00ff00ff00, + 0xffffff00ff0000ff, 0xffffff00ff000001, 0xffffff00ff000100, 0xffffff00ff000101, + 0xffffff00ff010000, 0xffffff0000ffff00, 0xffffff0000ff0001, 0xffffff0000ff0100, + 0xffffff000000ff01, 0xffffff0000000000, 0xffffff0000000101, 0xffffff000001ff00, + 0xffffff00000100ff, 0xffffff0000010001, 0xffffff00000101ff, 0xffffff0001ff0000, + 0xffffff000100ff00, 0xffffff00010000ff, 0xffffff0001000001, 0xffffff0001010000, + 0xffffff01ffffffff, 0xffffff01ffffff01, 0xffffff01ffff01ff, 0xffffff01ffff0101, + 0xffffff01ff000000, 0xffffff01ff01ffff, 0xffffff01ff01ff01, 0xffffff01ff0101ff, + 0xffffff01ff010101, 0xffffff0100ff0000, 0xffffff010000ff00, 0xffffff0100000100, + 0xffffff01000100ff, 0xffffff0100010100, 0xffffff0101ffffff, 0xffffff0101ffff01, + 0xffffff0101ff01ff, 0xffffff0101ff0101, 0xffffff010100ff00, 0xffffff0101000000, + 0xffffff0101000100, 0xffffff010101ffff, 0xffffff010101ff01, 0xffffff01010101ff, + 0xffffff0101010101, 0xffff00ffff00ff00, 0xffff00ffff0000ff, 0xffff00ffff000001, + 0xffff00ffff010000, 0xffff00ff00ffff00, 0xffff00ff00ff0100, 0xffff00ff00000000, + 0xffff00ff00000101, 0xffff00ff000100ff, 0xffff00ff00010000, 0xffff00ff0100ff00, + 0xffff00ff01000100, 0xffff00ff01010000, 0xffff0000ffffff00, 0xffff0000ffff00ff, + 0xffff0000ffff0000, 0xffff0000ffff0001, 0xffff0000ff000000, 0xffff0000ff0001ff, + 0xffff0000ff000101, 0xffff0000ff010100, 0xffff000000ffffff, 0xffff000000ff0000, + 0xffff000000ff0101, 0xffff00000000ffff, 0xffff00000000ff00, 0xffff0000000000ff, + 0xffff000000000000, 0xffff000000000001, 0xffff000000000100, 0xffff00000001ffff, + 0xffff00000001ff01, 0xffff000000010000, 0xffff0000000101ff, 0xffff000000010101, + 0xffff000001ffff00, 0xffff00000100ff00, 0xffff000001000000, 0xffff0000010001ff, + 0xffff000001000101, 0xffff00000101ff00, 0xffff0000010100ff, 0xffff000001010000, + 0xffff000001010001, 0xffff000001010100, 0xffff0001ff0000ff, 0xffff0001ff000100, + 0xffff000100ffff00, 0xffff000100ff00ff, 0xffff00010000ffff, 0xffff00010000ff01, + 0xffff000100000000, 0xffff0001000001ff, 0xffff00010001ffff, 0xffff00010001ff00, + 0xffff000100010001, 0xffff000100010100, 0xffff000101ff0000, 0xffff00010100ff00, + 0xffff0001010000ff, 0xffff000101000100, 0xffff01ffffffffff, 0xffff01ffffffff01, + 0xffff01ffffff01ff, 0xffff01ffffff0101, 0xffff01ffff000000, 0xffff01ffff01ffff, + 0xffff01ffff01ff01, 0xffff01ffff0101ff, 0xffff01ffff010101, 0xffff01ff00ff0000, + 0xffff01ff0000ff00, 0xffff01ff00000001, 0xffff01ff00010000, 0xffff01ff01ffffff, + 0xffff01ff01ffff01, 0xffff01ff01ff01ff, 0xffff01ff01ff0101, 0xffff01ff01000000, + 0xffff01ff0101ffff, 0xffff01ff0101ff01, 0xffff01ff010101ff, 0xffff01ff01010101, + 0xffff0100ffff0000, 0xffff0100ff00ff00, 0xffff0100ff0000ff, 0xffff0100ff000100, + 0xffff0100ff0100ff, 0xffff0100ff010000, 0xffff010000ffff00, 0xffff01000000ffff, + 0xffff01000000ff00, 0xffff010000000000, 0xffff01000001ff00, 0xffff0100000100ff, + 0xffff010000010100, 0xffff01000100ff00, 0xffff0100010000ff, 0xffff010001000001, + 0xffff010001000100, 0xffff010001010000, 0xffff0101ffffffff, 0xffff0101ffffff01, + 0xffff0101ffff01ff, 0xffff0101ffff0101, 0xffff0101ff000000, 0xffff0101ff01ffff, + 0xffff0101ff01ff01, 0xffff0101ff0101ff, 0xffff0101ff010101, 0xffff010100ff0000, + 0xffff01010000ff00, 0xffff010100000100, 0xffff01010001ff00, 0xffff010100010000, + 0xffff010101ffffff, 0xffff010101ffff01, 0xffff010101ff0000, 0xffff010101ff01ff, + 0xffff010101ff0101, 0xffff010101000000, 0xffff01010101ffff, 0xffff01010101ff01, + 0xffff0101010101ff, 0xffff010101010101, 0xff00ffffff00ffff, 0xff00ffffff00ff00, + 0xff00ffffff0000ff, 0xff00ffffff000100, 0xff00ffffff0100ff, 0xff00ffffff010000, + 0xff00ffff00ffff00, 0xff00ffff00ff00ff, 0xff00ffff0000ffff, 0xff00ffff00000000, + 0xff00ffff000001ff, 0xff00ffff0001ff00, 0xff00ffff000100ff, 0xff00ffff00010000, + 0xff00ffff00010100, 0xff00ffff0100ff00, 0xff00ffff010000ff, 0xff00ffff01000001, + 0xff00ffff0101ff00, 0xff00ffff01010000, 0xff00ff00ffffff00, 0xff00ff00ffff00ff, + 0xff00ff00ffff0001, 0xff00ff00ffff0100, 0xff00ff00ff00ffff, 0xff00ff00ff00ff01, + 0xff00ff00ff000000, 0xff00ff00ff0001ff, 0xff00ff00ff01ff00, 0xff00ff00ff0100ff, + 0xff00ff00ff010100, 0xff00ff0000ff0000, 0xff00ff0000ff0101, 0xff00ff000000ffff, + 0xff00ff000000ff00, 0xff00ff000000ff01, 0xff00ff00000000ff, 0xff00ff0000000000, + 0xff00ff0000000001, 0xff00ff0000000100, 0xff00ff000001ffff, 0xff00ff0000010000, + 0xff00ff0001ff00ff, 0xff00ff000100ff01, 0xff00ff0001000000, 0xff00ff000101ff00, + 0xff00ff00010100ff, 0xff00ff01ff00ff00, 0xff00ff01ff0000ff, 0xff00ff01ff000001, + 0xff00ff01ff010000, 0xff00ff0100ffffff, 0xff00ff0100ff0001, 0xff00ff0100ff0100, + 0xff00ff010000ff01, 0xff00ff0100000000, 0xff00ff01000001ff, 0xff00ff0100000101, + 0xff00ff01000100ff, 0xff00ff0100010001, 0xff00ff0101ff0000, 0xff00ff010100ff00, + 0xff00ff01010000ff, 0xff00ff0101000001, 0xff00ff0101010000, 0xff0000ffffffff00, + 0xff0000ffffff0001, 0xff0000ffffff0100, 0xff0000ffff0000ff, 0xff0000ffff000000, + 0xff0000ffff0001ff, 0xff0000ffff000100, 0xff0000ffff01ff00, 0xff0000ffff010001, + 0xff0000ff00ffff00, 0xff0000ff00ff0000, 0xff0000ff00ff0001, 0xff0000ff00ff01ff, + 0xff0000ff00ff0101, 0xff0000ff0000ff00, 0xff0000ff000000ff, 0xff0000ff00000000, + 0xff0000ff00000001, 0xff0000ff00000100, 0xff0000ff0001ff01, 0xff0000ff00010000, + 0xff0000ff000101ff, 0xff0000ff01ff00ff, 0xff0000ff01ff0100, 0xff0000ff0100ffff, + 0xff0000ff010000ff, 0xff0000ff01000000, 0xff0000ff010001ff, 0xff0000ff01000100, + 0xff0000ff01000101, 0xff0000ff0101ff00, 0xff0000ff010100ff, 0xff0000ff01010000, + 0xff0000ff01010100, 0xff000000ffffff01, 0xff000000ffff0000, 0xff000000ffff0101, + 0xff000000ff00ff00, 0xff000000ff0000ff, 0xff000000ff000000, 0xff000000ff000001, + 0xff000000ff000100, 0xff000000ff01ffff, 0xff000000ff01ff01, 0xff000000ff010000, + 0xff000000ff0101ff, 0xff000000ff010101, 0xff00000000ffff00, 0xff00000000ff00ff, + 0xff00000000ff0000, 0xff00000000ff0001, 0xff0000000000ff00, 0xff0000000000ff01, + 0xff000000000000ff, 0xff00000000000000, 0xff00000000000001, 0xff00000000000100, + 0xff00000000000101, 0xff0000000001ff00, 0xff000000000100ff, 0xff00000000010000, + 0xff00000000010001, 0xff00000000010100, 0xff00000001ffffff, 0xff00000001ffff01, + 0xff00000001ff00ff, 0xff00000001ff0000, 0xff00000001ff01ff, 0xff00000001ff0101, + 0xff0000000100ffff, 0xff0000000100ff00, 0xff000000010000ff, 0xff00000001000000, + 0xff00000001000001, 0xff00000001000100, 0xff00000001000101, 0xff0000000101ffff, + 0xff0000000101ff01, 0xff00000001010000, 0xff000001ffffff00, 0xff000001ffff00ff, + 0xff000001ffff0000, 0xff000001ffff0001, 0xff000001ff000000, 0xff000001ff000001, + 0xff000001ff0001ff, 0xff000001ff000101, 0xff000001ff01ff00, 0xff000001ff010001, + 0xff00000100ffffff, 0xff00000100ffff01, 0xff00000100ff00ff, 0xff00000100ff0000, + 0xff00000100ff01ff, 0xff00000100ff0101, 0xff0000010000ff00, 0xff00000100000000, + 0xff00000100000001, 0xff000001000001ff, 0xff00000100000100, 0xff0000010001ff00, + 0xff000001000100ff, 0xff00000100010000, 0xff000001000101ff, 0xff00000100010100, + 0xff00000100010101, 0xff00000101ff0001, 0xff00000101ff0101, 0xff0000010100ff01, + 0xff00000101000000, 0xff000001010100ff, 0xff00000101010100, 0xff0001ffff00ff00, + 0xff0001ffff000001, 0xff0001ffff010000, 0xff0001ff00ffff00, 0xff0001ff00ff00ff, + 0xff0001ff00ff0001, 0xff0001ff00ff0100, 0xff0001ff0000ffff, 0xff0001ff00000000, + 0xff0001ff000001ff, 0xff0001ff00000101, 0xff0001ff0001ffff, 0xff0001ff0001ff00, + 0xff0001ff000100ff, 0xff0001ff00010001, 0xff0001ff00010100, 0xff0001ff01ff0000, + 0xff0001ff0100ff00, 0xff0001ff010000ff, 0xff0001ff01010000, 0xff000100ff00ffff, + 0xff000100ff00ff01, 0xff000100ff000000, 0xff000100ff000101, 0xff000100ff01ff00, + 0xff000100ff010000, 0xff00010000ffff01, 0xff00010000ff00ff, 0xff00010000ff0000, + 0xff00010000ff01ff, 0xff0001000000ff00, 0xff000100000000ff, 0xff00010000000000, + 0xff00010000000001, 0xff00010000000100, 0xff00010000000101, 0xff0001000001ffff, + 0xff00010000010000, 0xff00010000010101, 0xff00010001ff0100, 0xff0001000100ff00, + 0xff0001000100ff01, 0xff00010001000000, 0xff000100010001ff, 0xff0001000101ff00, + 0xff00010001010001, 0xff00010001010100, 0xff000101ffff0100, 0xff000101ff000001, + 0xff000101ff0100ff, 0xff000101ff010001, 0xff00010100ff00ff, 0xff00010100ff0001, + 0xff00010100ff0100, 0xff0001010000ffff, 0xff0001010000ff01, 0xff00010100000000, + 0xff000101000001ff, 0xff0001010001ff00, 0xff00010100010001, 0xff00010100010100, + 0xff00010101ff0000, 0xff0001010100ff00, 0xff00010101000001, 0xff00010101000101, + 0xff01ffffffffffff, 0xff01ffffffffff01, 0xff01ffffffff01ff, 0xff01ffffffff0101, + 0xff01ffffff000000, 0xff01ffffff01ffff, 0xff01ffffff01ff01, 0xff01ffffff010000, + 0xff01ffffff0101ff, 0xff01ffffff010101, 0xff01ffff00ff0000, 0xff01ffff0000ff00, + 0xff01ffff00000100, 0xff01ffff0001ff00, 0xff01ffff00010000, 0xff01ffff01ffffff, + 0xff01ffff01ffff01, 0xff01ffff01ff01ff, 0xff01ffff01ff0101, 0xff01ffff01000000, + 0xff01ffff0101ffff, 0xff01ffff0101ff01, 0xff01ffff01010000, 0xff01ffff010101ff, + 0xff01ffff01010101, 0xff01ff00ffff0000, 0xff01ff00ff00ff00, 0xff01ff00ff0000ff, + 0xff01ff00ff000100, 0xff01ff00ff010000, 0xff01ff0000ffff01, 0xff01ff0000ff00ff, + 0xff01ff0000ff0100, 0xff01ff0000000000, 0xff01ff00000001ff, 0xff01ff0000000101, + 0xff01ff000001ff00, 0xff01ff00000100ff, 0xff01ff0000010000, 0xff01ff0000010001, + 0xff01ff0001ff0000, 0xff01ff000100ffff, 0xff01ff0001000001, 0xff01ff0001000100, + 0xff01ff0001010000, 0xff01ff01ffffff00, 0xff01ff01ffff01ff, 0xff01ff01ffff0101, + 0xff01ff01ff00ff00, 0xff01ff01ff000000, 0xff01ff01ff01ffff, 0xff01ff01ff01ff01, + 0xff01ff01ff0101ff, 0xff01ff01ff010101, 0xff01ff0100ff0000, 0xff01ff010000ff00, + 0xff01ff0100000001, 0xff01ff0100000100, 0xff01ff0100010000, 0xff01ff0101ffff00, + 0xff01ff0101ff01ff, 0xff01ff0101ff0101, 0xff01ff010100ff00, 0xff01ff0101000000, + 0xff01ff010101ffff, 0xff01ff010101ff01, 0xff01ff01010101ff, 0xff01ff0101010101, + 0xff0100ffffff0000, 0xff0100ffff0000ff, 0xff0100ffff000001, 0xff0100ffff000100, + 0xff0100ffff010000, 0xff0100ff00ff00ff, 0xff0100ff00ff0000, 0xff0100ff00ff0001, + 0xff0100ff00ff0100, 0xff0100ff0000ff01, 0xff0100ff00000000, 0xff0100ff000001ff, + 0xff0100ff00000101, 0xff0100ff00010001, 0xff0100ff01ff0000, 0xff0100ff0100ff00, + 0xff0100ff010000ff, 0xff0100ff01000100, 0xff0100ff0101ff00, 0xff0100ff01010000, + 0xff010000ffff0100, 0xff010000ff000000, 0xff010000ff01ff00, 0xff010000ff010100, + 0xff01000000ffffff, 0xff01000000ff0000, 0xff01000000ff01ff, 0xff0100000000ff00, + 0xff010000000000ff, 0xff01000000000000, 0xff01000000000100, 0xff0100000001ff01, + 0xff01000000010000, 0xff010000000101ff, 0xff01000001ff0100, 0xff0100000100ffff, + 0xff010000010000ff, 0xff01000001000000, 0xff010000010001ff, 0xff01000001000101, + 0xff0100000101ff00, 0xff010000010100ff, 0xff01000001010001, 0xff01000001010100, + 0xff010001ffff0000, 0xff010001ff00ffff, 0xff010001ff00ff01, 0xff010001ff000100, + 0xff010001ff010000, 0xff01000100ffff00, 0xff01000100ff0100, 0xff01000100000000, + 0xff0100010001ffff, 0xff0100010001ff00, 0xff01000100010100, 0xff01000101ff00ff, + 0xff01000101ff0001, 0xff0100010100ffff, 0xff01000101000101, 0xff0101ffffffffff, + 0xff0101ffffffff01, 0xff0101ffffff01ff, 0xff0101ffffff0101, 0xff0101ffff000000, + 0xff0101ffff01ffff, 0xff0101ffff01ff01, 0xff0101ffff0101ff, 0xff0101ffff010101, + 0xff0101ff00ff0000, 0xff0101ff0000ff00, 0xff0101ff000000ff, 0xff0101ff00010000, + 0xff0101ff01ffffff, 0xff0101ff01ffff01, 0xff0101ff01ff01ff, 0xff0101ff01ff0101, + 0xff0101ff0101ffff, 0xff0101ff0101ff01, 0xff0101ff010101ff, 0xff0101ff01010101, + 0xff010100ffff0100, 0xff010100ff00ff00, 0xff010100ff0000ff, 0xff010100ff000100, + 0xff010100ff010000, 0xff01010000ff0001, 0xff01010000ff0100, 0xff0101000000ff01, + 0xff01010000000000, 0xff0101000001ff00, 0xff010100000100ff, 0xff01010000010001, + 0xff01010000010100, 0xff01010001ff0000, 0xff0101000100ffff, 0xff01010001000001, + 0xff01010001000100, 0xff010100010100ff, 0xff01010001010000, 0xff010101ffffffff, + 0xff010101ffffff01, 0xff010101ffff01ff, 0xff010101ffff0101, 0xff010101ff01ffff, + 0xff010101ff01ff01, 0xff010101ff0101ff, 0xff010101ff010101, 0xff01010100ff0000, + 0xff0101010000ff00, 0xff01010100000001, 0xff01010100000100, 0xff01010100010000, + 0xff01010101ffffff, 0xff01010101ffff01, 0xff01010101ff01ff, 0xff01010101ff0101, + 0xff01010101000000, 0xff0101010101ffff, 0xff0101010101ff01, 0xff010101010101ff, + 0xff01010101010101, 0x00ffffffffff0000, 0x00ffffffff00ff00, 0x00ffffffff000001, + 0x00ffffffff010000, 0x00ffffff00ff0100, 0x00ffffff0000ff01, 0x00ffffff00000000, + 0x00ffffff000001ff, 0x00ffffff00000101, 0x00ffffff0001ff00, 0x00ffffff000100ff, + 0x00ffffff00010001, 0x00ffffff010000ff, 0x00ffffff01000100, 0x00ffffff0101ff00, + 0x00ffffff01010001, 0x00ffff00ffffffff, 0x00ffff00ffffff00, 0x00ffff00ffff00ff, + 0x00ffff00ffff0001, 0x00ffff00ffff0100, 0x00ffff00ff00ff01, 0x00ffff00ff000000, + 0x00ffff00ff000001, 0x00ffff00ff0001ff, 0x00ffff00ff000101, 0x00ffff00ff01ff00, + 0x00ffff00ff010001, 0x00ffff00ff010100, 0x00ffff0000ff0000, 0x00ffff0000ff01ff, + 0x00ffff0000ff0101, 0x00ffff000000ff00, 0x00ffff00000000ff, 0x00ffff0000000000, + 0x00ffff0000000001, 0x00ffff0000000100, 0x00ffff0000000101, 0x00ffff0000010000, + 0x00ffff00000101ff, 0x00ffff0000010101, 0x00ffff0001ffff00, 0x00ffff0001ff00ff, + 0x00ffff0001ff0001, 0x00ffff000100ffff, 0x00ffff000100ff01, 0x00ffff0001000000, + 0x00ffff000101ffff, 0x00ffff000101ff00, 0x00ffff000101ff01, 0x00ffff01ffff0000, + 0x00ffff01ff00ff00, 0x00ffff01ff0000ff, 0x00ffff01ff000001, 0x00ffff01ff010000, + 0x00ffff0100ffff00, 0x00ffff010000ff01, 0x00ffff0100000000, 0x00ffff0100000101, + 0x00ffff01000100ff, 0x00ffff0100010100, 0x00ffff0101ff0100, 0x00ffff01010000ff, + 0x00ffff0101010000, 0x00ff00ffffffff00, 0x00ff00ffff000000, 0x00ff00ffff000100, + 0x00ff00ffff010100, 0x00ff00ff00ff0000, 0x00ff00ff00ff01ff, 0x00ff00ff00ff0101, + 0x00ff00ff0000ff00, 0x00ff00ff000000ff, 0x00ff00ff00000000, 0x00ff00ff00000001, + 0x00ff00ff0001ff00, 0x00ff00ff0001ff01, 0x00ff00ff00010000, 0x00ff00ff000101ff, + 0x00ff00ff00010101, 0x00ff00ff01ffff00, 0x00ff00ff01ff0001, 0x00ff00ff01ff0100, + 0x00ff00ff0100ffff, 0x00ff00ff0100ff01, 0x00ff00ff01000000, 0x00ff00ff0101ffff, + 0x00ff00ff0101ff00, 0x00ff00ff01010100, 0x00ff0000ffffff00, 0x00ff0000ffffff01, + 0x00ff0000ffff0000, 0x00ff0000ffff0101, 0x00ff0000ff00ff00, 0x00ff0000ff0000ff, + 0x00ff0000ff000000, 0x00ff0000ff000001, 0x00ff0000ff000100, 0x00ff0000ff01ffff, + 0x00ff0000ff010000, 0x00ff0000ff010101, 0x00ff000000ffff00, 0x00ff000000ff00ff, + 0x00ff000000ff0000, 0x00ff000000ff0001, 0x00ff000000ff0100, 0x00ff00000000ffff, + 0x00ff00000000ff00, 0x00ff0000000000ff, 0x00ff000000000000, 0x00ff000000000001, + 0x00ff0000000001ff, 0x00ff000000000100, 0x00ff00000001ff00, 0x00ff0000000100ff, + 0x00ff000000010000, 0x00ff000000010001, 0x00ff000000010100, 0x00ff000001ffff01, + 0x00ff000001ff00ff, 0x00ff000001ff0000, 0x00ff000001ff01ff, 0x00ff00000100ff00, + 0x00ff0000010000ff, 0x00ff000001000000, 0x00ff000001000001, 0x00ff000001000100, + 0x00ff000001000101, 0x00ff000001010000, 0x00ff0000010101ff, 0x00ff000001010101, + 0x00ff0001ffffff00, 0x00ff0001ffff0000, 0x00ff0001ffff0100, 0x00ff0001ff0000ff, + 0x00ff0001ff000000, 0x00ff0001ff0001ff, 0x00ff0001ff000101, 0x00ff0001ff01ff00, + 0x00ff0001ff0100ff, 0x00ff0001ff010100, 0x00ff000100ffffff, 0x00ff000100ffff01, + 0x00ff000100ff0000, 0x00ff000100ff01ff, 0x00ff00010000ffff, 0x00ff00010000ff00, + 0x00ff00010000ff01, 0x00ff000100000000, 0x00ff000100000001, 0x00ff000100000100, + 0x00ff00010001ff01, 0x00ff000100010000, 0x00ff0001000101ff, 0x00ff000101ffff00, + 0x00ff000101ff0000, 0x00ff000101ff0101, 0x00ff0001010000ff, 0x00ff000101000000, + 0x00ff00010101ff00, 0x00ff0001010100ff, 0x00ff000101010001, 0x00ff01ffffff0000, + 0x00ff01ffff00ff00, 0x00ff01ffff000000, 0x00ff01ffff000101, 0x00ff01ffff010000, + 0x00ff01ff00ffff01, 0x00ff01ff00ff0100, 0x00ff01ff0000ffff, 0x00ff01ff00000000, + 0x00ff01ff000001ff, 0x00ff01ff0001ff00, 0x00ff01ff000100ff, 0x00ff01ff00010001, + 0x00ff01ff00010100, 0x00ff01ff01ff0000, 0x00ff01ff0100ff00, 0x00ff01ff010000ff, + 0x00ff01ff01000001, 0x00ff01ff01000100, 0x00ff01ff01010000, 0x00ff0100ffffff00, + 0x00ff0100ffff0000, 0x00ff0100ffff0001, 0x00ff0100ffff0101, 0x00ff0100ff00ffff, + 0x00ff0100ff0000ff, 0x00ff0100ff000000, 0x00ff0100ff0001ff, 0x00ff0100ff01ff00, + 0x00ff0100ff0100ff, 0x00ff0100ff010001, 0x00ff010000ffffff, 0x00ff010000ff0000, + 0x00ff010000ff0101, 0x00ff01000000ff00, 0x00ff01000000ff01, 0x00ff0100000000ff, + 0x00ff010000000000, 0x00ff010000000001, 0x00ff010000000100, 0x00ff01000001ffff, + 0x00ff01000001ff01, 0x00ff010000010000, 0x00ff010000010001, 0x00ff010000010101, + 0x00ff010001ff0001, 0x00ff010001ff0100, 0x00ff01000100ff01, 0x00ff010001000000, + 0x00ff010001000001, 0x00ff0100010001ff, 0x00ff01000101ff00, 0x00ff0100010100ff, + 0x00ff010001010001, 0x00ff010001010100, 0x00ff0101ff000001, 0x00ff010100ff00ff, + 0x00ff010100ff0001, 0x00ff010100ff0100, 0x00ff010100000000, 0x00ff0101000001ff, + 0x00ff010100000101, 0x00ff0101000100ff, 0x00ff010100010100, 0x00ff0101010000ff, + 0x00ff010101010000, 0x0000ffffffffff00, 0x0000ffffffff00ff, 0x0000ffffffff0000, + 0x0000ffffffff0001, 0x0000ffffffff0100, 0x0000ffffff00ff01, 0x0000ffffff000000, + 0x0000ffffff000101, 0x0000ffffff01ff00, 0x0000ffffff0100ff, 0x0000ffffff010100, + 0x0000ffff00ffffff, 0x0000ffff00ff0000, 0x0000ffff00ff01ff, 0x0000ffff0000ff00, + 0x0000ffff000000ff, 0x0000ffff00000000, 0x0000ffff00000001, 0x0000ffff00000100, + 0x0000ffff00010000, 0x0000ffff000101ff, 0x0000ffff01ff0001, 0x0000ffff01ff0100, + 0x0000ffff01000000, 0x0000ffff010001ff, 0x0000ffff0101ffff, 0x0000ffff0101ff00, + 0x0000ffff01010001, 0x0000ffff01010100, 0x0000ff00ffff0000, 0x0000ff00ffff01ff, + 0x0000ff00ffff0100, 0x0000ff00ffff0101, 0x0000ff00ff00ff00, 0x0000ff00ff0000ff, + 0x0000ff00ff000000, 0x0000ff00ff000001, 0x0000ff00ff0001ff, 0x0000ff00ff000100, + 0x0000ff00ff01ffff, 0x0000ff00ff010000, 0x0000ff00ff010001, 0x0000ff00ff0101ff, + 0x0000ff00ff010101, 0x0000ff0000ffff00, 0x0000ff0000ff00ff, 0x0000ff0000ff0000, + 0x0000ff0000ff0001, 0x0000ff0000ff0100, 0x0000ff000000ffff, 0x0000ff000000ff00, + 0x0000ff000000ff01, 0x0000ff00000000ff, 0x0000ff0000000000, 0x0000ff0000000001, + 0x0000ff00000001ff, 0x0000ff0000000100, 0x0000ff0000000101, 0x0000ff000001ff00, + 0x0000ff00000100ff, 0x0000ff0000010000, 0x0000ff0000010001, 0x0000ff0000010100, + 0x0000ff0001ffff01, 0x0000ff0001ff0000, 0x0000ff000100ff00, 0x0000ff00010000ff, + 0x0000ff0001000000, 0x0000ff0001000001, 0x0000ff0001000100, 0x0000ff000101ffff, + 0x0000ff0001010000, 0x0000ff0001010101, 0x0000ff01ffffff00, 0x0000ff01ffff0001, + 0x0000ff01ff00ff01, 0x0000ff01ff000000, 0x0000ff01ff000101, 0x0000ff01ff01ff00, + 0x0000ff01ff0100ff, 0x0000ff0100ffff01, 0x0000ff0100ff0000, 0x0000ff0100ff0101, + 0x0000ff010000ff00, 0x0000ff01000000ff, 0x0000ff0100000000, 0x0000ff0100000001, + 0x0000ff0100000100, 0x0000ff010001ff01, 0x0000ff0100010000, 0x0000ff0101ff0000, + 0x0000ff010100ffff, 0x0000ff010100ff01, 0x0000ff0101000000, 0x0000ff0101000100, + 0x0000ff0101000101, 0x0000ff01010100ff, 0x000000ffffff00ff, 0x000000ffffff0000, + 0x000000ffff00ff00, 0x000000ffff0000ff, 0x000000ffff000000, 0x000000ffff000001, + 0x000000ffff0001ff, 0x000000ffff000100, 0x000000ffff01ff00, 0x000000ffff010000, + 0x000000ffff0101ff, 0x000000ffff010101, 0x000000ff00ffff00, 0x000000ff00ff00ff, + 0x000000ff00ff0000, 0x000000ff00ff0001, 0x000000ff00ff0100, 0x000000ff00ff0101, + 0x000000ff0000ffff, 0x000000ff0000ff00, 0x000000ff000000ff, 0x000000ff00000000, + 0x000000ff00000001, 0x000000ff000001ff, 0x000000ff00000100, 0x000000ff00000101, + 0x000000ff0001ff00, 0x000000ff0001ff01, 0x000000ff000100ff, 0x000000ff00010000, + 0x000000ff00010001, 0x000000ff00010100, 0x000000ff01ffffff, 0x000000ff01ff01ff, + 0x000000ff01ff0101, 0x000000ff0100ff00, 0x000000ff010000ff, 0x000000ff01000000, + 0x000000ff01000001, 0x000000ff01000100, 0x000000ff0101ff00, 0x000000ff010100ff, + 0x000000ff01010000, 0x000000ff01010101, 0x00000000ffffff00, 0x00000000ffffff01, + 0x00000000ffff00ff, 0x00000000ffff0000, 0x00000000ffff0001, 0x00000000ffff0100, + 0x00000000ff00ffff, 0x00000000ff00ff00, 0x00000000ff00ff01, 0x00000000ff0000ff, + 0x00000000ff000000, 0x00000000ff000001, 0x00000000ff000100, 0x00000000ff000101, + 0x00000000ff01ff00, 0x00000000ff0100ff, 0x00000000ff010000, 0x00000000ff010001, + 0x00000000ff010100, 0x0000000000ffffff, 0x0000000000ffff00, 0x0000000000ffff01, + 0x0000000000ff00ff, 0x0000000000ff0000, 0x0000000000ff0001, 0x0000000000ff01ff, + 0x0000000000ff0100, 0x000000000000ffff, 0x000000000000ff00, 0x000000000000ff01, + 0x00000000000000ff, 0x0000000000000000, 0x0000000000000001, 0x00000000000001ff, + 0x0000000000000100, 0x0000000000000101, 0x000000000001ffff, 0x000000000001ff00, + 0x00000000000100ff, 0x0000000000010000, 0x0000000000010001, 0x00000000000101ff, + 0x0000000000010100, 0x0000000000010101, 0x0000000001ffff00, 0x0000000001ff00ff, + 0x0000000001ff0000, 0x0000000001ff0100, 0x0000000001ff0101, 0x000000000100ffff, + 0x000000000100ff00, 0x00000000010000ff, 0x0000000001000000, 0x0000000001000001, + 0x00000000010001ff, 0x0000000001000100, 0x000000000101ff00, 0x00000000010100ff, + 0x0000000001010000, 0x0000000001010001, 0x0000000001010100, 0x00000001ffffffff, + 0x00000001ffffff00, 0x00000001ffffff01, 0x00000001ffff00ff, 0x00000001ffff0001, + 0x00000001ffff01ff, 0x00000001ffff0100, 0x00000001ff00ff00, 0x00000001ff0000ff, + 0x00000001ff000000, 0x00000001ff0001ff, 0x00000001ff000100, 0x00000001ff01ffff, + 0x00000001ff01ff00, 0x00000001ff01ff01, 0x00000001ff0100ff, 0x00000001ff010000, + 0x00000001ff010001, 0x00000001ff0101ff, 0x00000001ff010100, 0x0000000100ffff00, + 0x0000000100ff0000, 0x0000000100ff0001, 0x0000000100ff01ff, 0x0000000100ff0100, + 0x0000000100ff0101, 0x000000010000ffff, 0x000000010000ff00, 0x000000010000ff01, + 0x00000001000000ff, 0x0000000100000000, 0x0000000100000001, 0x00000001000001ff, + 0x0000000100000100, 0x0000000100000101, 0x000000010001ff00, 0x00000001000100ff, + 0x0000000100010000, 0x0000000100010100, 0x0000000101ffff01, 0x0000000101ff0000, + 0x0000000101ff0001, 0x0000000101ff01ff, 0x0000000101ff0100, 0x0000000101ff0101, + 0x000000010100ff00, 0x0000000101000000, 0x0000000101000101, 0x000000010101ff01, + 0x0000000101010000, 0x0000000101010001, 0x00000001010101ff, 0x0000000101010100, + 0x000001ffffff00ff, 0x000001ffffff0000, 0x000001ffffff0001, 0x000001ffffff0100, + 0x000001ffff00ffff, 0x000001ffff000000, 0x000001ffff0001ff, 0x000001ffff01ff00, + 0x000001ffff010101, 0x000001ff00ff0000, 0x000001ff00ff01ff, 0x000001ff00ff0101, + 0x000001ff0000ff00, 0x000001ff000000ff, 0x000001ff00000000, 0x000001ff00000001, + 0x000001ff000001ff, 0x000001ff00000100, 0x000001ff0001ffff, 0x000001ff0001ff01, + 0x000001ff000100ff, 0x000001ff00010000, 0x000001ff01ffff01, 0x000001ff01ff0100, + 0x000001ff0100ffff, 0x000001ff0100ff01, 0x000001ff01000000, 0x000001ff010001ff, + 0x000001ff0101ff00, 0x000001ff01010100, 0x00000100ffffff00, 0x00000100ffffff01, + 0x00000100ffff0000, 0x00000100ffff0101, 0x00000100ff00ff00, 0x00000100ff0000ff, + 0x00000100ff000000, 0x00000100ff000001, 0x00000100ff000100, 0x00000100ff010000, + 0x0000010000ffff00, 0x0000010000ff00ff, 0x0000010000ff0000, 0x0000010000ff0001, + 0x0000010000ff0100, 0x000001000000ffff, 0x000001000000ff00, 0x000001000000ff01, + 0x00000100000000ff, 0x0000010000000000, 0x0000010000000001, 0x00000100000001ff, + 0x0000010000000100, 0x0000010000000101, 0x000001000001ff00, 0x00000100000100ff, + 0x0000010000010000, 0x0000010000010001, 0x0000010000010100, 0x0000010001ffff00, + 0x0000010001ff0000, 0x0000010001ff0100, 0x000001000100ff00, 0x00000100010000ff, + 0x0000010001000000, 0x0000010001000001, 0x00000100010001ff, 0x0000010001000100, + 0x0000010001010000, 0x00000101ffff00ff, 0x00000101ffff01ff, 0x00000101ff000000, + 0x00000101ff000101, 0x00000101ff01ffff, 0x00000101ff010000, 0x00000101ff010001, + 0x00000101ff010100, 0x0000010100ff0000, 0x0000010100ff01ff, 0x0000010100ff0100, + 0x000001010000ff00, 0x0000010100000000, 0x0000010100000001, 0x00000101000001ff, + 0x0000010100000100, 0x000001010001ff01, 0x0000010100010000, 0x00000101000101ff, + 0x0000010100010101, 0x0000010101ffff00, 0x0000010101ff0101, 0x000001010100ff01, + 0x0000010101000000, 0x0000010101000001, 0x00000101010001ff, 0x0000010101000101, + 0x000001010101ff00, 0x0001ffffffff0000, 0x0001ffffff0000ff, 0x0001ffffff000001, + 0x0001ffffff000100, 0x0001ffffff010000, 0x0001ffff00ff00ff, 0x0001ffff0000ffff, + 0x0001ffff00000000, 0x0001ffff00000001, 0x0001ffff000001ff, 0x0001ffff00000101, + 0x0001ffff0001ff00, 0x0001ffff000100ff, 0x0001ffff00010001, 0x0001ffff00010100, + 0x0001ffff01ffff00, 0x0001ffff01000001, 0x0001ffff01010000, 0x0001ff00ffffff00, + 0x0001ff00ffff00ff, 0x0001ff00ffff0001, 0x0001ff00ffff0100, 0x0001ff00ff00ff01, + 0x0001ff00ff000000, 0x0001ff00ff01ff00, 0x0001ff00ff01ff01, 0x0001ff00ff010001, + 0x0001ff00ff010100, 0x0001ff0000ff0000, 0x0001ff0000ff0100, 0x0001ff000000ff00, + 0x0001ff0000000000, 0x0001ff0000000001, 0x0001ff0000000100, 0x0001ff0000010000, + 0x0001ff0000010001, 0x0001ff0000010101, 0x0001ff0001ff00ff, 0x0001ff0001ff0101, + 0x0001ff000100ff01, 0x0001ff0001000000, 0x0001ff000101ff00, 0x0001ff0001010001, + 0x0001ff0001010100, 0x0001ff01ff00ff00, 0x0001ff01ff000001, 0x0001ff01ff000100, + 0x0001ff0100ffffff, 0x0001ff0100ffff00, 0x0001ff0100ff0001, 0x0001ff0100000000, + 0x0001ff0100000001, 0x0001ff01000001ff, 0x0001ff010001ffff, 0x0001ff0101ff0000, + 0x0001ff010100ff00, 0x0001ff0101000001, 0x0001ff0101010000, 0x000100ffff00ff00, + 0x000100ffff00ff01, 0x000100ffff000000, 0x000100ffff000001, 0x000100ffff000101, + 0x000100ffff01ff00, 0x000100ffff010001, 0x000100ffff010100, 0x000100ff00ffffff, + 0x000100ff00ffff01, 0x000100ff00ff0000, 0x000100ff00ff01ff, 0x000100ff00ff0101, + 0x000100ff0000ff00, 0x000100ff000000ff, 0x000100ff00000000, 0x000100ff00000001, + 0x000100ff00000100, 0x000100ff00000101, 0x000100ff0001ffff, 0x000100ff0001ff01, + 0x000100ff00010000, 0x000100ff01ff00ff, 0x000100ff01ff0000, 0x000100ff01ff0100, + 0x000100ff0100ffff, 0x000100ff0100ff01, 0x000100ff010000ff, 0x000100ff01000000, + 0x000100ff01000001, 0x000100ff010001ff, 0x000100ff01000101, 0x000100ff0101ff00, + 0x000100ff010100ff, 0x000100ff01010100, 0x00010000ffff0000, 0x00010000ffff01ff, + 0x00010000ffff0101, 0x00010000ff00ff00, 0x00010000ff000000, 0x00010000ff000001, + 0x00010000ff000100, 0x0001000000ff00ff, 0x0001000000ff0000, 0x0001000000ff0001, + 0x0001000000ff0100, 0x000100000000ffff, 0x000100000000ff00, 0x00010000000000ff, + 0x0001000000000000, 0x0001000000000001, 0x0001000000000100, 0x000100000001ff00, + 0x00010000000100ff, 0x0001000000010000, 0x0001000000010001, 0x0001000000010100, + 0x0001000001ff0001, 0x0001000001ff0100, 0x0001000001ff0101, 0x000100000100ff00, + 0x0001000001000000, 0x0001000001000001, 0x0001000001000100, 0x0001000001000101, + 0x000100000101ff01, 0x0001000001010000, 0x0001000001010001, 0x00010000010101ff, + 0x00010001ffffff01, 0x00010001ffff0100, 0x00010001ff000000, 0x00010001ff01ffff, + 0x00010001ff010001, 0x00010001ff0101ff, 0x00010001ff010100, 0x0001000100ffffff, + 0x0001000100ff0000, 0x0001000100ff01ff, 0x0001000100ff0101, 0x000100010000ff00, + 0x00010001000000ff, 0x0001000100000000, 0x0001000100000001, 0x00010001000001ff, + 0x0001000100000101, 0x000100010001ffff, 0x0001000100010000, 0x00010001000101ff, + 0x0001000101ffffff, 0x0001000101ffff01, 0x0001000101ff0000, 0x0001000101ff0101, + 0x00010001010000ff, 0x0001000101000001, 0x00010001010001ff, 0x0001000101000100, + 0x000100010101ffff, 0x00010001010100ff, 0x0001000101010001, 0x0001000101010101, + 0x000101ffff000001, 0x000101ffff000100, 0x000101ffff010000, 0x000101ff00ffff00, + 0x000101ff0000ff01, 0x000101ff00000000, 0x000101ff00000101, 0x000101ff0001ff00, + 0x000101ff00010100, 0x000101ff01ff0000, 0x000101ff0100ff00, 0x000101ff010001ff, + 0x000101ff01010001, 0x00010100ffffff00, 0x00010100ffff00ff, 0x00010100ff00ffff, + 0x00010100ff000000, 0x00010100ff01ff00, 0x00010100ff0100ff, 0x00010100ff010001, + 0x00010100ff010100, 0x0001010000ffffff, 0x0001010000ffff00, 0x0001010000ff0000, + 0x0001010000ff0001, 0x0001010000ff01ff, 0x000101000000ff00, 0x00010100000000ff, + 0x0001010000000000, 0x0001010000000001, 0x0001010000000100, 0x000101000001ffff, + 0x0001010000010000, 0x0001010000010101, 0x0001010001ffff01, 0x0001010001ff00ff, + 0x0001010001ff0101, 0x0001010001000000, 0x000101000101ff00, 0x00010100010100ff, + 0x0001010001010000, 0x0001010001010100, 0x00010101ff00ff00, 0x00010101ff000001, + 0x00010101ff0001ff, 0x0001010100ffff00, 0x0001010100ff00ff, 0x0001010100ff0100, + 0x000101010000ffff, 0x0001010100000000, 0x00010101000001ff, 0x0001010100000101, + 0x00010101000100ff, 0x0001010100010000, 0x0001010100010100, 0x0001010101ff0001, + 0x00010101010000ff, 0x00010101010001ff, 0x0001010101000101, 0x0001010101010001, + 0x01ffffffffffffff, 0x01ffffffffffff01, 0x01ffffffffff01ff, 0x01ffffffffff0101, + 0x01ffffffff01ffff, 0x01ffffffff01ff01, 0x01ffffffff0101ff, 0x01ffffffff010101, + 0x01ffffff00ff0000, 0x01ffffff0000ffff, 0x01ffffff0000ff00, 0x01ffffff000000ff, + 0x01ffffff00000001, 0x01ffffff00000100, 0x01ffffff00010000, 0x01ffffff01ffffff, + 0x01ffffff01ffff01, 0x01ffffff01ff01ff, 0x01ffffff01ff0101, 0x01ffffff01000000, + 0x01ffffff0101ffff, 0x01ffffff0101ff01, 0x01ffffff010101ff, 0x01ffffff01010101, + 0x01ffff00ffff0000, 0x01ffff00ff00ff00, 0x01ffff00ff0000ff, 0x01ffff00ff000001, + 0x01ffff00ff000100, 0x01ffff00ff010000, 0x01ffff0000ffff00, 0x01ffff0000ff00ff, + 0x01ffff0000ff0100, 0x01ffff000000ffff, 0x01ffff000000ff01, 0x01ffff0000000000, + 0x01ffff0000000001, 0x01ffff00000001ff, 0x01ffff0000000100, 0x01ffff00000100ff, + 0x01ffff0000010001, 0x01ffff0000010100, 0x01ffff0001ff0000, 0x01ffff0001ff0100, + 0x01ffff00010000ff, 0x01ffff0001000001, 0x01ffff0001000100, 0x01ffff0001010000, + 0x01ffff01ffffffff, 0x01ffff01ffffff01, 0x01ffff01ffff01ff, 0x01ffff01ffff0101, + 0x01ffff01ff000000, 0x01ffff01ff01ffff, 0x01ffff01ff01ff01, 0x01ffff01ff0101ff, + 0x01ffff01ff010101, 0x01ffff010000ff00, 0x01ffff01000000ff, 0x01ffff0100000100, + 0x01ffff0100010000, 0x01ffff0101ffffff, 0x01ffff0101ffff01, 0x01ffff0101ff01ff, + 0x01ffff0101ff0101, 0x01ffff0101000000, 0x01ffff010101ffff, 0x01ffff010101ff01, + 0x01ffff01010101ff, 0x01ffff0101010101, 0x01ff00ffff0000ff, 0x01ff00ffff000100, + 0x01ff00ff00ffff00, 0x01ff00ff00ff00ff, 0x01ff00ff0000ff00, 0x01ff00ff00000000, + 0x01ff00ff00000101, 0x01ff00ff0001ff00, 0x01ff00ff000100ff, 0x01ff00ff00010100, + 0x01ff00ff010000ff, 0x01ff00ff01000100, 0x01ff0000ffffff00, 0x01ff0000ffff0100, + 0x01ff0000ff00ff01, 0x01ff0000ff000000, 0x01ff0000ff000101, 0x01ff0000ff010001, + 0x01ff0000ff010100, 0x01ff000000ffffff, 0x01ff000000ffff00, 0x01ff000000ff0000, + 0x01ff000000ff01ff, 0x01ff00000000ff00, 0x01ff0000000000ff, 0x01ff000000000000, + 0x01ff000000000001, 0x01ff000000000100, 0x01ff000000000101, 0x01ff000000010000, + 0x01ff000000010001, 0x01ff0000000101ff, 0x01ff000000010101, 0x01ff000001ffff00, + 0x01ff000001ff00ff, 0x01ff000001ff0001, 0x01ff000001ff0100, 0x01ff00000100ffff, + 0x01ff00000100ff01, 0x01ff000001000000, 0x01ff0000010001ff, 0x01ff000001010001, + 0x01ff0001ff00ff00, 0x01ff0001ff000001, 0x01ff0001ff000100, 0x01ff0001ff010000, + 0x01ff000100ffff00, 0x01ff000100ff00ff, 0x01ff000100ff0100, 0x01ff000100ff0101, + 0x01ff00010000ffff, 0x01ff000100000000, 0x01ff000100000100, 0x01ff000100000101, + 0x01ff00010001ff00, 0x01ff000100010001, 0x01ff000100010101, 0x01ff000101ff0000, + 0x01ff00010100ff00, 0x01ff000101000101, 0x01ff0001010100ff, 0x01ff01ffffffffff, + 0x01ff01ffffffff01, 0x01ff01ffffff01ff, 0x01ff01ffffff0101, 0x01ff01ffff000000, + 0x01ff01ffff01ffff, 0x01ff01ffff01ff01, 0x01ff01ffff0101ff, 0x01ff01ffff010101, + 0x01ff01ff00ffff00, 0x01ff01ff00ff0000, 0x01ff01ff0000ff00, 0x01ff01ff000000ff, + 0x01ff01ff00000100, 0x01ff01ff00010000, 0x01ff01ff00010100, 0x01ff01ff01ffffff, + 0x01ff01ff01ffff01, 0x01ff01ff01ff01ff, 0x01ff01ff01ff0101, 0x01ff01ff01000000, + 0x01ff01ff0101ffff, 0x01ff01ff0101ff01, 0x01ff01ff010101ff, 0x01ff01ff01010101, + 0x01ff0100ffff0000, 0x01ff0100ffff0001, 0x01ff0100ff00ff00, 0x01ff0100ff0000ff, + 0x01ff0100ff000001, 0x01ff0100ff010000, 0x01ff010000ffff00, 0x01ff010000ff00ff, + 0x01ff010000ff0001, 0x01ff010000ff0100, 0x01ff01000000ffff, 0x01ff01000000ff01, + 0x01ff010000000000, 0x01ff010000000101, 0x01ff01000001ff00, 0x01ff0100000100ff, + 0x01ff010001ff0000, 0x01ff010001000001, 0x01ff010001000100, 0x01ff010001010000, + 0x01ff0101ffffffff, 0x01ff0101ffffff01, 0x01ff0101ffff01ff, 0x01ff0101ffff0101, + 0x01ff0101ff000000, 0x01ff0101ff01ffff, 0x01ff0101ff01ff01, 0x01ff0101ff0101ff, + 0x01ff0101ff010101, 0x01ff010100ff0000, 0x01ff01010000ff00, 0x01ff0101000000ff, + 0x01ff010100000001, 0x01ff010101ffffff, 0x01ff010101ffff01, 0x01ff010101ff01ff, + 0x01ff010101ff0101, 0x01ff010101000000, 0x01ff01010101ffff, 0x01ff01010101ff01, + 0x01ff0101010101ff, 0x01ff010101010101, 0x0100ffffffff0000, 0x0100ffffff00ff00, + 0x0100ffffff000001, 0x0100ffffff0001ff, 0x0100ffffff000100, 0x0100ffffff010000, + 0x0100ffff00ffff00, 0x0100ffff00ff0001, 0x0100ffff00ff0100, 0x0100ffff00000000, + 0x0100ffff000001ff, 0x0100ffff00000101, 0x0100ffff00010100, 0x0100ffff00010101, + 0x0100ffff01ff0000, 0x0100ffff0100ff00, 0x0100ffff010000ff, 0x0100ffff01000001, + 0x0100ffff01000100, 0x0100ffff01010000, 0x0100ff00ffffff00, 0x0100ff00ffff00ff, + 0x0100ff00ffff0001, 0x0100ff00ffff0100, 0x0100ff00ff00ffff, 0x0100ff00ff000000, + 0x0100ff00ff0001ff, 0x0100ff00ff000101, 0x0100ff00ff01ff00, 0x0100ff00ff0100ff, + 0x0100ff00ff010001, 0x0100ff00ff010100, 0x0100ff0000ffffff, 0x0100ff0000ff0000, + 0x0100ff000000ffff, 0x0100ff000000ff00, 0x0100ff00000000ff, 0x0100ff0000000000, + 0x0100ff0000000001, 0x0100ff0000000100, 0x0100ff000001ff01, 0x0100ff0000010000, + 0x0100ff0001ff00ff, 0x0100ff0001ff0001, 0x0100ff000100ff01, 0x0100ff0001000000, + 0x0100ff00010001ff, 0x0100ff000101ff00, 0x0100ff00010100ff, 0x0100ff0001010001, + 0x0100ff0001010100, 0x0100ff01ffff0000, 0x0100ff01ff00ff00, 0x0100ff01ff0000ff, + 0x0100ff01ff000100, 0x0100ff01ff010000, 0x0100ff0100ff00ff, 0x0100ff0100ff0001, + 0x0100ff0100ff0100, 0x0100ff010000ffff, 0x0100ff010000ff01, 0x0100ff0100000000, + 0x0100ff01000001ff, 0x0100ff0100010001, 0x0100ff0100010100, 0x0100ff0101ff0000, + 0x0100ff01010000ff, 0x0100ff0101000001, 0x0100ff0101010100, 0x010000ffffffff00, + 0x010000ffffff00ff, 0x010000ffffff0001, 0x010000ffff00ffff, 0x010000ffff000000, + 0x010000ffff0001ff, 0x010000ffff010001, 0x010000ff00ffffff, 0x010000ff00ff0101, + 0x010000ff0000ff00, 0x010000ff000000ff, 0x010000ff00000000, 0x010000ff00000001, + 0x010000ff000001ff, 0x010000ff00000100, 0x010000ff0001ffff, 0x010000ff0001ff00, + 0x010000ff0001ff01, 0x010000ff00010000, 0x010000ff01ff00ff, 0x010000ff01ff0001, + 0x010000ff0100ff01, 0x010000ff010000ff, 0x010000ff01000000, 0x010000ff010001ff, + 0x010000ff0101ff00, 0x010000ff01010100, 0x01000000ffffffff, 0x01000000ffff0000, + 0x01000000ffff01ff, 0x01000000ffff0101, 0x01000000ff00ffff, 0x01000000ff00ff00, + 0x01000000ff0000ff, 0x01000000ff000000, 0x01000000ff000001, 0x01000000ff000100, + 0x01000000ff01ff00, 0x01000000ff010000, 0x01000000ff010100, 0x01000000ff010101, + 0x0100000000ffff00, 0x0100000000ff00ff, 0x0100000000ff0000, 0x0100000000ff0001, + 0x0100000000ff0100, 0x010000000000ffff, 0x010000000000ff00, 0x010000000000ff01, + 0x01000000000000ff, 0x0100000000000000, 0x0100000000000001, 0x01000000000001ff, + 0x0100000000000100, 0x0100000000000101, 0x010000000001ff00, 0x01000000000100ff, + 0x0100000000010000, 0x0100000000010001, 0x0100000000010100, 0x0100000001ffff00, + 0x0100000001ff0000, 0x0100000001ff01ff, 0x010000000100ff00, 0x010000000100ff01, + 0x01000000010000ff, 0x0100000001000000, 0x0100000001000001, 0x0100000001000100, + 0x0100000001000101, 0x010000000101ffff, 0x010000000101ff01, 0x0100000001010000, + 0x01000000010101ff, 0x0100000001010101, 0x01000001ffffff00, 0x01000001ffff00ff, + 0x01000001ff00ffff, 0x01000001ff000000, 0x01000001ff000100, 0x01000001ff01ffff, + 0x01000001ff010001, 0x01000001ff010100, 0x0100000100ff0000, 0x0100000100ff01ff, + 0x0100000100ff0100, 0x010000010000ff00, 0x010000010000ff01, 0x0100000100000000, + 0x0100000100000001, 0x0100000100000100, 0x0100000100010000, 0x01000001000101ff, + 0x0100000101ffff01, 0x0100000101ff00ff, 0x0100000101ff0100, 0x0100000101ff0101, + 0x010000010100ff01, 0x01000001010000ff, 0x0100000101000000, 0x01000001010100ff, + 0x0100000101010001, 0x0100000101010100, 0x010001ffffff0000, 0x010001ffff000001, + 0x010001ffff000100, 0x010001ffff010000, 0x010001ff00ffff00, 0x010001ff00ff0001, + 0x010001ff0000ffff, 0x010001ff0000ff01, 0x010001ff00000000, 0x010001ff00000001, + 0x010001ff00000101, 0x010001ff000100ff, 0x010001ff00010000, 0x010001ff01ff0000, + 0x010001ff0100ff00, 0x010001ff01000001, 0x010001ff01000100, 0x010001ff01010000, + 0x01000100ffff00ff, 0x01000100ffff0001, 0x01000100ffff0100, 0x01000100ff00ffff, + 0x01000100ff00ff01, 0x01000100ff000000, 0x01000100ff0001ff, 0x01000100ff000101, + 0x01000100ff01ffff, 0x01000100ff01ff00, 0x01000100ff0100ff, 0x01000100ff010001, + 0x0100010000ffffff, 0x0100010000ffff01, 0x0100010000ff0000, 0x0100010000ff01ff, + 0x0100010000ff0101, 0x010001000000ff00, 0x01000100000000ff, 0x0100010000000000, + 0x0100010000000001, 0x0100010000000100, 0x010001000001ff01, 0x0100010000010000, + 0x0100010000010001, 0x0100010000010101, 0x0100010001ffff00, 0x0100010001ff00ff, + 0x010001000100ffff, 0x010001000100ff01, 0x0100010001000000, 0x0100010001000101, + 0x010001000101ff00, 0x0100010001010001, 0x01000101ffff0000, 0x01000101ff000000, + 0x01000101ff010000, 0x0100010100ff00ff, 0x0100010100ff0001, 0x0100010100ff0100, + 0x010001010000ffff, 0x0100010100000000, 0x01000101000001ff, 0x010001010001ff00, + 0x0100010101ff0000, 0x010001010100ff00, 0x01000101010000ff, 0x0100010101000000, + 0x0100010101000001, 0x0101ffffffffffff, 0x0101ffffffffff01, 0x0101ffffffff01ff, + 0x0101ffffffff0101, 0x0101ffffff000000, 0x0101ffffff01ffff, 0x0101ffffff01ff01, + 0x0101ffffff0101ff, 0x0101ffffff010101, 0x0101ffff00ff0000, 0x0101ffff0000ff00, + 0x0101ffff000000ff, 0x0101ffff00000001, 0x0101ffff00000100, 0x0101ffff01ffffff, + 0x0101ffff01ffff01, 0x0101ffff01ff01ff, 0x0101ffff01ff0101, 0x0101ffff01000000, + 0x0101ffff0101ffff, 0x0101ffff0101ff01, 0x0101ffff010101ff, 0x0101ffff01010101, + 0x0101ff00ffff0000, 0x0101ff00ffff0100, 0x0101ff00ff00ff00, 0x0101ff00ff0000ff, + 0x0101ff00ff000001, 0x0101ff00ff000100, 0x0101ff00ff000101, 0x0101ff0000ff0001, + 0x0101ff0000ff0100, 0x0101ff000000ff00, 0x0101ff0000000000, 0x0101ff00000001ff, + 0x0101ff0000000101, 0x0101ff000001ff00, 0x0101ff00000100ff, 0x0101ff0001ff0000, + 0x0101ff000100ffff, 0x0101ff000100ff01, 0x0101ff0001000001, 0x0101ff0001000100, + 0x0101ff01ffffff01, 0x0101ff01ffff01ff, 0x0101ff01ffff0101, 0x0101ff01ff00ffff, + 0x0101ff01ff000100, 0x0101ff01ff01ff01, 0x0101ff01ff0101ff, 0x0101ff01ff010101, + 0x0101ff0100ff0000, 0x0101ff010000ff00, 0x0101ff0100000001, 0x0101ff0100000100, + 0x0101ff0100010000, 0x0101ff0101ffffff, 0x0101ff0101ffff01, 0x0101ff0101ff01ff, + 0x0101ff0101ff0101, 0x0101ff0101000000, 0x0101ff010101ffff, 0x0101ff010101ff01, + 0x0101ff01010101ff, 0x0101ff0101010101, 0x010100ffff000100, 0x010100ffff010000, + 0x010100ff00ffff00, 0x010100ff00ff00ff, 0x010100ff0000ffff, 0x010100ff000000ff, + 0x010100ff00000000, 0x010100ff000001ff, 0x010100ff00000101, 0x010100ff0001ff00, + 0x010100ff00010000, 0x010100ff00010001, 0x010100ff000101ff, 0x010100ff00010100, + 0x010100ff01ff0000, 0x01010000ffff0001, 0x01010000ffff0100, 0x01010000ff00ffff, + 0x01010000ff00ff01, 0x01010000ff000000, 0x01010000ff0001ff, 0x01010000ff010001, + 0x01010000ff010100, 0x0101000000ffff01, 0x0101000000ff0000, 0x010100000000ff00, + 0x01010000000000ff, 0x0101000000000000, 0x0101000000000001, 0x0101000000000100, + 0x0101000000010000, 0x0101000000010101, 0x0101000001ffff00, 0x0101000001ff00ff, + 0x0101000001ff0000, 0x0101000001ff0001, 0x0101000001ff0100, 0x010100000100ff01, + 0x0101000001000000, 0x01010000010001ff, 0x01010001ffff0000, 0x01010001ff00ff00, + 0x01010001ff000001, 0x01010001ff000101, 0x01010001ff01ff00, 0x01010001ff010000, + 0x0101000100ff00ff, 0x0101000100ff0001, 0x0101000100ff0101, 0x010100010000ff01, + 0x0101000100000000, 0x0101000100000001, 0x01010001000001ff, 0x010100010001ffff, + 0x010100010001ff01, 0x0101000101ff0001, 0x010100010100ffff, 0x0101000101000000, + 0x0101000101000001, 0x0101000101000100, 0x010100010101ff00, 0x01010001010100ff, + 0x0101000101010001, 0x010101ffffffffff, 0x010101ffffffff01, 0x010101ffffff01ff, + 0x010101ffffff0101, 0x010101ffff01ffff, 0x010101ffff01ff01, 0x010101ffff0101ff, + 0x010101ffff010101, 0x010101ff0000ff00, 0x010101ff000000ff, 0x010101ff00000001, + 0x010101ff00000100, 0x010101ff01ffffff, 0x010101ff01ffff01, 0x010101ff01ff01ff, + 0x010101ff01ff0101, 0x010101ff01000000, 0x010101ff0101ffff, 0x010101ff0101ff01, + 0x010101ff010101ff, 0x010101ff01010101, 0x01010100ffff0000, 0x01010100ff0000ff, + 0x01010100ff000100, 0x01010100ff01ff00, 0x01010100ff010000, 0x0101010000ffff00, + 0x010101000000ffff, 0x0101010000000000, 0x0101010000000101, 0x010101000001ff00, + 0x0101010000010001, 0x0101010000010100, 0x010101000100ffff, 0x0101010001000001, + 0x01010101ffffffff, 0x01010101ffffff01, 0x01010101ffff01ff, 0x01010101ffff0101, + 0x01010101ff01ffff, 0x01010101ff01ff01, 0x01010101ff0101ff, 0x01010101ff010101, + 0x010101010000ff00, 0x01010101000000ff, 0x0101010100000001, 0x0101010101ffffff, + 0x0101010101ffff01, 0x0101010101ff01ff, 0x0101010101ff0101, 0x0101010101000000, + 0x010101010101ffff, 0x010101010101ff01, 0x01010101010101ff, 0x0101010101010101}; + +inline constexpr uint8_t ksigns_iq2xs[128] = { + 0, 129, 130, 3, 132, 5, 6, 135, 136, 9, 10, 139, 12, 141, 142, 15, + 144, 17, 18, 147, 20, 149, 150, 23, 24, 153, 154, 27, 156, 29, 30, 159, + 160, 33, 34, 163, 36, 165, 166, 39, 40, 169, 170, 43, 172, 45, 46, 175, + 48, 177, 178, 51, 180, 53, 54, 183, 184, 57, 58, 187, 60, 189, 190, 63, + 192, 65, 66, 195, 68, 197, 198, 71, 72, 201, 202, 75, 204, 77, 78, 207, + 80, 209, 210, 83, 212, 85, 86, 215, 216, 89, 90, 219, 92, 221, 222, 95, + 96, 225, 226, 99, 228, 101, 102, 231, 232, 105, 106, 235, 108, 237, 238, 111, + 240, 113, 114, 243, 116, 245, 246, 119, 120, 249, 250, 123, 252, 125, 126, 255}; + +inline constexpr uint8_t kmask_iq2xs[8] = { + 1, 2, 4, 8, 16, 32, 64, 128}; + +} // namespace vllm::gguf::cpu_iq diff --git a/csrc/gguf_expert_cache.cpp b/csrc/gguf_expert_cache.cpp index 4ab5282062..dc84cc9d38 100644 --- a/csrc/gguf_expert_cache.cpp +++ b/csrc/gguf_expert_cache.cpp @@ -29,12 +29,18 @@ #include #include #include +#include #include #include #include #include #include +#include "gguf_shared_expert_dma.h" +#include "gguf_shared_expert_lane.h" +#include "gguf_shared_expert_jobs.h" +#include "gguf_shared_engram.h" + #if defined(__x86_64__) || defined(__i386__) #include #endif @@ -769,9 +775,133 @@ void copy_h2d(int64_t dst, int64_t src, int64_t bytes, int64_t stream_ptr) { "copy host to device"); } +// Batched strided host->device copies (one Python call per prefill tile). +void copy_h2d_2d(const std::vector& dst, const std::vector& src, + const std::vector& dpitch, + const std::vector& spitch, + const std::vector& width, + const std::vector& height, int64_t stream_ptr) { + const size_t n = dst.size(); + if (src.size() != n || dpitch.size() != n || spitch.size() != n || + width.size() != n || height.size() != n) + throw std::invalid_argument("copy_h2d_2d: mismatched descriptor lists"); + auto stream = reinterpret_cast(stream_ptr); + for (size_t i = 0; i < n; ++i) + check(cudaMemcpy2DAsync(reinterpret_cast(dst[i]), dpitch[i], + reinterpret_cast(src[i]), spitch[i], + width[i], height[i], cudaMemcpyHostToDevice, stream), + "copy host to device 2d"); +} + } // namespace PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + using namespace vllm::gguf; + m.def("shared_engram_control_layout", [](uint64_t maximum) { + return std::vector{SharedEngramReader::control_words(maximum), + SharedEngramReader::channel_words(maximum), + SharedEngramReader::rank_words(maximum)}; + }); + pybind11::class_(m, "SharedEngramTable") + .def(pybind11::init([](int fd, uint64_t offset, uint64_t rows) { + return SharedEngramTable{fd, offset, rows}; + })); + pybind11::class_(m, "SharedEngramReader") + .def(pybind11::init([](uint64_t data, uint64_t bytes, uint64_t control, + uint64_t maximum, uint64_t row_bytes, + std::vector tables, + uint64_t timeout_ns) { + return std::make_unique( + reinterpret_cast(data), bytes, + reinterpret_cast(control), maximum, row_bytes, + std::move(tables), timeout_ns); + })) + .def("stop", &SharedEngramReader::stop, + pybind11::call_guard()) + .def("error", &SharedEngramReader::error); + pybind11::class_(m, "SharedExpertDmaLane") + .def(pybind11::init>, uint64_t, unsigned, + unsigned, unsigned, unsigned, unsigned, uint64_t>()) + .def("stop", &SharedExpertDmaLane::stop, + pybind11::call_guard()) + .def("error", &SharedExpertDmaLane::error) + .def("completed", &SharedExpertDmaLane::completed); + pybind11::class_(m, "SharedExpertJobLane") + .def(pybind11::init>, + std::vector>, uint64_t>()) + .def("stop", &SharedExpertJobLane::stop, + pybind11::call_guard()) + .def("error", &SharedExpertJobLane::error) + .def("completed", &SharedExpertJobLane::completed); + m.def("shared_expert_hint_layout", [] { + return std::vector{kSharedHintStart, kSharedHintHeader, + kSharedHintEntries, kSharedHintWords, + kSharedHintTokens, kSharedHintCands}; + }); + m.def("shared_expert_job_layout", [] { + return std::vector{kSharedJobStart, kSharedJobHeader, + kSharedJobEntries, kSharedJobWords}; + }); + m.def("shared_expert_control_layout", [] { + return std::vector{kSharedControlWords, kSharedRankStart, + kSharedRankWords, kSharedOpStart, + kSharedOpWords, kSharedMaxCopies}; + }); + m.def("shared_expert_hit_layout", [] { + return std::vector{kSharedHitStart, kSharedHitHeader, + kSharedHitEntries, kSharedHitWords, + kSharedTotalWords}; + }); + pybind11::class_(m, "SharedExpertLayer") + .def( + pybind11::init([](uint64_t bytes, int32_t experts, int32_t resident) { + return SharedExpertLayer{bytes, experts, resident}; + })); + pybind11::class_(m, "SharedExpertCopy") + .def_readonly("slot", &SharedExpertCopy::slot) + .def_readonly("incoming", &SharedExpertCopy::incoming) + .def_readonly("victim", &SharedExpertCopy::victim) + .def_readonly("load_offset", &SharedExpertCopy::load_offset) + .def_readonly("save_offset", &SharedExpertCopy::save_offset); + pybind11::class_(m, "SharedExpertPlanner") + .def( + pybind11::init([](uint64_t pointer, uint64_t bytes, + std::vector layers, + unsigned ranks, unsigned routes, uint64_t reserve) { + return std::make_unique( + reinterpret_cast(pointer), bytes, std::move(layers), + ranks, routes, reserve); + })) + .def("plan", &SharedExpertPlanner::plan, + pybind11::call_guard()) + .def("acknowledge", &SharedExpertPlanner::acknowledge) + .def("ready", &SharedExpertPlanner::ready) + .def("finish", &SharedExpertPlanner::finish, + pybind11::call_guard()) + .def("slot", &SharedExpertPlanner::slot) + .def("cold_offset", &SharedExpertPlanner::cold_offset) + .def("acquire", &SharedExpertPlanner::acquire) + .def("release", &SharedExpertPlanner::release) + .def("stop", &SharedExpertPlanner::stop); + pybind11::class_(m, "SharedExpertCoordinator") + .def(pybind11::init([](uint64_t arena, uint64_t bytes, uint64_t control, + std::vector layers, + unsigned ranks, unsigned routes, uint64_t reserve, + uint64_t timeout) { + return std::make_unique( + reinterpret_cast(arena), bytes, + reinterpret_cast(control), std::move(layers), ranks, + routes, reserve, timeout); + })) + .def("stop", &SharedExpertCoordinator::stop, + pybind11::call_guard()) + .def("error", &SharedExpertCoordinator::error) + .def("slot", &SharedExpertCoordinator::slot) + .def("cold_offset", &SharedExpertCoordinator::cold_offset) + .def("acquire", &SharedExpertCoordinator::acquire) + .def("release", &SharedExpertCoordinator::release); pybind11::class_(m, "ExpertCacheLane") .def(pybind11::init, std::vector, @@ -795,4 +925,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("host_unregister", &host_unregister); m.def("mapped_view", &mapped_view); m.def("copy_h2d", ©_h2d); + m.def("copy_h2d_2d", ©_h2d_2d, + pybind11::call_guard()); } diff --git a/csrc/gguf_shared_engram.h b/csrc/gguf_shared_engram.h new file mode 100644 index 0000000000..1624bf8741 --- /dev/null +++ b/csrc/gguf_shared_engram.h @@ -0,0 +1,289 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#pragma once + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace vllm::gguf { + +struct SharedEngramTable { + int fd; + uint64_t offset, rows; +}; + +// Two independently prefetched table channels, one raw row buffer per table. +// Each GPU writes its own request record and consumption sequence. The CPU +// lane never calls CUDA; it publishes readiness only after all preads and +// duplicate-row scatters complete. A buffer is reused only after all ranks +// consume its previous sequence. +class SharedEngramReader { + public: + static constexpr uint64_t kChannels = 2, kRanks = 4, kHeader = 8; + + static uint64_t rank_words(uint64_t maximum) { return maximum + 8; } + static uint64_t channel_words(uint64_t maximum) { + return kHeader + kRanks * rank_words(maximum); + } + static uint64_t control_words(uint64_t maximum) { + return kChannels * channel_words(maximum); + } + + SharedEngramReader(uint8_t* data, uint64_t bytes, uint64_t* control, + uint64_t maximum, uint64_t row_bytes, + std::vector tables, uint64_t timeout_ns) + : data_(data), + control_(control), + maximum_(maximum), + row_bytes_(row_bytes), + timeout_(timeout_ns), + tables_(std::move(tables)) { + if (!data || !control || reinterpret_cast(control) % 64 || + !maximum || maximum > 16384 * 24 || row_bytes != 144 || + tables_.size() != kChannels || !timeout_ || + bytes < kChannels * maximum * row_bytes) + throw std::invalid_argument("invalid shared Engram geometry"); + try { + for (const auto& table : tables_) { + struct stat info{}; + if (table.fd < 0 || !table.rows || fstat(table.fd, &info) || + info.st_size < 0 || table.offset > uint64_t(info.st_size) || + table.rows > (uint64_t(info.st_size) - table.offset) / row_bytes) + throw std::invalid_argument("shared Engram table exceeds its file"); + const int fd = fcntl(table.fd, F_DUPFD_CLOEXEC, 0); + if (fd < 0) + throw std::runtime_error("cannot duplicate Engram source fd"); + fds_.push_back(fd); + } + std::fill_n(control_, control_words(maximum), 0); + // A short decode reads at most 144 rows per channel. Creating and + // joining eight threads for each batch can outlast the cached reads. + // Keep the same bounded I/O parallelism for both hot and cold rows. + // Cold rows are random 144-byte NVMe reads: queue depth (threads) + // bounds IOPS. VLLM_V41_ENGRAM_READERS overrides the default. + size_t readers = 64; + if (const char* v = std::getenv("VLLM_V41_ENGRAM_READERS")) readers = std::max(1, std::atoi(v)); + for (size_t i = 0; i < readers; ++i) + readers_.emplace_back([this] { read_worker(); }); + thread_ = std::thread([this] { run(); }); + } catch (...) { + stop_readers(); + for (int fd : fds_) close(fd); + throw; + } + } + + ~SharedEngramReader() { + stop(); + for (int fd : fds_) close(fd); + } + + void stop() { + stopped_.store(true, std::memory_order_release); + if (thread_.joinable()) thread_.join(); + stop_readers(); + for (size_t c = 0; c < kChannels; ++c) store(header(c) + 1, 1); + } + + std::string error() { + std::lock_guard guard(mutex_); + return error_; + } + + private: + static uint64_t now() { + return std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count(); + } + static uint64_t load(const uint64_t* p) { + return __atomic_load_n(p, __ATOMIC_ACQUIRE); + } + static void store(uint64_t* p, uint64_t v) { + __atomic_store_n(p, v, __ATOMIC_RELEASE); + } + uint64_t* header(size_t c) { return control_ + c * channel_words(maximum_); } + uint64_t* rank(size_t c, size_t r) { + return header(c) + kHeader + r * rank_words(maximum_); + } + void check(uint64_t begin) { + if (stopped_.load(std::memory_order_acquire)) + throw std::runtime_error("shared Engram reader stopped during request"); + if (now() - begin >= timeout_) + throw std::runtime_error("shared Engram request timed out"); + } + + void stop_readers() { + { + std::lock_guard guard(read_mutex_); + readers_stopped_ = true; + } + read_ready_.notify_all(); + for (auto& reader : readers_) + if (reader.joinable()) reader.join(); + } + + void read_worker() { + uint64_t completed = 0; + std::unique_lock guard(read_mutex_); + for (;;) { + read_ready_.wait(guard, [&] { + return readers_stopped_ || read_generation_ != completed; + }); + if (readers_stopped_) return; + completed = read_generation_; + // The coordinator cannot replace this task until all workers have + // acknowledged it. Its captures remain alive in read_rows below. + guard.unlock(); + read_task_(); + guard.lock(); + if (!--read_remaining_) read_done_.notify_one(); + } + } + + void read_rows(size_t channel, uint64_t count, uint64_t begin) { + auto* raw = data_ + channel * maximum_ * row_bytes_; + const auto* ids = rank(channel, 0) + 8; + std::unordered_map first; + std::vector> unique, duplicate; + first.reserve(count); + unique.reserve(count); + for (uint64_t i = 0; i < count; ++i) { + if (ids[i] >= tables_[channel].rows) + throw std::runtime_error("shared Engram row id outside table"); + const auto inserted = first.emplace(ids[i], i); + if (inserted.second) + unique.emplace_back(ids[i], i); + else + duplicate.emplace_back(i, inserted.first->second); + } + std::sort(unique.begin(), unique.end()); + std::atomic next{0}; + std::atomic failed{false}; + { + std::lock_guard guard(read_mutex_); + read_task_ = [&] { + for (;;) { + const size_t i = next.fetch_add(1, std::memory_order_relaxed); + if (i >= unique.size() || failed.load(std::memory_order_acquire)) + break; + const auto [row, output] = unique[i]; + auto* target = raw + output * row_bytes_; + uint64_t offset = tables_[channel].offset + row * row_bytes_; + uint64_t left = row_bytes_; + while (left) { + const auto got = pread(fds_[channel], target, left, off_t(offset)); + if (got <= 0) { + failed.store(true, std::memory_order_release); + break; + } + left -= uint64_t(got); + offset += uint64_t(got); + target += got; + } + } + }; + read_remaining_ = readers_.size(); + ++read_generation_; + } + read_ready_.notify_all(); + { + std::unique_lock guard(read_mutex_); + read_done_.wait(guard, [&] { return !read_remaining_; }); + read_task_ = {}; + } + if (failed.load(std::memory_order_acquire)) + throw std::runtime_error("shared Engram source returned a short read"); + check(begin); + for (const auto& [to, from] : duplicate) + std::memcpy(raw + to * row_bytes_, raw + from * row_bytes_, row_bytes_); + auto* h = header(channel); + h[2] = unique.size(); + h[3] = unique.size() * row_bytes_; + h[4] = now() - begin; + h[5] = count; + } + + void run() { + std::array sequence{1, 1}, started{0, 0}; + try { + while (!stopped_.load(std::memory_order_acquire)) { + bool work = false; + for (size_t c = 0; c < kChannels; ++c) { + bool all = true, consumed = true; + for (size_t r = 0; r < kRanks; ++r) { + auto* request = rank(c, r); + const auto tag = load(request); + if (tag > sequence[c]) + throw std::runtime_error("shared Engram sequence skipped"); + if (tag == sequence[c] && !started[c]) started[c] = now(); + all &= tag == sequence[c]; + consumed &= load(request + 1) == sequence[c] - 1; + } + if (!started[c]) continue; + check(started[c]); + if (!all || !consumed) continue; + const auto count = rank(c, 0)[2]; + if (!count || count > maximum_) + throw std::runtime_error("invalid shared Engram row count"); + for (size_t r = 1; r < kRanks; ++r) { + if (rank(c, r)[2] != count || + !std::equal(rank(c, 0) + 8, rank(c, 0) + 8 + count, + rank(c, r) + 8)) + throw std::runtime_error( + "shared Engram ranks requested different rows"); + } + read_rows(c, count, started[c]); + store(header(c), sequence[c]); + ++sequence[c]; + started[c] = 0; + work = true; + } + if (!work) std::this_thread::sleep_for(std::chrono::microseconds(10)); + } + } catch (const std::exception& e) { + { + std::lock_guard guard(mutex_); + error_ = e.what(); + } + for (size_t c = 0; c < kChannels; ++c) store(header(c) + 1, 1); + } + } + + uint8_t* data_; + uint64_t* control_; + uint64_t maximum_, row_bytes_, timeout_; + std::vector tables_; + std::vector fds_; + std::atomic stopped_{false}; + std::thread thread_; + std::vector readers_; + std::mutex read_mutex_; + std::condition_variable read_ready_, read_done_; + std::function read_task_; + uint64_t read_generation_ = 0; + size_t read_remaining_ = 0; + bool readers_stopped_ = false; + std::mutex mutex_; + std::string error_; +}; +} // namespace vllm::gguf diff --git a/csrc/gguf_shared_expert_cpu.h b/csrc/gguf_shared_expert_cpu.h new file mode 100644 index 0000000000..ad864b12a5 --- /dev/null +++ b/csrc/gguf_shared_expert_cpu.h @@ -0,0 +1,383 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#pragma once + +// CPU execution of cold decode experts (VLLM_V41_CPU_EXPERTS). The owner +// lane computes chosen cold experts straight from the raw GGUF bytes in the +// shared host arena while the GPUs transfer the others and run resident +// experts, so a miss costs neither uplink bandwidth nor an eviction. Each +// 256-weight IQ2_XXS / IQ2_XS / IQ1_M block is unpacked once per 32-weight +// group into unsigned magnitudes, a sign carrier and odd per-16 scales +// (w = d / 8 * s16 * v, the ggml formats), then every token does +// sign / maddubs / madd against Q8_K activations. The expert matches the GPU +// semantics (clamped SwiGLU times the route probability rounded to FP16, +// then down); activations use Q8_K blocks of 256 instead of Q8_1 blocks of +// 32, so results are close but not bit-identical. +#include "gguf_cpu_iq_tables.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#pragma GCC push_options +#pragma GCC target("avx2,fma,f16c") + +namespace vllm::gguf { + +#define V41_CPU_TARGET + +namespace cpu_moe { + +struct BlockIq2xxs { uint16_t d; uint16_t qs[32]; }; +struct BlockIq2xs { uint16_t d; uint16_t qs[32]; uint8_t scales[8]; }; +struct BlockIq1m { uint8_t qs[32]; uint8_t qh[16]; uint8_t scales[8]; }; +struct BlockQ8K { float d; int8_t qs[256]; int16_t bsums[16]; }; +static_assert(sizeof(BlockIq2xxs) == 66 && sizeof(BlockIq2xs) == 74 && sizeof(BlockIq1m) == 56); + +inline size_t block_bytes(int type) { + switch (type) { + case 16: return sizeof(BlockIq2xxs); + case 17: return sizeof(BlockIq2xs); + case 29: return sizeof(BlockIq1m); + default: return 0; + } +} + +inline const uint64_t* sign64() { + static const auto table = [] { + std::vector t(128); + for (int i = 0; i < 128; ++i) { + uint64_t v = 0; + for (int j = 0; j < 8; ++j) + v |= uint64_t(cpu_iq::ksigns_iq2xs[i] & cpu_iq::kmask_iq2xs[j] ? 0xFF : 0x01) << (8 * j); + t[i] = v; + } + return t; + }(); + return table.data(); +} + +V41_CPU_TARGET inline float h2f(uint16_t h) { return _cvtsh_ss(h); } + +// IQ1_M grid shift per qh byte: two 8-weight groups (bits 3 and 7), +1 / -1 +inline const __m128i* iq1m_delta() { + static const auto table = [] { + std::vector<__m128i> t(256); + for (int q = 0; q < 256; ++q) + t[q] = _mm_set_epi64x(q & 0x80 ? -1LL : 0x0101010101010101LL, q & 0x08 ? -1LL : 0x0101010101010101LL); + return t; + }(); + return table.data(); +} + +// ggml quantize_row_q8_K_ref +inline void quantize_q8k(const float* x, BlockQ8K* y, int k) { + for (int i = 0; i < k / 256; ++i, x += 256) { + float max = 0, amax = 0; + for (int j = 0; j < 256; ++j) + if (std::fabs(x[j]) > amax) { amax = std::fabs(x[j]); max = x[j]; } + if (!amax) { std::memset(&y[i], 0, sizeof(BlockQ8K)); continue; } + const float iscale = -127.f / max; + for (int j = 0; j < 256; ++j) { + const int v = int(std::nearbyint(iscale * x[j])); + y[i].qs[j] = int8_t(v > 127 ? 127 : v); + } + for (int j = 0; j < 16; ++j) { + int s = 0; + for (int l = 0; l < 16; ++l) s += y[i].qs[16 * j + l]; + y[i].bsums[j] = int16_t(s); + } + y[i].d = 1 / iscale; + } +} + +struct GroupIq2xxs { + using Block = BlockIq2xxs; + V41_CPU_TARGET static inline float d(const Block* x) { return h2f(x->d); } + V41_CPU_TARGET static inline void group(const Block* x, int ib, __m256i& u, __m256i& g, __m256i& sc) { + uint32_t a[2]; + std::memcpy(a, x->qs + 4 * ib, 8); + const uint8_t* i8 = reinterpret_cast(a); + const uint64_t* s = sign64(); + u = _mm256_set_epi64x(cpu_iq::iq2xxs_grid[i8[3]], cpu_iq::iq2xxs_grid[i8[2]], + cpu_iq::iq2xxs_grid[i8[1]], cpu_iq::iq2xxs_grid[i8[0]]); + g = _mm256_set_epi64x(s[(a[1] >> 21) & 127], s[(a[1] >> 14) & 127], s[(a[1] >> 7) & 127], + s[a[1] & 127]); + sc = _mm256_set1_epi16(int16_t(2 * (a[1] >> 28) + 1)); + } +}; + +struct GroupIq2xs { + using Block = BlockIq2xs; + V41_CPU_TARGET static inline float d(const Block* x) { return h2f(x->d); } + V41_CPU_TARGET static inline void group(const Block* x, int ib, __m256i& u, __m256i& g, __m256i& sc) { + const uint16_t* q = x->qs + 4 * ib; + const uint64_t* s = sign64(); + u = _mm256_set_epi64x(cpu_iq::iq2xs_grid[q[3] & 511], cpu_iq::iq2xs_grid[q[2] & 511], + cpu_iq::iq2xs_grid[q[1] & 511], cpu_iq::iq2xs_grid[q[0] & 511]); + g = _mm256_set_epi64x(s[q[3] >> 9], s[q[2] >> 9], s[q[1] >> 9], s[q[0] >> 9]); + sc = _mm256_set_m128i(_mm_set1_epi16(int16_t(2 * (x->scales[ib] >> 4) + 1)), + _mm_set1_epi16(int16_t(2 * (x->scales[ib] & 0xf) + 1))); + } +}; + +struct GroupIq1m { + using Block = BlockIq1m; + V41_CPU_TARGET static inline float d(const Block* x) { + uint16_t sc[4]; + std::memcpy(sc, x->scales, 8); + return h2f(uint16_t((sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | + (sc[3] & 0xf000))); + } + V41_CPU_TARGET static inline void group(const Block* x, int ib, __m256i& u, __m256i& g, __m256i& sc) { + uint16_t s[4]; + std::memcpy(s, x->scales, 8); + const uint8_t* qs = x->qs + 4 * ib; + const uint8_t* qh = x->qh + 2 * ib; + const __m256i grid = _mm256_set_epi64x( + cpu_iq::iq1s_grid[qs[3] | ((qh[1] << 4) & 0x700)], cpu_iq::iq1s_grid[qs[2] | ((qh[1] << 8) & 0x700)], + cpu_iq::iq1s_grid[qs[1] | ((qh[0] << 4) & 0x700)], cpu_iq::iq1s_grid[qs[0] | ((qh[0] << 8) & 0x700)]); + const __m128i* dt = iq1m_delta(); + const __m256i delta = _mm256_set_m128i(dt[qh[1]], dt[qh[0]]); + // (grid + delta) * 8 = 8 * grid +- 1 (grid in {-1, 0, 1}); never zero + g = _mm256_add_epi8(_mm256_sign_epi8(_mm256_set1_epi8(8), grid), delta); + u = _mm256_abs_epi8(g); + sc = _mm256_set_m128i(_mm_set1_epi16(int16_t(2 * ((s[ib / 2] >> (6 * (ib % 2) + 3)) & 0x7) + 1)), + _mm_set1_epi16(int16_t(2 * ((s[ib / 2] >> (6 * (ib % 2) + 0)) & 0x7) + 1))); + } +}; + +V41_CPU_TARGET inline float hsum8(__m256 v) { + __m128 x = _mm_add_ps(_mm256_castps256_ps128(v), _mm256_extractf128_ps(v, 1)); + x = _mm_add_ps(x, _mm_movehl_ps(x, x)); + x = _mm_add_ss(x, _mm_movehdup_ps(x)); + return _mm_cvtss_f32(x); +} + +// out[m * ldo + r] = row r . q[m] for rows [r0, r1) of a row-major matrix of +// nb blocks per row. +template +V41_CPU_TARGET void rows(const uint8_t* w, size_t row_bytes, int nb, int r0, int r1, + const BlockQ8K* const* q, float* out, int ldo) { + using Block = typename G::Block; + for (int r = r0; r < r1; ++r) { + const Block* row = reinterpret_cast(w + size_t(r) * row_bytes); + __m256 accf[M]; + for (int m = 0; m < M; ++m) accf[m] = _mm256_setzero_ps(); + for (int b = 0; b < nb; ++b) { + const Block* xb = row + b; + __m256i acc[M]; + for (int m = 0; m < M; ++m) acc[m] = _mm256_setzero_si256(); + for (int ib = 0; ib < 8; ++ib) { + __m256i u, g, sc; + G::group(xb, ib, u, g, sc); + for (int m = 0; m < M; ++m) { + const __m256i q8 = _mm256_loadu_si256(reinterpret_cast(q[m][b].qs + 32 * ib)); + acc[m] = _mm256_add_epi32( + acc[m], _mm256_madd_epi16(_mm256_maddubs_epi16(u, _mm256_sign_epi8(q8, g)), sc)); + } + } + const float dw = G::d(xb); + for (int m = 0; m < M; ++m) + accf[m] = _mm256_fmadd_ps(_mm256_set1_ps(dw * q[m][b].d), _mm256_cvtepi32_ps(acc[m]), accf[m]); + } + for (int m = 0; m < M; ++m) out[m * ldo + r] = 0.125f * hsum8(accf[m]); + } +} + +template +void rows_any(int m, const uint8_t* w, size_t row_bytes, int nb, int r0, int r1, + const BlockQ8K* const* q, float* out, int ldo) { + switch (m) { + case 1: return rows(w, row_bytes, nb, r0, r1, q, out, ldo); + case 2: return rows(w, row_bytes, nb, r0, r1, q, out, ldo); + case 3: return rows(w, row_bytes, nb, r0, r1, q, out, ldo); + case 4: return rows(w, row_bytes, nb, r0, r1, q, out, ldo); + case 5: return rows(w, row_bytes, nb, r0, r1, q, out, ldo); + case 6: return rows(w, row_bytes, nb, r0, r1, q, out, ldo); + case 7: return rows(w, row_bytes, nb, r0, r1, q, out, ldo); + default: return rows(w, row_bytes, nb, r0, r1, q, out, ldo); + } +} + +inline void dot_rows(int type, int m, const uint8_t* w, size_t row_bytes, int nb, int r0, int r1, + const BlockQ8K* const* q, float* out, int ldo) { + if (type == 16) rows_any(m, w, row_bytes, nb, r0, r1, q, out, ldo); + else if (type == 17) rows_any(m, w, row_bytes, nb, r0, r1, q, out, ldo); + else rows_any(m, w, row_bytes, nb, r0, r1, q, out, ldo); +} + +} // namespace cpu_moe + +// One cold expert of one plan: its raw bytes and the verify tokens routed to it. +struct SharedCpuExpert { + const uint8_t* bytes; // gate rows, up rows, down rows (original GGUF layout) + int kinds[3]; + int tokens; + int token[8]; + float weight[8]; +}; + +// Persistent spinning pool. run() is called by the owner coordinator lane, +// which also works as thread 0; no CUDA calls are made here. +class SharedExpertCpu { + public: + static constexpr int kMaxTokens = 8; + + SharedExpertCpu(int hidden, int intermediate, int threads) + : hidden_(hidden), inter_(intermediate), threads_(threads < 1 ? 1 : threads) { + if (hidden % 256 || intermediate % 256 || hidden > 8192 || intermediate > 8192) + throw std::invalid_argument("CPU experts need 256-multiple dimensions"); + if (!__builtin_cpu_supports("avx2") || !__builtin_cpu_supports("fma") || + !__builtin_cpu_supports("f16c")) + throw std::runtime_error("CPU experts need AVX2/FMA/F16C"); + qx_.resize(kMaxTokens * (hidden / 256)); + parts_.resize(threads_); + for (auto& p : parts_) p.resize(size_t(kMaxTokens) * hidden); + cpu_moe::sign64(); + cpu_moe::iq1m_delta(); + for (int t = 1; t < threads_; ++t) pool_.emplace_back([this, t] { worker(t); }); + } + + ~SharedExpertCpu() { + stop_.store(true, std::memory_order_release); + generation_.fetch_add(1, std::memory_order_acq_rel); + for (auto& t : pool_) t.join(); + } + + // x: [tokens, hidden] FP16 (host); y: [tokens, hidden] FP16 (host), the sum + // over experts of weight * expert(x[token]) for the routed tokens. + void run(const uint16_t* x, int tokens, const std::vector& experts, uint16_t* y) { + if (tokens < 1 || tokens > kMaxTokens) throw std::invalid_argument("CPU expert tokens"); + x_ = x; y_ = y; tokens_ = tokens; experts_ = &experts; + const size_t n = experts.size(); + if (gate_.size() < n * kMaxTokens * inter_) { + gate_.resize(n * kMaxTokens * inter_); + up_.resize(n * kMaxTokens * inter_); + qh_.resize(n * kMaxTokens * (inter_ / 256)); + } + if (acc_.size() < size_t(kMaxTokens) * hidden_) acc_.resize(size_t(kMaxTokens) * hidden_); + const uint64_t job = generation_.fetch_add(1, std::memory_order_acq_rel) + 1; + work(0, job); + } + + private: + // Arrivals only grow: job j phase p completes at ((j - 1) * 4 + p) * threads. + void barrier(uint64_t job, int phase) { + const uint64_t target = ((job - 1) * 4 + phase) * uint64_t(threads_); + arrived_.fetch_add(1, std::memory_order_acq_rel); + while (arrived_.load(std::memory_order_acquire) < target) _mm_pause(); + } + + void worker(int id) { + uint64_t seen = 0; // jobs are numbered from 1; never miss one posted before start + auto idle_since = std::chrono::steady_clock::now(); + while (true) { + const uint64_t now = generation_.load(std::memory_order_acquire); + if (now != seen) { + seen = now; + if (stop_.load(std::memory_order_acquire)) return; + work(id, now); + idle_since = std::chrono::steady_clock::now(); + continue; + } + if (std::chrono::steady_clock::now() - idle_since < std::chrono::milliseconds(4)) + _mm_pause(); + else + std::this_thread::sleep_for(std::chrono::microseconds(20)); + } + } + + void work(int id, uint64_t job) { + const auto& ex = *experts_; + const int n = int(ex.size()); + const int hb = hidden_ / 256, ib = inter_ / 256; + // phase 1: quantize the input tokens + for (int t = id; t < tokens_; t += threads_) { + float row[8192]; + for (int c = 0; c < hidden_; ++c) row[c] = cpu_moe::h2f(x_[size_t(t) * hidden_ + c]); + cpu_moe::quantize_q8k(row, qx_.data() + t * hb, hidden_); + } + barrier(job, 1); + // phase 2: gate and up rows of every expert (rows split across threads) + { + const int total = 2 * n * inter_; + const int r0 = int(int64_t(total) * id / threads_), r1 = int(int64_t(total) * (id + 1) / threads_); + for (int at = r0; at < r1;) { + const int unit = at / inter_, e = unit / 2, mat = unit % 2, row = at % inter_; + const int stop = std::min(r1, (unit + 1) * inter_) - unit * inter_; + const auto& x = ex[e]; + const cpu_moe::BlockQ8K* q[kMaxTokens]; + for (int m = 0; m < x.tokens; ++m) q[m] = qx_.data() + x.token[m] * hb; + const size_t rb = hb * cpu_moe::block_bytes(x.kinds[mat]); + const uint8_t* w = x.bytes + (mat ? size_t(inter_) * hb * cpu_moe::block_bytes(x.kinds[0]) : 0); + float* out = (mat ? up_ : gate_).data() + size_t(e) * kMaxTokens * inter_; + cpu_moe::dot_rows(x.kinds[mat], x.tokens, w, rb, hb, row, stop, q, out, inter_); + at = unit * inter_ + stop; + } + } + barrier(job, 2); + // phase 3: clamped SwiGLU x probability (FP16 rounding), quantize per (expert, token) + for (int job = id; job < n * kMaxTokens; job += threads_) { + const int e = job / kMaxTokens, m = job % kMaxTokens; + if (m >= ex[e].tokens) continue; + float h[8192]; + const float* g = gate_.data() + (size_t(e) * kMaxTokens + m) * inter_; + const float* u = up_.data() + (size_t(e) * kMaxTokens + m) * inter_; + const float p = ex[e].weight[m]; + for (int c = 0; c < inter_; ++c) { + const float gc = std::min(g[c], 10.f), uc = std::max(std::min(u[c], 10.f), -10.f); + const float v = gc / (1.f + std::exp(-gc)) * uc * p; + h[c] = _cvtsh_ss(_cvtss_sh(v, 0)); + } + cpu_moe::quantize_q8k(h, qh_.data() + (size_t(e) * kMaxTokens + m) * ib, inter_); + } + barrier(job, 3); + // phase 4: down rows (hidden rows split across threads), summed over experts + { + const int r0 = int(int64_t(hidden_) * id / threads_), r1 = int(int64_t(hidden_) * (id + 1) / threads_); + std::vector& part = parts_[id]; + for (int t = 0; t < tokens_; ++t) + std::memset(acc_.data() + size_t(t) * hidden_ + r0, 0, sizeof(float) * (r1 - r0)); + for (int e = 0; e < n; ++e) { + const auto& x = ex[e]; + const cpu_moe::BlockQ8K* q[kMaxTokens]; + for (int m = 0; m < x.tokens; ++m) q[m] = qh_.data() + (size_t(e) * kMaxTokens + m) * ib; + const size_t gu = size_t(inter_) * hb * (cpu_moe::block_bytes(x.kinds[0]) + cpu_moe::block_bytes(x.kinds[1])); + cpu_moe::dot_rows(x.kinds[2], x.tokens, x.bytes + gu, ib * cpu_moe::block_bytes(x.kinds[2]), ib, r0, r1, q, + part.data(), hidden_); + for (int m = 0; m < x.tokens; ++m) + for (int r = r0; r < r1; ++r) + acc_[size_t(x.token[m]) * hidden_ + r] += part[size_t(m) * hidden_ + r]; + } + for (int t = 0; t < tokens_; ++t) + for (int r = r0; r < r1; ++r) + y_[size_t(t) * hidden_ + r] = _cvtss_sh(acc_[size_t(t) * hidden_ + r], 0); + } + barrier(job, 4); + } + + const int hidden_, inter_, threads_; + std::vector pool_; + std::atomic generation_{0}; + std::atomic arrived_{0}; + std::atomic stop_{false}; + const uint16_t* x_ = nullptr; + uint16_t* y_ = nullptr; + int tokens_ = 0; + const std::vector* experts_ = nullptr; + std::vector qx_, qh_; + std::vector gate_, up_, acc_; + std::vector> parts_; +}; + +} // namespace vllm::gguf + +#pragma GCC pop_options diff --git a/csrc/gguf_shared_expert_dma.h b/csrc/gguf_shared_expert_dma.h new file mode 100644 index 0000000000..6bcf2d03fd --- /dev/null +++ b/csrc/gguf_shared_expert_dma.h @@ -0,0 +1,186 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#pragma once + +// Per-rank writeback lane for the in-place expert exchange. After every rank +// of the group has read a plan's incoming bytes (each rank publishes its +// read-done sequence in the control page), this host thread copies the +// victims' exclusive bytes from the rank's device scratch into the shared +// arena with copy-engine DMA, then publishes the rank's completion sequence. +// SM zero-copy stores here would throttle the GPU memory pipeline for every +// concurrently running kernel; the copy engine does not. +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#if defined(__x86_64__) || defined(__i386__) + #include +#endif + +namespace vllm::gguf { + +class SharedExpertDmaLane { + public: + // layouts[layer] = 3 x (S, R, F, D, W, N, P), as the device swap kernels. + SharedExpertDmaLane(uint64_t control, uint64_t arena, uint64_t scratch, + int device, unsigned rank, unsigned ranks, + std::vector> layouts, + uint64_t scratch_per_op, unsigned rank_start, + unsigned rank_words, unsigned op_start, unsigned op_words, + unsigned flag_word, uint64_t timeout_ns) + : control_(reinterpret_cast(control)), + arena_(reinterpret_cast(arena)), + scratch_(reinterpret_cast(scratch)), + device_(device), rank_(rank), ranks_(ranks), + layouts_(std::move(layouts)), per_op_(scratch_per_op), + rank_start_(rank_start), rank_words_(rank_words), op_start_(op_start), + op_words_(op_words), flag_word_(flag_word), timeout_ns_(timeout_ns) { + if (!control_ || !arena_ || !scratch_ || rank >= ranks || ranks > 8) + throw std::invalid_argument("invalid shared expert DMA lane"); + for (const auto& l : layouts_) + if (l.size() != 21) throw std::invalid_argument("DMA lane layout"); + last_ = load(flag(rank_)); + thread_ = std::thread([this] { run(); }); + } + + ~SharedExpertDmaLane() { stop(); } + + void stop() { + stopped_.store(true, std::memory_order_release); + if (thread_.joinable()) thread_.join(); + } + + std::string error() { + std::lock_guard guard(mutex_); + return error_; + } + + uint64_t completed() const { return done_.load(std::memory_order_acquire); } + + private: + static uint64_t load(const uint64_t* p) { + return __atomic_load_n(p, __ATOMIC_ACQUIRE); + } + static void relax() { +#if defined(__x86_64__) || defined(__i386__) + _mm_pause(); +#else + std::this_thread::yield(); +#endif + } + static uint64_t now() { + return std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count(); + } + uint64_t* rank_words(unsigned r) { return control_ + rank_start_ + r * rank_words_; } + uint64_t* flag(unsigned r) { return rank_words(r) + flag_word_; } + + void check(cudaError_t e, const char* what) { + if (e != cudaSuccess) + throw std::runtime_error(std::string("DMA lane ") + what + ": " + + cudaGetErrorString(e)); + } + + void run() { + try { + check(cudaSetDevice(device_), "set device"); + // The model captures CUDA graphs in global mode on another thread; the + // copies and syncs on this lane's own stream must not be prohibited. + cudaStreamCaptureMode mode = cudaStreamCaptureModeRelaxed; + check(cudaThreadExchangeStreamCaptureMode(&mode), "capture mode"); + cudaStream_t stream; + check(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), "stream"); + uint64_t active_until = 0; + while (!stopped_.load(std::memory_order_acquire)) { + const uint64_t seq = load(flag(rank_)); + if (seq == last_) { + if (now() < active_until) relax(); + else std::this_thread::sleep_for(std::chrono::microseconds(10)); + continue; + } + const uint64_t begin = now(); + if (debug_) + std::fprintf(stderr, "[dma r%u] seq %llu seen (last %llu) flags %llu %llu %llu %llu\n", rank_, + (unsigned long long)seq, (unsigned long long)last_, + (unsigned long long)load(flag(0)), (unsigned long long)load(flag(1)), + (unsigned long long)load(flag(2)), (unsigned long long)load(flag(3))); + // Boundary down blocks are read by two ranks: no rank may overwrite + // the incoming allocation before every rank has read it. + for (;;) { + bool all = true; + for (unsigned r = 0; r < ranks_; ++r) all &= load(flag(r)) >= seq; + if (all) break; + if (stopped_.load(std::memory_order_acquire)) return; + if (now() - begin > timeout_ns_) + throw std::runtime_error("DMA lane: ranks did not finish reading"); + relax(); + } + const uint64_t layer = rank_words(rank_)[2]; + const uint64_t copies = load(control_ + 3); + if (layer >= layouts_.size() || copies > 64) + throw std::runtime_error("DMA lane: invalid plan"); + const auto& L = layouts_[layer]; + uint64_t prefix[3], off = 0; + for (int i = 0; i < 3; ++i) { prefix[i] = off; off += uint64_t(L[7 * i + 5] * L[7 * i + 4]); } + const uint64_t layer_per_op = (off + 255) / 256 * 256; // as swap_stage + for (uint64_t op = 0; op < copies; ++op) { + const uint64_t dst0 = load(control_ + op_start_ + op * op_words_ + 3); + for (int i = 0; i < 3; ++i) { + const int64_t S = L[7 * i], R = L[7 * i + 1], F = L[7 * i + 2], + W = L[7 * i + 4], N = L[7 * i + 5], P = L[7 * i + 6]; + if (!W || !N) continue; + check(cudaMemcpy2DAsync(arena_ + dst0 + R * S + F + P, S, + scratch_ + op * layer_per_op + prefix[i], W, W, N, + cudaMemcpyDeviceToHost, stream), + "writeback"); + } + } + if (debug_) + std::fprintf(stderr, "[dma r%u] seq %llu flags ok after %llu us, layer %llu copies %llu\n", rank_, + (unsigned long long)seq, (unsigned long long)((now() - begin) / 1000), + (unsigned long long)layer, (unsigned long long)copies); + check(cudaStreamSynchronize(stream), "writeback sync"); + if (debug_) + std::fprintf(stderr, "[dma r%u] seq %llu acked after %llu us\n", rank_, + (unsigned long long)seq, (unsigned long long)((now() - begin) / 1000)); + __atomic_store_n(rank_words(rank_) + 1, seq, __ATOMIC_RELEASE); + last_ = seq; + done_.fetch_add(1, std::memory_order_release); + active_until = now() + 4'000'000; + } + cudaStreamDestroy(stream); + } catch (const std::exception& e) { + std::fprintf(stderr, "[SharedExpertDmaLane rank %u] %s\n", rank_, e.what()); + std::lock_guard guard(mutex_); + error_ = e.what(); + } + } + + uint64_t* const control_; + uint8_t* const arena_; + uint8_t* const scratch_; + const int device_; + const unsigned rank_, ranks_; + const std::vector> layouts_; + const uint64_t per_op_; + const unsigned rank_start_, rank_words_, op_start_, op_words_, flag_word_; + const uint64_t timeout_ns_; + uint64_t last_ = 0; + const bool debug_ = std::getenv("VLLM_V41_DMA_DEBUG") != nullptr; + std::atomic stopped_{false}; + std::atomic done_{0}; + std::thread thread_; + std::mutex mutex_; + std::string error_; +}; + +} // namespace vllm::gguf diff --git a/csrc/gguf_shared_expert_jobs.h b/csrc/gguf_shared_expert_jobs.h new file mode 100644 index 0000000000..58562d63c2 --- /dev/null +++ b/csrc/gguf_shared_expert_jobs.h @@ -0,0 +1,235 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#pragma once + +// Per-rank background job lane (VLLM_V41_ASYNC_EVICT). The coordinator +// publishes a save only after every rank's GPU committed the plan that +// dropped the expert, so decode kernels no longer read that slot; a load +// (prefetch) targets an empty slot that no table references. This thread +// copies each job's rank-local bytes with copy-engine DMA, saves and loads on +// two streams (both copy engines), and completion counts per kind are +// written by the streams themselves (cuStreamWriteValue64). No GPU kernel +// waits on this lane; the coordinator may wait for a load only once every +// rank has issued it, which then depends on DMA progress alone. +#include +#include + +#include "gguf_shared_expert_lane.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace vllm::gguf { + +class SharedExpertJobLane { + public: + // banks[layer] = 3 bank base pointers then 3 slot strides (bytes); + // layouts[layer] = 3 x (S, R, F, D, W, N, P), as the device swap kernels. + SharedExpertJobLane(uint64_t control, uint64_t arena, int device, unsigned rank, + std::vector> banks, + std::vector> layouts, uint64_t timeout_ns) + : control_(reinterpret_cast(control)), + arena_(reinterpret_cast(arena)), + device_(device), rank_(rank), banks_(std::move(banks)), + layouts_(std::move(layouts)), timeout_ns_(timeout_ns) { + if (!control_ || !arena_ || rank >= kSharedMaxRanks || + banks_.size() != layouts_.size()) + throw std::invalid_argument("invalid shared expert job lane"); + for (size_t i = 0; i < banks_.size(); ++i) + if (banks_[i].size() != 6 || layouts_[i].size() != 21) + throw std::invalid_argument("shared expert job lane layout"); + issued_ = load(issued_word()); + for (int k = 0; k < 2; ++k) ordinal_[k] = load(done_word(k)); + thread_ = std::thread([this] { run(); }); + } + + ~SharedExpertJobLane() { stop(); } + + void stop() { + stopped_.store(true, std::memory_order_release); + if (thread_.joinable()) thread_.join(); + } + + std::string error() { + std::lock_guard guard(mutex_); + return error_; + } + + uint64_t completed() { return load(done_word(0)) + load(done_word(1)); } + + private: + static uint64_t load(const uint64_t* p) { return __atomic_load_n(p, __ATOMIC_ACQUIRE); } + static uint64_t now() { + return std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count(); + } + uint64_t* jobs() { return control_ + kSharedJobStart; } + uint64_t* done_word(int kind) { return jobs() + 1 + 4 * kind + rank_; } + uint64_t* issued_word() { return jobs() + 9 + rank_; } + + using WriteValue64 = CUresult (*)(CUstream, CUdeviceptr, cuuint64_t, unsigned int); + WriteValue64 stream_write_ = nullptr; + CUdeviceptr done_device_[2] = {0, 0}; + void init_stream_write() { + if (std::getenv("VLLM_V41_JOB_EVENTS")) return; + void* fn = nullptr; + cudaDriverEntryPointQueryResult status; + if (cudaGetDriverEntryPoint("cuStreamWriteValue64", &fn, cudaEnableDefault, &status) != + cudaSuccess || status != cudaDriverEntryPointSuccess || !fn) + return; + for (int k = 0; k < 2; ++k) { + void* device = nullptr; + if (cudaHostGetDevicePointer(&device, done_word(k), 0) != cudaSuccess) { + cudaGetLastError(); + return; + } + done_device_[k] = reinterpret_cast(device); + } + stream_write_ = reinterpret_cast(fn); + } + + void check(cudaError_t e, const char* what) { + if (e != cudaSuccess) + throw std::runtime_error(std::string("job lane ") + what + ": " + + cudaGetErrorString(e)); + } + + void issue(const uint64_t* entry, cudaStream_t stream) { + const uint64_t kind = entry[1], layer = entry[2], slot = entry[3], at = entry[4]; + const auto& L = layouts_[layer]; + const auto& B = banks_[layer]; + for (int i = 0; i < 3; ++i) { + const int64_t S = L[7 * i], R = L[7 * i + 1], F = L[7 * i + 2], D = L[7 * i + 3], + W = L[7 * i + 4], N = L[7 * i + 5], P = L[7 * i + 6]; + if (!N) continue; + auto* dev = reinterpret_cast(B[i]) + int64_t(slot) * B[3 + i]; + auto* host = arena_ + at + P + R * S + F; + if (kind == 0) { + if (W) check(cudaMemcpy2DAsync(host, S, dev, D, W, N, cudaMemcpyDeviceToHost, stream), + "save"); + } else { + check(cudaMemcpy2DAsync(dev, D, host, S, D, N, cudaMemcpyHostToDevice, stream), + "load"); + } + } + } + + void run() { + try { + check(cudaSetDevice(device_), "set device"); + // The model captures CUDA graphs in global mode on another thread. + cudaStreamCaptureMode mode = cudaStreamCaptureModeRelaxed; + check(cudaThreadExchangeStreamCaptureMode(&mode), "capture mode"); + cudaStream_t streams[2]; + for (auto& s : streams) + check(cudaStreamCreateWithFlags(&s, cudaStreamNonBlocking), "stream"); + init_stream_write(); + if (std::getenv("VLLM_V41_DMA_DEBUG")) + std::fprintf(stderr, "[job lane r%u] stream write %s\n", rank_, + stream_write_ ? "on" : "off (events)"); + std::deque> inflight[2]; + std::vector spare; + uint64_t active_until = 0; + while (!stopped_.load(std::memory_order_acquire)) { + bool busy = false; + const uint64_t head = load(jobs()); + if (issued_ < head) { + bool used[2] = {false, false}; + for (; issued_ < head; ++issued_) { + const uint64_t* entry = + jobs() + kSharedJobHeader + (issued_ % kSharedJobEntries) * kSharedJobWords; + if (load(entry) != issued_ + 1) + throw std::runtime_error("job lane: torn or lapped job entry"); + const uint64_t kind = entry[1]; + if (entry[2] >= layouts_.size() || kind > 1 || entry[6] != ordinal_[kind]) + throw std::runtime_error("job lane: invalid job"); + issue(entry, streams[kind]); + ++ordinal_[kind]; + used[kind] = true; + } + for (int k = 0; k < 2; ++k) { + if (!used[k]) continue; + if (stream_write_) { + if (stream_write_(reinterpret_cast(streams[k]), done_device_[k], + ordinal_[k], 0) != CUDA_SUCCESS) + throw std::runtime_error("job lane: stream write"); + continue; + } + cudaEvent_t event; + if (spare.empty()) { + check(cudaEventCreateWithFlags(&event, cudaEventDisableTiming), "event"); + } else { + event = spare.back(); + spare.pop_back(); + } + check(cudaEventRecord(event, streams[k]), "record"); + inflight[k].push_back({event, ordinal_[k]}); + } + __atomic_store_n(issued_word(), issued_, __ATOMIC_RELEASE); + if (debug_) + std::fprintf(stderr, "[job lane r%u] issued %llu ord %llu/%llu\n", rank_, + (unsigned long long)issued_, (unsigned long long)ordinal_[0], + (unsigned long long)ordinal_[1]); + busy = true; + } + for (int k = 0; k < 2; ++k) { + while (!inflight[k].empty()) { + const auto status = cudaEventQuery(inflight[k].front().first); + if (status == cudaErrorNotReady) break; + check(status, "query"); + __atomic_store_n(done_word(k), inflight[k].front().second, __ATOMIC_RELEASE); + spare.push_back(inflight[k].front().first); + inflight[k].pop_front(); + busy = true; + } + } + if (busy) { + active_until = now() + 4'000'000; + continue; + } + if (!inflight[0].empty() || !inflight[1].empty() || now() < active_until) { +#if defined(__x86_64__) || defined(__i386__) + _mm_pause(); +#else + std::this_thread::yield(); +#endif + } else { + std::this_thread::sleep_for(std::chrono::microseconds(20)); + } + } + for (auto& q : inflight) + for (auto& e : q) cudaEventDestroy(e.first); + for (auto e : spare) cudaEventDestroy(e); + for (auto s : streams) cudaStreamDestroy(s); + } catch (const std::exception& e) { + std::fprintf(stderr, "[SharedExpertJobLane rank %u] %s\n", rank_, e.what()); + std::lock_guard guard(mutex_); + error_ = e.what(); + } + } + + uint64_t* const control_; + uint8_t* const arena_; + const int device_; + const unsigned rank_; + const std::vector> banks_, layouts_; + const uint64_t timeout_ns_; + uint64_t issued_ = 0; + uint64_t ordinal_[2] = {0, 0}; + const bool debug_ = std::getenv("VLLM_V41_DMA_DEBUG") != nullptr; + std::atomic stopped_{false}; + std::thread thread_; + std::mutex mutex_; + std::string error_; +}; + +} // namespace vllm::gguf diff --git a/csrc/gguf_shared_expert_lane.h b/csrc/gguf_shared_expert_lane.h new file mode 100644 index 0000000000..bdb8a62a23 --- /dev/null +++ b/csrc/gguf_shared_expert_lane.h @@ -0,0 +1,620 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#pragma once + +#include "gguf_shared_expert_planner.h" + +#include +#include +#include +#include +#include +#include +#if defined(__x86_64__) || defined(__i386__) + #include +#endif + +namespace vllm::gguf { + +// All offsets below are uint64 words within one separate MAP_SHARED page. +// Only the owner CPU lane writes the plan/commit/failure header. Each rank +// writes its own request and completion sequence with a system-scope release. +constexpr size_t kSharedMaxRanks = 4; +constexpr size_t kSharedMaxRoutes = 36; +constexpr size_t kSharedMaxCopies = 44; +constexpr size_t kSharedOpStart = 8; +constexpr size_t kSharedOpWords = 5; // slot, incoming, victim, load, save +constexpr size_t kSharedRankStart = 232; +constexpr size_t kSharedRankWords = 64; +constexpr size_t kSharedControlWords = 512; +static_assert(kSharedOpStart + kSharedMaxCopies * kSharedOpWords <= + kSharedRankStart); +static_assert(kSharedRankStart + kSharedMaxRanks * kSharedRankWords <= + kSharedControlWords); +// GPU-resolved hit log, after the control words: word 0 is the producer +// head (entries written), then a ring of entries {tag=index+1, layer, count, +// ids[36]}. Rank 0 alone produces; the lane consumes it for LRU uses only. +constexpr size_t kSharedHitStart = kSharedControlWords; +constexpr size_t kSharedHitHeader = 8; +constexpr size_t kSharedHitEntries = 256; +constexpr size_t kSharedHitWords = 40; +constexpr size_t kSharedHitTotalWords = + kSharedHitStart + kSharedHitHeader + kSharedHitEntries * kSharedHitWords; +static_assert(3 + kSharedMaxRoutes <= kSharedHitWords); +// Background job ring (async eviction / prefetch), after the hit log: header +// word 0 = jobs published (coordinator), word 1 + 4 * kind + r = jobs of that +// kind completed on rank r (written by that rank's copy stream itself when +// supported), word 9 + r = jobs whose copies rank r's lane has issued; +// entries {tag=index+1, kind, layer, slot, offset, expert, kind ordinal}. +constexpr size_t kSharedJobStart = (kSharedHitTotalWords + 7) / 8 * 8; +constexpr size_t kSharedJobHeader = 16; +constexpr size_t kSharedJobEntries = 256; +constexpr size_t kSharedJobWords = 8; +// Per-rank prefetch hint rings (VLLM_V41_PREFETCH), after the job ring: each +// rank publishes, per decode layer i, its own expert quarter's top candidates +// per verify token for layers i+1 and i+2: header word 0 = entries written; +// entry {tag=index+1, layer, tokens, trunc, words[2][8][8] = expert << 32 | +// float bits of the biased selection score}. +constexpr size_t kSharedHintStart = + kSharedJobStart + kSharedJobHeader + kSharedJobEntries * kSharedJobWords; +constexpr size_t kSharedHintHeader = 8; +constexpr size_t kSharedHintEntries = 32; +constexpr size_t kSharedHintTokens = 8; +constexpr size_t kSharedHintCands = 8; +constexpr size_t kSharedHintWords = 4 + 2 * kSharedHintTokens * kSharedHintCands + 4; +constexpr size_t kSharedHintRankWords = + kSharedHintHeader + kSharedHintEntries * kSharedHintWords; +constexpr size_t kSharedTotalWords = kSharedHintStart + kSharedMaxRanks * kSharedHintRankWords; + +class SharedExpertCoordinator { + public: + SharedExpertCoordinator(uint8_t* arena, uint64_t arena_bytes, + uint64_t* control, + std::vector layers, unsigned ranks, + unsigned max_routes, uint64_t exchange_bytes, + uint64_t timeout_ns) + : control_(control), + ranks_(ranks), + timeout_ns_(timeout_ns), + planner_(arena, arena_bytes, std::move(layers), ranks, max_routes, + exchange_bytes) { + if (!control || reinterpret_cast(control) % 64 || ranks != 4 || + max_routes > kSharedMaxRoutes || !timeout_ns) + throw std::invalid_argument("invalid shared expert control mapping"); + // Construct before sharing the descriptor or publishing any GPU requests. + // The control mapping must cover the hit log (kSharedHitTotalWords). + std::fill_n(control_, kSharedTotalWords, 0); + control_[4] = max_routes; + layer_count_ = planner_.configs().size(); + for (auto& v : last_hint_) v.resize(layer_count_); + open_trace(); + thread_ = std::thread([this] { run(); }); + } + + ~SharedExpertCoordinator() { + stop(); + if (trace_) std::fclose(trace_); + } + + void stop() { + stopped_.store(true, std::memory_order_release); + if (thread_.joinable()) thread_.join(); + planner_.stop(); + store(control_ + 2, 1); + } + + std::string error() { + std::lock_guard guard(error_mutex_); + return error_; + } + + int32_t slot(size_t layer, int32_t expert) { + return planner_.slot(layer, expert); + } + uint64_t cold_offset(size_t layer, int32_t expert) { + return planner_.cold_offset(layer, expert); + } + std::pair acquire(size_t layer, int32_t expert) { + return planner_.acquire(layer, expert); + } + void release(uint64_t lease) { planner_.release(lease); } + + private: + static uint64_t now() { + return std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count(); + } + + static uint64_t load(const uint64_t* value) { + return __atomic_load_n(value, __ATOMIC_ACQUIRE); + } + + static void store(uint64_t* destination, uint64_t value) { + __atomic_store_n(destination, value, __ATOMIC_RELEASE); + } + + static void relax() { +#if defined(__x86_64__) || defined(__i386__) + _mm_pause(); +#else + std::this_thread::yield(); +#endif + } + + uint64_t* rank(unsigned index) { + return control_ + kSharedRankStart + index * kSharedRankWords; + } + + void check(uint64_t started, const char* where = "request") { + if (stopped_.load(std::memory_order_acquire)) + throw std::runtime_error( + "shared expert coordinator stopped during a plan"); + if (now() - started >= timeout_ns_) { + char text[512]; + std::snprintf(text, sizeof(text), + "shared expert group %s timed out: req %llu %llu %llu %llu " + "done %llu %llu %llu %llu layer %llu %llu %llu %llu copies %llu", + where, (unsigned long long)rank(0)[0], (unsigned long long)rank(1)[0], + (unsigned long long)rank(2)[0], (unsigned long long)rank(3)[0], + (unsigned long long)rank(0)[1], (unsigned long long)rank(1)[1], + (unsigned long long)rank(2)[1], (unsigned long long)rank(3)[1], + (unsigned long long)rank(0)[2], (unsigned long long)rank(1)[2], + (unsigned long long)rank(2)[2], (unsigned long long)rank(3)[2], + (unsigned long long)control_[3]); + throw std::runtime_error(text); + } + } + + // Debug route trace (VLLM_V41_ROUTE_TRACE=): one fixed record + // per decode layer request: kind (0 GPU-resolved hit, 1 plan), layer, route + // count, copies, routes[36], miss bitmask over routes. Header: layer configs. + struct TraceRecord { + int32_t kind, layer, count, copies; + int32_t routes[kSharedMaxRoutes]; + uint64_t miss_mask; + }; + void open_trace() { + const char* prefix = std::getenv("VLLM_V41_ROUTE_TRACE"); + if (!prefix || !*prefix) return; + char path[512]; + std::snprintf(path, sizeof(path), "%s.%d.bin", prefix, int(getpid())); + trace_ = std::fopen(path, "wb"); + if (!trace_) return; + const auto configs = planner_.configs(); + const int32_t magic = 0x56344154, n = int32_t(configs.size()); + std::fwrite(&magic, 4, 1, trace_); + std::fwrite(&n, 4, 1, trace_); + for (const auto& c : configs) { + const int32_t v[2] = {c.experts, c.resident}; + std::fwrite(v, 4, 2, trace_); + } + std::fflush(trace_); + } + void trace(int32_t kind, uint64_t layer, const int32_t* routes, size_t count, + const std::vector& incoming) { + if (!trace_) return; + TraceRecord r{}; + r.kind = kind; + r.layer = int32_t(layer); + r.count = int32_t(count); + r.copies = int32_t(incoming.size()); + for (size_t i = 0; i < count && i < kSharedMaxRoutes; ++i) { + r.routes[i] = routes[i]; + for (auto e : incoming) + if (e == routes[i]) r.miss_mask |= uint64_t{1} << i; + } + std::fwrite(&r, sizeof(r), 1, trace_); + if (++trace_count_ % 256 == 0) std::fflush(trace_); + } + + // Apply GPU-resolved hits to the LRU clock. A lapped reader skips to the + // oldest entry still present; an entry is used only if its tag matches. + void drain_hits() { + const uint64_t head = load(control_ + kSharedHitStart); + if (head - hits_read_ > kSharedHitEntries) + hits_read_ = head - kSharedHitEntries; + int32_t routes[kSharedMaxRoutes]; + while (hits_read_ < head) { + const uint64_t* entry = control_ + kSharedHitStart + kSharedHitHeader + + (hits_read_ % kSharedHitEntries) * kSharedHitWords; + const uint64_t tag = load(entry); + const uint64_t layer = entry[1], count = entry[2]; + if (tag == hits_read_ + 1 && count && count <= kSharedMaxRoutes) { + for (size_t i = 0; i < count; ++i) routes[i] = int32_t(entry[3 + i]); + if (load(entry) == tag) { + planner_.touch(layer, routes, count); + evaluate(layer, routes, count); + trace(0, layer, routes, count, {}); + } + } + ++hits_read_; + } + } + + uint64_t* jobs() { return control_ + kSharedJobStart; } + + uint64_t kind_done(int kind) { + uint64_t done = load(jobs() + 1 + 4 * kind); + for (unsigned r = 1; r < ranks_; ++r) + done = std::min(done, load(jobs() + 1 + 4 * kind + r)); + return done; + } + + // Free slots return once every rank's lane completed a job. + void drain_jobs() { + if (!jobs_published_) return; + for (int kind = 0; kind < 2; ++kind) { + const uint64_t done = kind_done(kind); + if (done > kind_completed_[kind]) { + planner_.complete_jobs(kind, done); + kind_completed_[kind] = done; + } + } + uint64_t issued = load(jobs() + 9); + for (unsigned r = 1; r < ranks_; ++r) issued = std::min(issued, load(jobs() + 9 + r)); + jobs_completed_ = issued; // lanes copy an entry when they issue it + } + + // A prefetch every rank has issued completes on DMA progress alone (the + // copy streams write completion), so a bounded wait cannot deadlock and + // is cheaper than transferring the expert again. + void wait_prefetches(uint64_t layer, const std::vector& routes) { + if (!prefetch_) return; + const auto waits = planner_.inflight_for(layer, routes); + if (waits.empty()) return; + uint64_t issued = load(jobs() + 9); + for (unsigned r = 1; r < ranks_; ++r) issued = std::min(issued, load(jobs() + 9 + r)); + const uint64_t begin = now(); + for (const auto& [id, ordinal] : waits) { + if (issued <= id) continue; + while (kind_done(1) <= ordinal && now() - begin < prefetch_wait_ns_) relax(); + } + stat_wait_ns_ += now() - begin; + ++stat_waits_; + drain_jobs(); + } + + // Prefetch evaluation (VLLM_V41_PREFETCH_EVAL): of the candidates that + // were cold when hinted, how many the layer's next request routes, per + // lookahead and weight bucket. + static int bucket(float w) { + static const float edges[] = {0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f}; + int b = 0; + while (b < 7 && w >= edges[b]) ++b; + return b; + } + void evaluate(uint64_t layer, const int32_t* routes, size_t count) { + if (!prefetch_eval_ || layer >= layer_count_) return; + for (int d = 0; d < 2; ++d) { + auto& keep = last_hint_[d][layer]; + for (const auto& c : keep) { + const bool used = std::find(routes, routes + count, c.e) != routes + count; + ++eval_[d][bucket(c.w)][used]; + } + keep.clear(); + } + } + void print_eval() { + if (!prefetch_eval_) return; + for (int d = 0; d < 2; ++d) { + std::fprintf(stderr, "[coord] eval d%d", d + 1); + for (int b = 0; b < 8; ++b) + std::fprintf(stderr, " %llu/%llu", (unsigned long long)eval_[d][b][1], + (unsigned long long)(eval_[d][b][0] + eval_[d][b][1])); + std::fprintf(stderr, "\n"); + } + } + + uint64_t* hints(unsigned r) { + return control_ + kSharedHintStart + r * kSharedHintRankWords; + } + + // Merge the ranks' per-quarter candidates for each verified position into + // the global top-k, weigh by position, and prefetch the best for layers + // i+1 and i+2. + void drain_hints() { + if (!prefetch_) return; + uint64_t head = load(hints(0)); + for (unsigned r = 1; r < ranks_; ++r) head = std::min(head, load(hints(r))); + if (head - hints_read_ > kSharedHintEntries) hints_read_ = head - kSharedHintEntries; + const size_t layers = layer_count_; + while (hints_read_ < head) { + const size_t at = kSharedHintHeader + (hints_read_ % kSharedHintEntries) * kSharedHintWords; + bool ok = true; + for (unsigned r = 0; r < ranks_; ++r) ok &= load(hints(r) + at) == hints_read_ + 1; + const uint64_t layer = hints(0)[at + 1], tokens = hints(0)[at + 2]; + const uint64_t trunc = hints(0)[at + 3]; + ++hints_read_; + if (!ok || layer >= layers) continue; + const size_t n = std::min({tokens, trunc, kSharedHintTokens}); + for (int d = 0; d < 2; ++d) { + const size_t target = layer + 1 + d; + if (target >= layers) continue; + std::unordered_map weight; + for (size_t t = 0; t < n; ++t) { + std::pair cand[kSharedMaxRanks * kSharedHintCands]; + size_t m = 0; + for (unsigned r = 0; r < ranks_; ++r) + for (size_t c = 0; c < kSharedHintCands; ++c) { + const uint64_t w = + hints(r)[at + 4 + (d * kSharedHintTokens + t) * kSharedHintCands + c]; + float key; + const uint32_t bits = uint32_t(w); + std::memcpy(&key, &bits, 4); + cand[m++] = {key, int32_t(w >> 32)}; + } + const size_t k = std::min(prefetch_topk_, m); + std::partial_sort(cand, cand + k, cand + m, + [](const auto& a, const auto& b) { return a.first > b.first; }); + const float pos = 1.0f / float(1 + t); + for (size_t c = 0; c < k; ++c) + weight[cand[c].second] += pos * (1.0f + float(k - c) / float(k)); + } + std::vector> order; + for (const auto& [e, w] : weight) order.push_back({w, e}); + std::sort(order.begin(), order.end(), + [](const auto& a, const auto& b) { return a.first > b.first; }); + std::vector ids; + for (const auto& [w, e] : order) + if (w >= prefetch_min_weight_) ids.push_back(e); + if (prefetch_eval_) { + auto& keep = last_hint_[d][target]; + keep.clear(); + for (const auto& [w, e] : order) + if (planner_.cold(target, e)) keep.push_back({e, w}); + } + planner_.prefetch(target, ids, prefetch_budget_, prefetch_inflight_); + } + } + publish_jobs(); + } + + void publish_jobs() { + auto list = planner_.take_jobs(); + if (list.empty()) return; + for (const auto& job : list) { + if (job.id != jobs_published_ || job.id - jobs_completed_ >= kSharedJobEntries) + throw std::runtime_error("shared expert job ring overflow"); + auto* entry = jobs() + kSharedJobHeader + (job.id % kSharedJobEntries) * kSharedJobWords; + entry[1] = uint64_t(job.kind); + entry[2] = uint64_t(job.layer); + entry[3] = uint64_t(job.slot); + entry[4] = job.offset; + entry[5] = uint64_t(job.expert); + entry[6] = job.ordinal; + store(entry, job.id + 1); + ++jobs_published_; + } + store(jobs(), jobs_published_); + } + + void run() { + uint64_t sequence = 1; + uint64_t active_until = 0; + try { + while (!stopped_.load(std::memory_order_acquire)) { + drain_hits(); + drain_jobs(); + drain_hints(); + bool started = false; + for (unsigned r = 0; r < ranks_; ++r) { + const auto tag = load(rank(r)); + if (tag > sequence) + throw std::runtime_error( + "shared expert request skipped a sequence"); + started |= tag == sequence; + } + if (!started) { + // A running decode reaches the next layer within a few ms. Avoid + // paying a scheduler wakeup on every plan, while sleeping once the + // group becomes idle. This lane still never submits CUDA work. + if (now() < active_until) + relax(); + else + std::this_thread::sleep_for(std::chrono::microseconds(10)); + continue; + } + const auto begin = now(); + for (;;) { + bool all = true; + for (unsigned r = 0; r < ranks_; ++r) { + const auto tag = load(rank(r)); + if (tag > sequence) + throw std::runtime_error("shared expert ranks diverged"); + all &= tag == sequence; + } + if (all) break; + check(begin, "publish"); + relax(); + } + const uint64_t layer = rank(0)[2], count = rank(0)[3]; + if (!count || count > kSharedMaxRoutes) + throw std::runtime_error("invalid shared expert route count"); + std::vector routes; + for (size_t i = 0; i < count; ++i) { + if (rank(0)[4 + i] > uint64_t(std::numeric_limits::max())) + throw std::runtime_error("invalid shared expert route id"); + routes.push_back(int32_t(rank(0)[4 + i])); + } + for (unsigned r = 1; r < ranks_; ++r) { + if (rank(r)[2] != layer || rank(r)[3] != count || + !std::equal(rank(0) + 4, rank(0) + 4 + count, rank(r) + 4)) + throw std::runtime_error("shared expert rank routes differ"); + } + evaluate(layer, routes.data(), routes.size()); + wait_prefetches(layer, routes); + const int32_t cutoff = int32_t(load(control_ + 5)); + auto copies = planner_.plan(layer, sequence, routes, begin, timeout_ns_, + cutoff > 0 && cutoff < 64 ? cutoff : 0, 6, + rank(0)[41]); + control_[6] = uint64_t(planner_.trunc()); + if (coord_debug_) + std::fprintf(stderr, "[coord] seq %llu layer %llu routes %zu copies %zu trunc %d cutoff %d round %llu\n", + (unsigned long long)sequence, (unsigned long long)layer, routes.size(), + copies.size(), planner_.trunc(), cutoff, (unsigned long long)rank(0)[41]); + if (trace_) { + std::vector incoming; + for (const auto& c : copies) incoming.push_back(c.incoming); + trace(1, layer, routes.data(), routes.size(), incoming); + } + // Demand loads by the copy engines: every rank's lane must ISSUE the + // copies within a bound (a stalled lane falls back to the GPU's SM + // copy, which rewrites the same bytes); once issued, completion is + // written by the copy streams, so waiting depends on DMA alone. + std::vector by_dma(copies.size(), false); + if (dma_demand_) { + const auto dj = planner_.demand_jobs(); + uint64_t last_id = 0, last_ord = 0; + bool any = false; + for (const auto& [id, ord] : dj) + if (id != ~uint64_t{0}) { any = true; last_id = std::max(last_id, id); last_ord = std::max(last_ord, ord); } + if (any) { + publish_jobs(); + const uint64_t t0 = now(); + auto issued = [&] { + uint64_t v = load(jobs() + 9); + for (unsigned r = 1; r < ranks_; ++r) v = std::min(v, load(jobs() + 9 + r)); + return v; + }; + while (issued() <= last_id && now() - t0 < dma_issue_ns_) relax(); + // Completion is bounded too: a lane stream that shares a hardware + // queue with the waiting decode stream (false dependency) cannot + // finish until the GPU moves on, so fall back to the SM copy (same + // bytes; the late DMA is benign and its slot reuse is deferred). + bool done = false; + if (issued() > last_id) { + while (kind_done(1) <= last_ord && now() - t0 < dma_done_ns_) relax(); + done = kind_done(1) > last_ord; + } + if (done) { + for (size_t i = 0; i < dj.size(); ++i) by_dma[i] = dj[i].first != ~uint64_t{0}; + ++stat_dma_; + } else { + ++stat_dma_fallback_; + } + stat_dma_ns_ += now() - t0; + } + } + const auto evicts = planner_.evicts(); + if (copies.size() + evicts.size() > kSharedMaxCopies) + throw std::runtime_error("shared expert plan exceeds op capacity"); + for (size_t i = 0; i < copies.size() + evicts.size(); ++i) { + auto* op = control_ + kSharedOpStart + i * kSharedOpWords; + const auto& copy = i < copies.size() ? copies[i] : evicts[i - copies.size()]; + op[0] = uint64_t(int64_t(copy.slot)); + op[1] = uint64_t(int64_t(copy.incoming)); + op[2] = uint64_t(int64_t(copy.victim)); + op[3] = i < copies.size() && by_dma[i] ? uint64_t(-2) : copy.load_offset; + op[4] = copy.save_offset; + } + // Low 16 bits: copies (swap kernel ops); high: evictions after them. + control_[3] = copies.size() | (evicts.size() << 16); + store(control_, sequence); + uint64_t completed = 0; + const uint64_t all_ranks = (uint64_t{1} << ranks_) - 1; + // A plan without copies (every miss truncated) has no writeback, so + // no rank publishes a completion for it. + if (copies.empty() && evicts.empty()) { + for (unsigned r = 0; r < ranks_; ++r) planner_.acknowledge(r, sequence, now()); + completed = all_ranks; + } + while (completed != all_ranks) { + for (unsigned r = 0; r < ranks_; ++r) { + const uint64_t tag = load(rank(r) + 1), bit = uint64_t{1} << r; + if (tag > sequence) + throw std::runtime_error( + "shared expert completion skipped a sequence"); + if (tag == sequence && !(completed & bit)) { + planner_.acknowledge(r, sequence, now()); + completed |= bit; + } + } + check(begin, "ack"); + relax(); + } + while (!planner_.ready(sequence, now())) { + check(begin, "ready"); + relax(); + } + planner_.finish(sequence, now()); + { + const auto background = planner_.take_background(); + if (!background.empty()) + stat_background_ += planner_.prefetch(layer, background, background.size(), bg_inflight_); + } + publish_jobs(); + if (plan_stats_ && sequence % 2000 == 0) + std::fprintf(stderr, "[coord] plans %llu free %llu exchange %llu rescue %llu ghost %llu evict %llu pending %llu prefetch %llu hit %llu waits %llu wait_ms %.1f\n", + (unsigned long long)sequence, (unsigned long long)planner_.stat_free_, + (unsigned long long)planner_.stat_exchange_, (unsigned long long)planner_.stat_rescue_, + (unsigned long long)planner_.stat_ghost_, (unsigned long long)planner_.stat_evict_, + (unsigned long long)planner_.pending_jobs(), + (unsigned long long)planner_.stat_prefetch_, + (unsigned long long)planner_.stat_prefetch_hit_, + (unsigned long long)stat_waits_, stat_wait_ns_ / 1e6); + if (plan_stats_ && sequence % 2000 == 0 && dma_demand_) + std::fprintf(stderr, "[coord] dma plans %llu fallback %llu dma_ms %.1f\n", + (unsigned long long)stat_dma_, (unsigned long long)stat_dma_fallback_, + stat_dma_ns_ / 1e6); + if (plan_stats_ && sequence % 2000 == 0) print_eval(); + if (plan_stats_ && sequence % 2000 == 0) + std::fprintf(stderr, "[coord] skip adopt %llu background loads %llu\n", + (unsigned long long)planner_.stat_adopt_, (unsigned long long)stat_background_); + store(control_ + 1, sequence); + if (sequence == std::numeric_limits::max()) + throw std::runtime_error("shared expert request sequence exhausted"); + ++sequence; + active_until = now() + 4'000'000; + } + } catch (const std::exception& exception) { + std::fprintf(stderr, "[SharedExpertCoordinator] %s\n", exception.what()); + { + std::lock_guard guard(error_mutex_); + error_ = exception.what(); + } + planner_.stop(); + store(control_ + 2, 1); + } + } + + uint64_t* const control_; + uint64_t hits_read_ = 0; + uint64_t jobs_published_ = 0, jobs_completed_ = 0; + uint64_t kind_completed_[2] = {0, 0}; + uint64_t hints_read_ = 0, stat_waits_ = 0, stat_wait_ns_ = 0; + size_t layer_count_ = 0; + struct HintCand { int32_t e; float w; }; + std::vector> last_hint_[2]; + uint64_t eval_[2][8][2] = {}; + const bool prefetch_eval_ = std::getenv("VLLM_V41_PREFETCH_EVAL") != nullptr; + static double env_number(const char* name, double fallback) { + const char* v = std::getenv(name); + return v && *v ? std::atof(v) : fallback; + } + const bool prefetch_ = env_number("VLLM_V41_PREFETCH", 0) > 0; + const bool dma_demand_ = env_number("VLLM_V41_DMA_DEMAND", 0) > 0; + const uint64_t dma_issue_ns_ = uint64_t(env_number("VLLM_V41_DMA_ISSUE_US", 300) * 1000); + const uint64_t dma_done_ns_ = uint64_t(env_number("VLLM_V41_DMA_DONE_US", 4000) * 1000); + uint64_t stat_dma_ = 0, stat_dma_fallback_ = 0, stat_dma_ns_ = 0, stat_background_ = 0; + const size_t bg_inflight_ = size_t(env_number("VLLM_V41_BG_INFLIGHT", 8)); + const size_t prefetch_topk_ = size_t(env_number("VLLM_V41_PREFETCH_TOPK", 8)); + const size_t prefetch_budget_ = size_t(env_number("VLLM_V41_PREFETCH_BUDGET", 3)); + const size_t prefetch_inflight_ = size_t(env_number("VLLM_V41_PREFETCH_INFLIGHT", 6)); + const float prefetch_min_weight_ = float(env_number("VLLM_V41_PREFETCH_MIN_WEIGHT", 1.5)); + const uint64_t prefetch_wait_ns_ = + uint64_t(env_number("VLLM_V41_PREFETCH_WAIT_US", 3000) * 1000); + const unsigned ranks_; + const uint64_t timeout_ns_; + SharedExpertPlanner planner_; + std::atomic stopped_{false}; + const bool coord_debug_ = std::getenv("VLLM_V41_COORD_DEBUG") != nullptr; + const bool plan_stats_ = std::getenv("VLLM_V41_PLAN_STATS") != nullptr; + FILE* trace_ = nullptr; + uint64_t trace_count_ = 0; + std::thread thread_; + std::mutex error_mutex_; + std::string error_; +}; + +} // namespace vllm::gguf diff --git a/csrc/gguf_shared_expert_planner.h b/csrc/gguf_shared_expert_planner.h new file mode 100644 index 0000000000..01faae802b --- /dev/null +++ b/csrc/gguf_shared_expert_planner.h @@ -0,0 +1,653 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#pragma once + +#include "gguf_shared_expert_state.h" + +#include +#include +#include +#include +#include +#include + +namespace vllm::gguf { + +struct SharedExpertLayer { + uint64_t expert_bytes; + int32_t experts; + int32_t resident; +}; + +struct SharedExpertCopy { + int32_t slot; + int32_t incoming; + int32_t victim; + uint64_t load_offset; + uint64_t save_offset; +}; + +// Background copy-engine job for every rank's job lane (VLLM_V41_ASYNC_EVICT): +// kind 0 saves slot -> host offset (device to host, the rank's exclusive +// bytes), kind 1 loads host offset -> slot. +struct SharedExpertJob { + uint64_t id; // ring index; ordinal counts jobs of this kind + int32_t kind, layer, slot, expert; + uint64_t offset, ordinal; +}; + +constexpr uint64_t kSharedNoOffset = ~uint64_t{0}; + +// One owner per TP group, called by its native coordination lane. Persistent +// cold slots contain only original, unsliced experts. The exchange reserve +// contains independent victim destinations during a plan. Once every rank +// and prefetch reader completes, compact victims into the vacated cold slots +// before reusing exchange space. This bounded host byte copy avoids allocator +// fragmentation across layers with different expert encodings/sizes; it does +// not create a second persistent cold bank or perform CPU matrix arithmetic. +class SharedExpertPlanner { + public: + SharedExpertPlanner(uint8_t* arena, uint64_t arena_bytes, + std::vector layers, unsigned ranks, + unsigned max_routes, uint64_t exchange_bytes) + : arena_(arena), max_routes_(max_routes), fence_(arena_bytes, ranks) { + if (!arena || layers.empty() || !max_routes || exchange_bytes > arena_bytes) + throw std::invalid_argument("invalid shared expert planner dimensions"); + uint64_t offset = 0; + for (auto config : layers) { + if (!config.expert_bytes || config.experts <= 1 || config.resident <= 0 || + config.resident >= config.experts || + max_routes > unsigned(config.resident)) + throw std::invalid_argument("invalid shared expert layer capacity"); + const auto cold = uint64_t(config.experts - config.resident); + if (config.expert_bytes > (arena_bytes - offset) / cold) + throw std::invalid_argument("cold experts exceed the shared arena"); + Layer layer{config, offset, {}, {}, {}}; + layer.slot.assign(config.experts, -1); + layer.cold.assign(config.experts, -1); // absolute arena offsets + layer.use.assign(config.experts, 0); + for (int32_t expert = 0; expert < config.experts; ++expert) { + if (expert < config.resident) + layer.slot[expert] = expert; + else + layer.cold[expert] = + int64_t(offset + uint64_t(expert - config.resident) * config.expert_bytes); + } + layers_.push_back(std::move(layer)); + offset += cold * config.expert_bytes; + if (max_routes > exchange_bytes / config.expert_bytes) + throw std::invalid_argument( + "shared exchange cannot hold maximum routes"); + } + if (offset > arena_bytes - exchange_bytes) + throw std::invalid_argument("cold experts overlap the exchange reserve"); + exchange_offset_ = offset; + if (free_mode_) { + // Free-slot exchange (VLLM_V41_SWAP_FREE): the exchange reserve becomes + // per-size-class free cold allocations. A victim is saved into a free + // slot while the incoming is read (no aliasing, no host compaction); the + // incoming's vacated allocation joins its class list after every reader. + const char* env = std::getenv("VLLM_V41_FREE_SLOTS"); + const uint64_t per_class = env ? std::strtoull(env, nullptr, 10) : 12; + // Async eviction keeps async_free_ empty GPU slots per layer, so as + // many more experts per layer live on the host. + std::unordered_map layers_of; + for (const auto& layer : layers_) ++layers_of[layer.config.expert_bytes]; + uint64_t at = exchange_offset_; + for (const auto& layer : layers_) { + const uint64_t bytes = layer.config.expert_bytes; + if (free_.count(bytes)) continue; + auto& list = free_[bytes]; + const uint64_t count = per_class + uint64_t(async_free_) * layers_of[bytes]; + for (uint64_t i = 0; i < count; ++i) { + if (at + bytes > exchange_offset_ + exchange_bytes) + throw std::invalid_argument("free-slot exchange exceeds the reserve"); + list.push_back(at); + at += bytes; + } + } + } + } + + // Exact truncated verification (cutoff > 0): a cold expert whose earliest + // route position (route / top_k = verify token) is >= cutoff is not fetched + // and truncates the verify block there; experts used only at positions + // >= trunc are never fetched. trunc resets when the layer index wraps. + int32_t trunc() const { return trunc_; } + + std::vector plan(size_t layer_index, uint64_t sequence, + const std::vector& raw_routes, + uint64_t now_ns, uint64_t timeout_ns, + int32_t cutoff = 0, int32_t top_k = 6, + uint64_t round = 0) { + std::lock_guard guard(mutex_); + if (pending_) + throw std::runtime_error("shared expert plan already pending"); + auto& layer = layers_.at(layer_index); + // Deferred drop candidates (bit 30, VLLM_V41_SKIP_DEFER): an expert a + // background load already placed (ghost) or whose eviction is still in + // its slot is adopted without a copy; the others stay dropped this time + // and are queued for a background load, so later uses compute exactly. + std::vector routes(raw_routes); + background_.clear(); + for (auto& r : routes) { + if (r == kSkipRoute || !(r & kDeferFlag)) continue; + const int32_t e = r & ~kDeferFlag; + if (e < 0 || e >= layer.config.experts) + throw std::invalid_argument("shared expert deferred id outside layer"); + if (layer.ghost.count(e) || layer.evicting.count(e)) { + r = e; + ++stat_adopt_; + } else { + r = kSkipRoute; + if (!layer.prefetching.count(e) && + std::find(background_.begin(), background_.end(), e) == background_.end()) + background_.push_back(e); + } + } + if (routes.empty() || routes.size() > max_routes_) + throw std::invalid_argument( + "shared expert routes exceed request capacity"); + std::unordered_set wanted; + for (auto expert : routes) { + if (expert == kSkipRoute) continue; // dropped on the GPU (lossy skip) + if (expert < 0 || expert >= layer.config.experts) + throw std::invalid_argument("shared expert id outside layer"); + wanted.insert(expert); + } + if (clock_ == std::numeric_limits::max()) { + fence_.stop(); + throw std::runtime_error("shared expert LRU sequence exhausted"); + } + // The GPU numbers its forwards; truncation never carries across them. + if (round != round_) { trunc_ = kNoTrunc; round_ = round; } + std::unordered_set skipped; + std::vector saves_by_order; + auto* pool = free_mode_ ? &free_.at(layer.config.expert_bytes) : nullptr; + size_t available = pool ? pool->size() : 0; + size_t free_slots = layer.free_slots.size(); + for (const auto& [e, s] : layer.ghost) + if (wanted.count(e)) --free_slots; // its rescue takes that free slot + if (cutoff > 0 || trunc_ != kNoTrunc || free_mode_) { + // earliest position per non-resident expert, visited in position order + std::vector> need; + for (size_t i = 0; i < routes.size(); ++i) { + const int32_t e = routes[i], pos = int32_t(i) / top_k; + if (e == kSkipRoute || layer.slot[e] >= 0) continue; + auto it = std::find_if(need.begin(), need.end(), + [&](const auto& v) { return v.second == e; }); + if (it == need.end()) need.push_back({pos, e}); + } + std::stable_sort(need.begin(), need.end()); + for (const auto& [pos, e] : need) { + if (pos >= trunc_) skipped.insert(e); + else if (cutoff > 0 && pos >= cutoff) { trunc_ = pos; skipped.insert(e); } + else if (layer.evicting.count(e) || layer.ghost.count(e)) { + // still intact in its old slot: reinstated without a copy + } else if (free_slots > 0) { + --free_slots; + } else if (pool && available == 0) { + // No free allocation of this class: exact truncation at the + // earliest position that still needs this expert. + if (pos == 0) + throw std::runtime_error("free-slot exchange cannot hold verify position 0"); + trunc_ = pos; + skipped.insert(e); + } else if (pool) { + --available; + } + } + } + std::vector copies; + std::vector inputs, outputs; + std::unordered_set evicting; + // Route order is stable; unordered_set is used only for membership. + for (auto incoming : routes) { + if (incoming == kSkipRoute || layer.slot[incoming] >= 0 || skipped.count(incoming) || + std::any_of(copies.begin(), copies.end(), + [&](const auto& op) { return op.incoming == incoming; })) + continue; + if (auto it = layer.evicting.find(incoming); it != layer.evicting.end()) { + // Rescue: the GPU table dropped it but the save only reads the slot. + copies.push_back({it->second.slot, incoming, -1, kSharedNoOffset, + kSharedNoOffset}); + ++stat_rescue_; + continue; + } + if (auto it = layer.ghost.find(incoming); it != layer.ghost.end()) { + // Victim cache: a free slot not yet reused still holds this expert. + const int32_t slot = it->second; + layer.free_slots.erase( + std::find(layer.free_slots.begin(), layer.free_slots.end(), slot)); + copies.push_back({slot, incoming, -1, kSharedNoOffset, kSharedNoOffset}); + ghost_rescues_.push_back({incoming, layer.cold[incoming]}); + layer.ghost.erase(it); + ++stat_ghost_; + if (prefetched_.erase({int32_t(layer_index), incoming})) ++stat_prefetch_hit_; + continue; + } + const uint64_t bytes = layer.config.expert_bytes; + const uint64_t load = uint64_t(layer.cold[incoming]); + // Load-only into the least recently freed empty slot whose previous + // expert (saved asynchronously) is not routed here; its ghost is dropped. + auto pick = layer.free_slots.end(); + for (auto f = layer.free_slots.begin(); f != layer.free_slots.end(); ++f) { + auto g = std::find_if(layer.ghost.begin(), layer.ghost.end(), + [&](const auto& v) { return v.second == *f; }); + if (g == layer.ghost.end() || !wanted.count(g->first)) { pick = f; break; } + } + if (pick != layer.free_slots.end()) { + const int32_t slot = *pick; + layer.free_slots.erase(pick); + for (auto g = layer.ghost.begin(); g != layer.ghost.end(); ++g) + if (g->second == slot) { + prefetched_.erase({int32_t(layer_index), g->first}); + layer.ghost.erase(g); + break; + } + copies.push_back({slot, incoming, -1, load, kSharedNoOffset}); + inputs.push_back({load, bytes}); + ++stat_free_; + continue; + } + if (pool && pool->empty()) + throw std::runtime_error("shared expert exchange pool exhausted"); + ++stat_exchange_; + int32_t victim = -1; + for (int32_t expert = 0; expert < layer.config.experts; ++expert) { + if (layer.slot[expert] < 0 || wanted.count(expert) || + evicting.count(expert)) + continue; + if (victim < 0 || layer.use[expert] < layer.use[victim]) + victim = expert; + } + if (victim < 0) + throw std::runtime_error("no evictable shared expert slot"); + // In-place exchange (VLLM_V41_SWAP_INPLACE): devices write the victim + // into the incoming expert's own cold allocation after every rank has + // read it (device-side barrier), so finish() needs no host compaction. + uint64_t save; + if (pool) { + save = pool->back(); + pool->pop_back(); + } else { + save = inplace_ ? load : exchange_offset_ + copies.size() * bytes; + } + copies.push_back({layer.slot[victim], incoming, victim, load, save}); + inputs.push_back({load, bytes}); + outputs.push_back({save, bytes}); + evicting.insert(victim); + } + // Refill this layer's empty slots in the background: LRU victims leave + // the GPU table with this plan's commit; every rank's job lane then saves + // its exclusive bytes with copy-engine DMA (no SM interference) and the + // slot becomes free once all ranks report completion. + evicts_.clear(); + if (async_free_ > 0 && pool) { + size_t have = layer.free_slots.size() + layer.evicting.size(); + while (have < size_t(async_free_) && pool->size() > kEvictReserve && + jobs_[0].size() + jobs_[1].size() + evicts_.size() < kMaxPendingJobs && + copies.size() + evicts_.size() < kMaxPlanOps) { + int32_t victim = -1; + for (int32_t expert = 0; expert < layer.config.experts; ++expert) { + if (layer.slot[expert] < 0 || wanted.count(expert) || + evicting.count(expert)) + continue; + if (victim < 0 || layer.use[expert] < layer.use[victim]) victim = expert; + } + if (victim < 0) break; + evicts_.push_back({layer.slot[victim], -1, victim, kSharedNoOffset, + pool->back()}); + ++stat_evict_; + pool->pop_back(); + evicting.insert(victim); + ++have; + } + } + fence_.begin(sequence, inputs, outputs, now_ns, timeout_ns); + ++clock_; + for (auto expert : routes) + if (expert != kSkipRoute) layer.use[expert] = clock_; + pending_ = true; + pending_layer_ = layer_index; + copies_ = copies; + return copies; + } + + // Dropped experts of the last plan to load in the background. + std::vector take_background() { + std::lock_guard guard(mutex_); + return std::exchange(background_, {}); + } + + // Evictions published with the pending plan (after its copies). + std::vector evicts() { + std::lock_guard guard(mutex_); + return evicts_; + } + + // Jobs created by the last finish(), to publish to the job lanes. + std::vector take_jobs() { + std::lock_guard guard(mutex_); + return std::exchange(new_jobs_, {}); + } + + // Every rank completed jobs of this kind with ordinal < done. + void complete_jobs(int kind, uint64_t done) { + std::lock_guard guard(mutex_); + auto& queue = jobs_[kind]; + while (!queue.empty() && queue.front().ordinal < done) { + const auto job = queue.front(); + queue.pop_front(); + auto& layer = layers_.at(job.layer); + if (kind == 1 && demand_.erase(job.ordinal)) continue; // committed by its plan + if (kind == 1) { + // Prefetch landed on every rank: the slot is free, holding a ghost + // of the expert while its host copy is still canonical. + auto it = layer.prefetching.find(job.expert); + if (it != layer.prefetching.end() && it->second == job.ordinal) + layer.prefetching.erase(it); + layer.free_slots.push_back(job.slot); + if (layer.slot[job.expert] < 0 && layer.cold[job.expert] == int64_t(job.offset) && + !layer.ghost.count(job.expert) && !layer.prefetching.count(job.expert)) { + layer.ghost[job.expert] = job.slot; + prefetched_.insert({job.layer, job.expert}); + } + continue; + } + auto it = layer.evicting.find(job.expert); + if (it != layer.evicting.end() && it->second.job == job.id) { + layer.cold[job.expert] = int64_t(job.offset); + layer.evicting.erase(it); + // A demand load that fell back to the SM copy may still land later; + // its slot stays out of the free pool until that DMA completes. + deferred_.push_back({job.layer, job.slot, job.expert, job.offset}); + } else { + // Rescued while saving: the copy is stale, recycle its allocation. + free_.at(layer.config.expert_bytes).push_back(job.offset); + } + } + for (size_t i = 0; i < deferred_.size();) { + const auto d = deferred_[i]; + const bool pending = std::any_of(jobs_[1].begin(), jobs_[1].end(), [&](const auto& j) { + return j.layer == d.layer && j.slot == d.slot; + }); + if (pending) { ++i; continue; } + auto& layer = layers_.at(d.layer); + layer.free_slots.push_back(d.slot); + if (layer.slot[d.expert] < 0 && layer.cold[d.expert] == int64_t(d.offset) && + !layer.ghost.count(d.expert) && !layer.prefetching.count(d.expert)) + layer.ghost[d.expert] = d.slot; + deferred_[i] = deferred_.back(); + deferred_.pop_back(); + } + } + + // Demand loads by copy-engine DMA (VLLM_V41_DMA_DEMAND): the pending + // plan's free-slot loads become lane jobs. Returns, per copy, the job's + // (ring id, ordinal), or (~0, ~0) for copies the GPU performs itself. + std::vector> demand_jobs() { + std::lock_guard guard(mutex_); + std::vector> out(copies_.size(), {~uint64_t{0}, ~uint64_t{0}}); + if (!pending_ || jobs_[0].size() + jobs_[1].size() + copies_.size() >= kMaxPendingJobs) + return out; + for (size_t i = 0; i < copies_.size(); ++i) { + const auto& c = copies_[i]; + if (c.victim >= 0 || c.load_offset == kSharedNoOffset) continue; + const SharedExpertJob job{next_job_++, 1, int32_t(pending_layer_), c.slot, c.incoming, + c.load_offset, next_ordinal_[1]++}; + demand_.insert(job.ordinal); + jobs_[1].push_back(job); + new_jobs_.push_back(job); + out[i] = {job.id, job.ordinal}; + } + return out; + } + + uint64_t pending_jobs() { + std::lock_guard guard(mutex_); + return jobs_[0].size() + jobs_[1].size(); + } + + // Predictive prefetch (between plans): copy predicted cold experts of a + // later layer into its empty slots with the job lanes. Candidates are in + // priority order. A landed prefetch becomes a ghost the plan reinstates + // without a copy; a wrong prediction only spends uplink bandwidth. + size_t prefetch(size_t layer_index, const std::vector& candidates, + size_t budget, size_t max_inflight) { + std::lock_guard guard(mutex_); + if (pending_) return 0; + auto& layer = layers_.at(layer_index); + size_t issued = 0; + for (auto e : candidates) { + if (issued >= budget || jobs_[1].size() >= max_inflight || + jobs_[0].size() + jobs_[1].size() >= kMaxPendingJobs) + break; + if (e < 0 || e >= layer.config.experts || layer.slot[e] >= 0 || + layer.cold[e] < 0 || layer.evicting.count(e) || layer.ghost.count(e) || + layer.prefetching.count(e)) + continue; + // Oldest free slot; a prefetched ghost is only displaced by a newer one. + auto pick = layer.free_slots.end(); + for (auto f = layer.free_slots.begin(); f != layer.free_slots.end(); ++f) { + auto g = std::find_if(layer.ghost.begin(), layer.ghost.end(), + [&](const auto& v) { return v.second == *f; }); + if (g == layer.ghost.end() || !prefetched_.count({int32_t(layer_index), g->first})) { + pick = f; + break; + } + } + if (pick == layer.free_slots.end()) break; + const int32_t slot = *pick; + layer.free_slots.erase(pick); + for (auto g = layer.ghost.begin(); g != layer.ghost.end(); ++g) + if (g->second == slot) { + prefetched_.erase({int32_t(layer_index), g->first}); + layer.ghost.erase(g); + break; + } + const SharedExpertJob job{next_job_++, 1, int32_t(layer_index), slot, e, + uint64_t(layer.cold[e]), next_ordinal_[1]++}; + layer.prefetching[e] = job.ordinal; + jobs_[1].push_back(job); + new_jobs_.push_back(job); + ++stat_prefetch_; + ++issued; + } + return issued; + } + + // In-flight prefetches (ring id, ordinal) of experts routed here, which + // the coordinator may wait for before planning. + std::vector> inflight_for(size_t layer_index, + const std::vector& routes) { + std::lock_guard guard(mutex_); + std::vector> out; + auto& layer = layers_.at(layer_index); + for (auto e : routes) { + auto it = layer.prefetching.find(e); + if (it == layer.prefetching.end()) continue; + for (const auto& job : jobs_[1]) + if (job.ordinal == it->second) { out.push_back({job.id, job.ordinal}); break; } + } + return out; + } + + // Requests resolved entirely on the GPU (every route resident) still count + // as uses, so LRU victims stay the least recently routed experts. + void touch(size_t layer_index, const int32_t* routes, size_t count) { + std::lock_guard guard(mutex_); + auto& layer = layers_.at(layer_index); + if (clock_ == std::numeric_limits::max()) return; + ++clock_; + for (size_t i = 0; i < count; ++i) + if (routes[i] >= 0 && routes[i] < layer.config.experts) + layer.use[routes[i]] = clock_; + } + + void acknowledge(unsigned rank, uint64_t sequence, uint64_t now_ns) { + fence_.acknowledge(rank, sequence, now_ns); + } + + bool ready(uint64_t sequence, uint64_t now_ns) { + return fence_.ready(sequence, now_ns); + } + + void finish(uint64_t sequence, uint64_t now_ns) { + std::lock_guard guard(mutex_); + fence_.finish(sequence, now_ns); + auto& layer = layers_.at(pending_layer_); + const auto bytes = layer.config.expert_bytes; + for (auto copy : copies_) { + // Every input lease and rank completion is resolved before overwriting + // the old canonical input, including overlapping TP boundary readers. + if (free_mode_) { + if (copy.victim >= 0) { + layer.cold[copy.victim] = int64_t(copy.save_offset); + layer.slot[copy.victim] = -1; + } + if (copy.load_offset != kSharedNoOffset) + free_.at(bytes).push_back(copy.load_offset); + else + layer.evicting.erase(copy.incoming); // rescued (or a ghost) + layer.cold[copy.incoming] = -1; + layer.slot[copy.incoming] = copy.slot; + continue; + } + if (copy.save_offset != copy.load_offset) + std::memcpy(arena_ + copy.load_offset, arena_ + copy.save_offset, bytes); + layer.cold[copy.victim] = layer.cold[copy.incoming]; + layer.cold[copy.incoming] = -1; + layer.slot[copy.incoming] = copy.slot; + layer.slot[copy.victim] = -1; + } + // A rescued ghost's saved host copy is no longer canonical. + for (const auto& [expert, at] : ghost_rescues_) + if (at >= 0) free_.at(bytes).push_back(uint64_t(at)); + ghost_rescues_.clear(); + for (const auto& ev : evicts_) { + const SharedExpertJob job{next_job_++, 0, int32_t(pending_layer_), ev.slot, + ev.victim, ev.save_offset, next_ordinal_[0]++}; + layer.slot[ev.victim] = -1; + layer.evicting[ev.victim] = {ev.slot, job.id}; + jobs_[0].push_back(job); + new_jobs_.push_back(job); + } + evicts_.clear(); + copies_.clear(); + pending_ = false; + } + + // Cold (needs a transfer): not resident, not reclaimable without a copy. + bool cold(size_t layer_index, int32_t expert) { + std::lock_guard guard(mutex_); + const auto& layer = layers_.at(layer_index); + return layer.slot.at(expert) < 0 && !layer.ghost.count(expert) && + !layer.evicting.count(expert) && !layer.prefetching.count(expert); + } + + int32_t slot(size_t layer, int32_t expert) { + std::lock_guard guard(mutex_); + return layers_.at(layer).slot.at(expert); + } + + uint64_t cold_offset(size_t layer, int32_t expert) { + std::lock_guard guard(mutex_); + return cold_offset_locked(layer, expert); + } + + std::pair acquire(size_t layer, int32_t expert) { + std::lock_guard guard(mutex_); + const uint64_t offset = cold_offset_locked(layer, expert); + const uint64_t lease = + fence_.acquire({offset, layers_.at(layer).config.expert_bytes}); + return {lease, offset}; + } + + std::vector configs() { + std::lock_guard guard(mutex_); + std::vector out; + for (const auto& layer : layers_) out.push_back(layer.config); + return out; + } + + void release(uint64_t lease) { fence_.release(lease); } + void stop() { fence_.stop(); } + + private: + uint64_t cold_offset_locked(size_t index, int32_t expert) { + const auto& layer = layers_.at(index); + const auto at = layer.cold.at(expert); + if (at < 0) throw std::invalid_argument("expert is resident, not cold"); + return uint64_t(at); + } + + struct Layer { + SharedExpertLayer config; + uint64_t offset; + std::vector slot; + std::vector cold; + std::vector use; + struct Evicting { int32_t slot; uint64_t job; }; + std::vector free_slots; // empty, saved on all ranks + std::unordered_map evicting; // dropped, save in flight + std::unordered_map ghost; // saved, still in a free slot + std::unordered_map prefetching; // load ordinal in flight + }; + uint8_t* const arena_; + const bool inplace_ = [] { + const char* v = std::getenv("VLLM_V41_SWAP_INPLACE"); + return v != nullptr && v[0] == '1'; + }(); + const bool free_mode_ = [] { + const char* v = std::getenv("VLLM_V41_SWAP_FREE"); + return v != nullptr && v[0] == '1'; + }(); + // VLLM_V41_ASYNC_EVICT=: empty GPU slots kept per layer (free mode only). + const int32_t async_free_ = [] { + const char* v = std::getenv("VLLM_V41_ASYNC_EVICT"); + const char* f = std::getenv("VLLM_V41_SWAP_FREE"); + return v && f && f[0] == '1' ? std::atoi(v) : 0; + }(); + static constexpr size_t kEvictReserve = 4; // host allocations kept for exchange + static constexpr size_t kMaxPendingJobs = 160; // < job ring entries + static constexpr size_t kMaxPlanOps = 44; // control page op capacity + std::deque jobs_[2]; + std::unordered_set demand_; + struct Deferred { int32_t layer, slot, expert; uint64_t offset; }; + std::vector deferred_; + uint64_t next_ordinal_[2] = {0, 0}; + struct PairHash { + size_t operator()(const std::pair& v) const { + return std::hash()((int64_t(v.first) << 32) | uint32_t(v.second)); + } + }; + std::unordered_set, PairHash> prefetched_; + std::vector new_jobs_; + std::vector evicts_; + std::vector background_; + uint64_t next_job_ = 0; + std::vector> ghost_rescues_; + public: + uint64_t stat_free_ = 0, stat_exchange_ = 0, stat_rescue_ = 0, stat_ghost_ = 0, + stat_evict_ = 0, stat_prefetch_ = 0, stat_prefetch_hit_ = 0, stat_adopt_ = 0; + private: + std::unordered_map> free_; + const unsigned max_routes_; + static constexpr int32_t kNoTrunc = 1 << 20; + public: + static constexpr int32_t kSkipRoute = 0x7FFFFFFF; + static constexpr int32_t kDeferFlag = 1 << 30; + private: + int32_t trunc_ = kNoTrunc; + uint64_t round_ = 0; + uint64_t exchange_offset_ = 0, clock_ = 0; + size_t pending_layer_ = 0; + bool pending_ = false; + std::mutex mutex_; + SharedExpertSwapState fence_; + std::vector layers_; + std::vector copies_; +}; + +} // namespace vllm::gguf diff --git a/csrc/gguf_shared_expert_state.h b/csrc/gguf_shared_expert_state.h new file mode 100644 index 0000000000..d489cda3f6 --- /dev/null +++ b/csrc/gguf_shared_expert_state.h @@ -0,0 +1,194 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#pragma once + +// Lifetime fence for the TP-group coordinator. This class does not issue CUDA +// work or store process-local pointers. The coordinator owns residency and the +// arena allocator, publishes copies using these 64-bit ranges, and recycles +// input ranges only after finish() succeeds. It is not a per-rank swap lane. +#include +#include +#include +#include +#include +#include +#include +#include + +namespace vllm::gguf { + +struct SharedExpertRange { + uint64_t offset; + uint64_t bytes; +}; + +class SharedExpertSwapState { + public: + SharedExpertSwapState(uint64_t arena_bytes, unsigned ranks) + : arena_bytes_(arena_bytes), ranks_(ranks) { + if (!arena_bytes || !ranks || ranks > 64) + throw std::invalid_argument("invalid shared expert group dimensions"); + } + + // The prefetcher must hold a lease until all its asynchronous reads finish. + // Lease acquisition and publication of a swap are serialized here. No new + // lease can extend the lifetime of an input after that swap is published. + uint64_t acquire(SharedExpertRange range) { + std::lock_guard guard(mutex_); + healthy(); + validate(range); + if (active_) { + for (auto input : inputs_) + if (overlap(range, input)) + throw std::runtime_error( + "prefetch input is fenced by an active swap"); + for (auto output : outputs_) + if (overlap(range, output)) + throw std::runtime_error( + "prefetch cannot read an unfinished writeback"); + } + if (next_lease_ == std::numeric_limits::max()) + poison("shared expert lease sequence exhausted"); + const uint64_t id = ++next_lease_; + leases_.emplace(id, range); + return id; + } + + void release(uint64_t lease) { + std::lock_guard guard(mutex_); + healthy(); + if (leases_.erase(lease) != 1) poison("unknown shared expert read lease"); + } + + // Input and output allocations must be distinct even when all four devices + // appear to execute the same plan. A slow rank may still read any input. + void begin(uint64_t sequence, std::vector inputs, + std::vector outputs, uint64_t now_ns, + uint64_t timeout_ns) { + std::lock_guard guard(mutex_); + healthy(); + if (active_) throw std::runtime_error("shared expert swap already active"); + if (!sequence || sequence_ == std::numeric_limits::max() || + sequence != sequence_ + 1) + poison("shared expert plan sequence mismatch"); + // Free-slot loads have an input without an output (no victim to save). + if (outputs.size() > inputs.size() || !timeout_ns || + timeout_ns > std::numeric_limits::max() - now_ns) + throw std::invalid_argument("invalid shared expert swap plan"); + for (size_t i = 0; i < inputs.size(); ++i) { + validate(inputs[i]); + for (size_t j = 0; j < i; ++j) + if (overlap(inputs[i], inputs[j])) + throw std::invalid_argument("shared expert plan allocations overlap"); + } + for (size_t i = 0; i < outputs.size(); ++i) { + validate(outputs[i]); + for (size_t j = 0; j < inputs.size(); ++j) { + const bool inplace = outputs[i].offset == inputs[j].offset && + outputs[i].bytes == inputs[j].bytes; + if (!inplace && overlap(outputs[i], inputs[j])) + throw std::invalid_argument("shared writeback aliases an input"); + } + for (size_t j = 0; j < i; ++j) + if (overlap(outputs[i], outputs[j])) + throw std::invalid_argument("shared expert plan allocations overlap"); + for (auto& entry : leases_) + if (overlap(outputs[i], entry.second)) + throw std::invalid_argument( + "shared writeback aliases a live prefetch"); + } + sequence_ = sequence; + deadline_ns_ = now_ns + timeout_ns; + inputs_ = std::move(inputs); + outputs_ = std::move(outputs); + completed_ = 0; + active_ = true; + } + + // A rank acknowledges only after its reads AND its unique-owner writes + // complete. CPU completion publication uses acquire/release outside this + // object; receiving a CUDA launch notification is insufficient. + void acknowledge(unsigned rank, uint64_t sequence, uint64_t now_ns) { + std::lock_guard guard(mutex_); + check_plan(sequence, now_ns); + if (rank >= ranks_) poison("invalid shared expert completion rank"); + const uint64_t bit = uint64_t{1} << rank; + if (completed_ & bit) poison("duplicate shared expert completion"); + completed_ |= bit; + } + + bool ready(uint64_t sequence, uint64_t now_ns) { + std::lock_guard guard(mutex_); + check_plan(sequence, now_ns); + return ready_locked(); + } + + // Returned ranges are the only old input allocations the caller may recycle. + // Retain all allocations on failure: devices may already have modified slots, + // so an aborted plan requires explicit cache reinitialization. + std::vector finish(uint64_t sequence, uint64_t now_ns) { + std::lock_guard guard(mutex_); + check_plan(sequence, now_ns); + if (!ready_locked()) + throw std::runtime_error("shared expert readers have not completed"); + active_ = false; + outputs_.clear(); + return std::exchange(inputs_, {}); + } + + void stop() { + std::lock_guard guard(mutex_); + failure_ = "shared expert coordinator stopped"; + } + + private: + static bool overlap(SharedExpertRange a, SharedExpertRange b) { + // Validated ranges cannot wrap. Use subtraction to remain overflow-safe + // even while checking an as-yet unvalidated prospective plan. + return a.offset <= b.offset ? b.offset - a.offset < a.bytes + : a.offset - b.offset < b.bytes; + } + + void validate(SharedExpertRange range) const { + if (!range.bytes || range.offset >= arena_bytes_ || + range.bytes > arena_bytes_ - range.offset) + throw std::invalid_argument("shared expert range exceeds arena"); + } + + void healthy() const { + if (!failure_.empty()) throw std::runtime_error(failure_); + } + + [[noreturn]] void poison(const char* reason) { + failure_ = reason; + throw std::runtime_error(failure_); + } + + void check_plan(uint64_t sequence, uint64_t now_ns) { + healthy(); + if (!active_ || sequence != sequence_) + poison("shared expert completion sequence mismatch"); + if (now_ns >= deadline_ns_) poison("shared expert plan timed out"); + } + + bool ready_locked() const { + const uint64_t expected = + ranks_ == 64 ? ~uint64_t{0} : (uint64_t{1} << ranks_) - 1; + if (completed_ != expected) return false; + for (auto& entry : leases_) + for (auto input : inputs_) + if (overlap(entry.second, input)) return false; + return true; + } + + const uint64_t arena_bytes_; + const unsigned ranks_; + std::mutex mutex_; + uint64_t next_lease_ = 0, sequence_ = 0, deadline_ns_ = 0, completed_ = 0; + bool active_ = false; + std::string failure_; + std::vector inputs_, outputs_; + std::unordered_map leases_; +}; + +} // namespace vllm::gguf diff --git a/csrc/libtorch_stable/quantization/gguf_upstream/runtime_dequant_blas.cu b/csrc/libtorch_stable/quantization/gguf_upstream/runtime_dequant_blas.cu index a0a9655145..c6876414e8 100644 --- a/csrc/libtorch_stable/quantization/gguf_upstream/runtime_dequant_blas.cu +++ b/csrc/libtorch_stable/quantization/gguf_upstream/runtime_dequant_blas.cu @@ -2,6 +2,7 @@ #include "ggml_dypes.cuh" #include "torch_context.cuh" #include "convert.cuh" +#include "dequantize.cuh" #include "ternary_dequant.cuh" #include @@ -13,15 +14,42 @@ #include #include #include +#include #include #include using torch::headeronly::ScalarType; namespace { +template +__global__ void dequantize_q5_k_blocks(const void* source, scalar_t* output, + int64_t blocks) { + // Reuse upstream arithmetic verbatim. Four independent 64-thread + // groups share a bounded persistent grid instead of scheduling one CTA + // per GGUF block. The groups have no cross-warp dependencies. + for (int64_t block = int64_t(blockIdx.x) * 4 + threadIdx.x / 64; + block < blocks; block += int64_t(gridDim.x) * 4) + dequantize_q5_K(source, block, output + block * QK_K, threadIdx.x % 64); +} + template void run_upstream_dequantize(const Tensor& W, Tensor& output, int64_t type, - int64_t total, cudaStream_t stream) { + int64_t total, cudaStream_t stream, + bool batched_blocks = false) { + if constexpr (std::is_same_v) { + if (batched_blocks && type == GGML_TYPE_Q5_K && + ggml_cuda_info().devices[W.get_device_index()].cc == + GGML_CUDA_CC_VOLTA) { + const int64_t blocks = total / QK_K; + const int64_t grid = std::min( + (blocks + 3) / 4, + int64_t(ggml_cuda_info().devices[W.get_device_index()].nsm) * 8); + dequantize_q5_k_blocks<<>>( + W.data_ptr(), static_cast(output.data_ptr()), blocks); + check_launch("batched Q5_K dequantization"); + return; + } + } if (type == GGML_TYPE_TQ1_0 || type == GGML_TYPE_TQ2_0) { const int64_t blocks = total / 256; if (type == GGML_TYPE_TQ1_0) { @@ -60,6 +88,109 @@ ScalarType blas_compute_dtype(int cc) { : ScalarType::Float; } +__device__ __forceinline__ void decode_q5_quarter_vector( + const void* weights, int64_t block, half* output, int quarter, int lane) { + const auto& w = static_cast(weights)[block]; + const uint32_t low = + reinterpret_cast(w.qs + quarter * 32)[lane]; + const uint32_t high = reinterpret_cast(w.qh)[lane]; + uint8_t scale, minimum; + get_scale_min_k4(quarter * 2, w.scales, scale, minimum); + const float d1 = __half2float(__low2half(w.dm)) * scale; + const float m1 = __half2float(__high2half(w.dm)) * minimum; + get_scale_min_k4(quarter * 2 + 1, w.scales, scale, minimum); + const float d2 = __half2float(__low2half(w.dm)) * scale; + const float m2 = __half2float(__high2half(w.dm)) * minimum; + half lower[4], upper[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int q = (low >> (8 * i)) & 255; + const int h = (high >> (8 * i)) & 255; + lower[i] = ggml_cuda_cast( + d1 * ((q & 15) + ((h >> (2 * quarter)) & 1) * 16) - m1); + upper[i] = ggml_cuda_cast( + d2 * ((q >> 4) + ((h >> (2 * quarter + 1)) & 1) * 16) - m2); + } + auto* first = reinterpret_cast(output + lane * 4); + first[0] = __halves2half2(lower[0], lower[1]); + first[1] = __halves2half2(lower[2], lower[3]); + first[16] = __halves2half2(upper[0], upper[1]); + first[17] = __halves2half2(upper[2], upper[3]); +} + +template +__global__ void q5_k_hmma_partials(const void* weights, const half* input, + float* partials, int m, int n, int k, + int splits, int groups) { +#if __CUDA_ARCH__ >= 700 + namespace wmma = nvcuda::wmma; + // V4.1's vector path pads shared rows to spread WMMA column accesses + // across banks. K64 tiles, fragments and accumulation order stay unchanged. + __shared__ __align__(32) half activation[16 * Stride]; + // A prefix makes all compensated upstream destination pointers stay + // inside this allocation while decoding just one 64-value quarter. + __shared__ __align__(32) half decoded[256 + 64 * Stride]; + const int warp = threadIdx.x / 32; + const int input_groups = Grouped ? groups : 1; + const int input_group = Grouped ? blockIdx.x / (n / groups / 64) : 0; + const int first = (k / 64) * blockIdx.y / splits; + const int last = (k / 64) * (blockIdx.y + 1) / splits; + wmma::fragment a; + wmma::fragment b; + wmma::fragment accum; + wmma::fill_fragment(accum, 0.f); + for (int tile = first; tile < last; ++tile) { + for (int i = threadIdx.x; i < 16 * 64; i += blockDim.x) { + const int row = i / 64, col = i % 64; + activation[row * Stride + col] = + row < m ? input[(int64_t(row) * input_groups + input_group) * k + + tile * 64 + col] + : __float2half_rn(0.f); + } + const int quarter = tile % 4; + if constexpr (Vector) { + for (int row = threadIdx.x / 8; row < 64; row += blockDim.x / 8) { + const int64_t block = + int64_t(blockIdx.x * 64 + row) * (k / QK_K) + tile / 4; + decode_q5_quarter_vector(weights, block, decoded + 256 + row * Stride, + quarter, threadIdx.x % 8); + } + } else { + for (int row = threadIdx.x / 16; row < 64; row += blockDim.x / 16) { + const int64_t block = + int64_t(blockIdx.x * 64 + row) * (k / QK_K) + tile / 4; + dequantize_q5_K(weights, block, + decoded + 256 + row * Stride - quarter * 64, + quarter * 16 + threadIdx.x % 16); + } + } + __syncthreads(); + #pragma unroll + for (int offset = 0; offset < 64; offset += 16) { + wmma::load_matrix_sync(a, activation + offset, Stride); + wmma::load_matrix_sync(b, decoded + 256 + warp * 16 * Stride + offset, + Stride); + wmma::mma_sync(accum, a, b, accum); + } + __syncthreads(); + } + wmma::store_matrix_sync( + partials + int64_t(blockIdx.y) * 16 * n + blockIdx.x * 64 + warp * 16, + accum, n, wmma::mem_row_major); +#endif +} + +__global__ void q5_k_hmma_reduce(const float* partials, half* output, int m, + int n, int splits) { + const int index = blockIdx.x * blockDim.x + threadIdx.x; + if (index >= m * n) return; + const int row = index / n, col = index % n; + float value = 0.f; + for (int split = 0; split < splits; ++split) + value = __fadd_rn(value, partials[(int64_t(split) * 16 + row) * n + col]); + output[index] = __float2half_rn(value); +} + } // namespace Tensor run_upstream_blas(const Tensor& W, const Tensor& X, int64_t type, @@ -81,7 +212,7 @@ Tensor run_upstream_blas(const Tensor& W, const Tensor& X, int64_t type, if (compute_dtype == ScalarType::Float) { run_upstream_dequantize(W, weights, type, row * k, stream); } else if (compute_dtype == ScalarType::Half) { - run_upstream_dequantize(W, weights, type, row * k, stream); + run_upstream_dequantize(W, weights, type, row * k, stream, true); } else { run_upstream_dequantize(W, weights, type, row * k, stream); } @@ -141,7 +272,8 @@ Tensor run_upstream_blas(const Tensor& W, const Tensor& X, int64_t type, } Tensor ggml_dequantize_upstream(Tensor W, int64_t type, int64_t m, int64_t n, - std::optional dtype) { + std::optional dtype, + bool batched_blocks) { STD_TORCH_CHECK(W.is_cuda(), "ggml_dequantize_upstream: W must be a CUDA tensor"); STD_TORCH_CHECK(W.dim() == 2 && W.is_contiguous(), @@ -206,10 +338,74 @@ Tensor ggml_dequantize_upstream(Tensor W, int64_t type, int64_t m, int64_t n, if (dtype_ == ScalarType::Float) { run_upstream_dequantize(W, output, type, total, stream); } else if (dtype_ == ScalarType::Half) { - run_upstream_dequantize(W, output, type, total, stream); + run_upstream_dequantize(W, output, type, total, stream, + batched_blocks); } else { run_upstream_dequantize(W, output, type, total, stream); } check_launch("upstream dequantize"); return output; } + +Tensor ggml_dense_q5_k_hmma(Tensor W, Tensor X, int64_t split_k, + int64_t workspace_bytes, bool grouped, + bool vector_decode) { + const int expected_dim = grouped ? 3 : 2; + STD_TORCH_CHECK( + W.is_cuda() && X.is_cuda() && + W.get_device_index() == X.get_device_index() && + W.scalar_type() == ScalarType::Byte && W.dim() == 2 && + W.is_contiguous() && X.scalar_type() == ScalarType::Half && + X.dim() == expected_dim && X.is_contiguous() && X.size(0) > 0 && + X.size(0) <= 6 && W.size(0) > 0 && W.size(0) <= 65536 && + W.size(0) % 64 == 0 && X.size(-1) > 0 && X.size(-1) <= 16384 && + X.size(-1) % QK_K == 0 && + W.size(1) == X.size(-1) / QK_K * sizeof(block_q5_K), + "Q5_K HMMA requires original packed rows and FP16 M1..6"); + const int groups = grouped ? X.size(1) : 1; + STD_TORCH_CHECK(groups > 0 && groups <= 8 && W.size(0) % (groups * 64) == 0, + "Q5_K grouped HMMA requires whole 64-row output groups"); + const auto device = X.get_device_index(); + const DeviceGuard guard(device); + const auto& info = ggml_cuda_info().devices[device]; + STD_TORCH_CHECK(info.cc == GGML_CUDA_CC_VOLTA, "Q5_K HMMA requires SM70"); + const int m = X.size(0), n = W.size(0), k = X.size(-1); + STD_TORCH_CHECK(workspace_bytes > 0 && workspace_bytes <= 32 * 1024 * 1024, + "Q5_K HMMA workspace must be in (0, 32 MiB]"); + const int capacity = workspace_bytes / (16 * n * sizeof(float)); + const int automatic = (info.nsm * 8 + n / 64 - 1) / (n / 64); + const int splits = + split_k ? split_k : std::min({64, k / 64, automatic, capacity}); + STD_TORCH_CHECK( + splits > 0 && splits <= 64 && splits <= k / 64 && + int64_t(splits) * 16 * n * sizeof(float) <= workspace_bytes, + "Q5_K HMMA partial workspace exceeds its 32 MiB bound"); + Tensor partials = + torch::stable::new_empty(X, {splits, 16, n}, ScalarType::Float); + Tensor output = torch::stable::new_empty(X, {m, n}, ScalarType::Half); + const auto stream = current_stream(device); + // Keep the original full-N split count and summation order. Each output + // group reads only its own input instead of computing every group and + // discarding the other rows afterward. + if (grouped && vector_decode) + q5_k_hmma_partials<<>>( + W.data_ptr(), static_cast(X.data_ptr()), + static_cast(partials.data_ptr()), m, n, k, splits, groups); + else if (grouped) + q5_k_hmma_partials<<>>( + W.data_ptr(), static_cast(X.data_ptr()), + static_cast(partials.data_ptr()), m, n, k, splits, groups); + else if (vector_decode) + q5_k_hmma_partials<<>>( + W.data_ptr(), static_cast(X.data_ptr()), + static_cast(partials.data_ptr()), m, n, k, splits, 1); + else + q5_k_hmma_partials<<>>( + W.data_ptr(), static_cast(X.data_ptr()), + static_cast(partials.data_ptr()), m, n, k, splits, 1); + q5_k_hmma_reduce<<<(m * n + 255) / 256, 256, 0, stream>>>( + static_cast(partials.data_ptr()), + static_cast(output.data_ptr()), m, n, splits); + check_launch("original Q5_K FP16 HMMA"); + return output; +} diff --git a/csrc/libtorch_stable/quantization/gguf_upstream/torch_bindings.cpp b/csrc/libtorch_stable/quantization/gguf_upstream/torch_bindings.cpp index 02895faebe..c0e95a134e 100644 --- a/csrc/libtorch_stable/quantization/gguf_upstream/torch_bindings.cpp +++ b/csrc/libtorch_stable/quantization/gguf_upstream/torch_bindings.cpp @@ -18,8 +18,12 @@ Tensor ggml_dense_mmvq(Tensor W, Tensor X, int64_t type, int64_t row); Tensor ggml_dense_mmq(Tensor W, Tensor X, int64_t type, int64_t row); Tensor ggml_dense_mmf(Tensor W, Tensor X, int64_t type, int64_t row); Tensor ggml_dense_blas(Tensor W, Tensor X, int64_t type, int64_t row); +Tensor ggml_dense_q5_k_hmma(Tensor W, Tensor X, int64_t split_k, + int64_t workspace_bytes, bool grouped, + bool vector_decode); Tensor ggml_dequantize_upstream(Tensor W, int64_t type, int64_t m, int64_t n, - std::optional dtype); + std::optional dtype, + bool batched_blocks); Tensor ggml_moe_upstream(Tensor X, Tensor W, Tensor ids, int64_t type, int64_t row, int64_t top_k, int64_t tokens); Tensor ggml_moe_mmvq(Tensor X, Tensor W, Tensor ids, int64_t type, int64_t row, @@ -44,9 +48,13 @@ STABLE_TORCH_LIBRARY(_C_gguf, ops) { ops.def("ggml_dense_mmf(Tensor W, Tensor X, int type, SymInt row) -> Tensor"); ops.def( "ggml_dense_blas(Tensor W, Tensor X, int type, SymInt row) -> Tensor"); + ops.def( + "ggml_dense_q5_k_hmma(Tensor W, Tensor X, int split_k=0, " + "int workspace_bytes=33554432, bool grouped=False, " + "bool vector_decode=False) -> Tensor"); ops.def( "ggml_dequantize_upstream(Tensor W, int type, SymInt m, SymInt n, " - "ScalarType? dtype) -> Tensor"); + "ScalarType? dtype, bool batched_blocks=False) -> Tensor"); ops.def( "ggml_moe_upstream(Tensor X, Tensor W, Tensor ids, int type, SymInt row, " "SymInt top_k, SymInt tokens) -> Tensor"); @@ -67,6 +75,7 @@ STABLE_TORCH_LIBRARY_IMPL(_C_gguf, CUDA, ops) { ops.impl("ggml_dense_mmq", TORCH_BOX(&ggml_dense_mmq)); ops.impl("ggml_dense_mmf", TORCH_BOX(&ggml_dense_mmf)); ops.impl("ggml_dense_blas", TORCH_BOX(&ggml_dense_blas)); + ops.impl("ggml_dense_q5_k_hmma", TORCH_BOX(&ggml_dense_q5_k_hmma)); ops.impl("ggml_dequantize_upstream", TORCH_BOX(&ggml_dequantize_upstream)); ops.impl("ggml_moe_upstream", TORCH_BOX(&ggml_moe_upstream)); ops.impl("ggml_moe_mmvq", TORCH_BOX(&ggml_moe_mmvq)); diff --git a/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu b/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu index a0686c97cf..da603bd0cf 100644 --- a/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu +++ b/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu @@ -9539,9 +9539,19 @@ void mxfp4_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, input.size(0) == num_experts && num_experts >= 12 && num_experts <= 48 && num_experts % 6 == 0 && ((k == 4096 && n == 512) || (k == 256 && n == 4096)); + // Any small one-row-per-route batch (e.g. the V4.1 DSpark draft, M5 x + // top-3 = 15 rows) runs as one grouped launch instead of a launch per row. + static const bool grouped_any = [] { + const char* v = std::getenv("VLLM_SM70_MXFP4_GROUPED_ANY"); + return v != nullptr && v[0] == '1'; + }(); + const bool grouped_any_shape = grouped_any && + input.size(0) == num_experts && + num_experts >= 2 && num_experts <= 48; if ((vllm::awq_sm70::mxfp4_moe_grouped_m8_enabled() && grouped_m8_shape) || (vllm::awq_sm70::mxfp4_moe_grouped_verifier_enabled() && - grouped_verifier_shape)) { + grouped_verifier_shape) || + grouped_any_shape) { // Keep every routed slot as an independent one-row group. Repeated experts // reuse the same weight pointer row, while the scheduler retains the exact // arithmetic contract already accepted for compact B1 decode. diff --git a/csrc/sm70_turbomind/ops/gguf_dp4a.cuh b/csrc/sm70_turbomind/ops/gguf_dp4a.cuh index 638d864c6f..6b941a23fd 100644 --- a/csrc/sm70_turbomind/ops/gguf_dp4a.cuh +++ b/csrc/sm70_turbomind/ops/gguf_dp4a.cuh @@ -128,6 +128,179 @@ struct LatticeDot { }; using IQ3SDot = LatticeDot<21>; +// V4.1's IQ1_M/IQ2 rows consume one Q8_1 group (32 activations) per call. +// IQ1's eighth-step delta is folded into signed integer codebook values, so +// both local scale groups remain exact integer dots before the shared scale. +template +struct V41LatticeDot { + static_assert(Type == 16 || Type == 17 || Type == 29); + using Book = turbomind::gemm::LatticeCodebook; + static constexpr int kBookWords = Book::kBytes / 4; + static constexpr int kBlockBytes = Type == 16 ? 66 : Type == 17 ? 74 : 56; + + __device__ static void initialize(uint32_t* book, uint32_t* masks) { + for (int i = threadIdx.x; i < kBookWords; i += blockDim.x) { + uint32_t word = Book::word(i) ^ 0x80808080U; + if constexpr (Type == 29) { + // IQ1's eighth-step grid is shared by every dot in this CTA. + word = __vadd4(word, word); + word = __vadd4(word, word); + word = __vadd4(word, word); + } + book[i] = word; + } + if (threadIdx.x < 16) + masks[threadIdx.x] = LatticeDot<22>::sign_mask(threadIdx.x); + __syncthreads(); + } + + __device__ static float dot(const uint8_t* row, int group, const Q8_1& x, + const uint32_t* book, const uint32_t* masks) { + const uint8_t* b = row + (group / 8) * kBlockBytes; + const int sub = group % 8; + const int* activation = reinterpret_cast(x.qs); + int sums[2] = {0, 0}; + const uint32_t meta = Type == 16 ? load_u32_2(b + 6 + sub * 8) : 0; +#pragma unroll + for (int octet = 0; octet < 4; ++octet) { + int index, signs = 0, delta = 0; + if constexpr (Type == 16) { + index = b[2 + sub * 8 + octet]; + signs = (meta >> (7 * octet)) & 127; + } else if constexpr (Type == 17) { + const uint16_t code = + *reinterpret_cast(b + 2 + (sub * 4 + octet) * 2); + index = code & 511; + signs = code >> 9; + } else { + const int t = sub * 4 + octet; + const int high = (b[32 + t / 2] >> (4 * (t % 2))) & 15; + index = b[t] | ((high & 7) << 8); + delta = high >> 3; + } + uint32_t w0 = book[index * 2], w1 = book[index * 2 + 1]; + if constexpr (Type == 29) { + const uint32_t bias = delta ? 0xffffffffU : 0x01010101U; + w0 = __vadd4(w0, bias); + w1 = __vadd4(w1, bias); + } else { + signs |= (__popc(signs) & 1) << 7; + const auto s0 = masks[signs & 15], s1 = masks[signs >> 4]; + w0 = __vsub4(w0 ^ s0, s0); + w1 = __vsub4(w1 ^ s1, s1); + } + sums[octet / 2] = __dp4a(int(w0), activation[2 * octet], sums[octet / 2]); + sums[octet / 2] = + __dp4a(int(w1), activation[2 * octet + 1], sums[octet / 2]); + } + float d; + int mult0, mult1; + if constexpr (Type == 29) { + const auto* scales = reinterpret_cast(b + 48); + const uint16_t bits = (scales[0] >> 12) | ((scales[1] >> 8) & 0xf0) | + ((scales[2] >> 4) & 0xf00) | (scales[3] & 0xf000); + d = __half2float(__ushort_as_half(bits)); + const int first = sub * 2, second = first + 1; + mult0 = 2 * ((scales[first / 4] >> (3 * (first % 4))) & 7) + 1; + mult1 = 2 * ((scales[second / 4] >> (3 * (second % 4))) & 7) + 1; + } else { + d = __half2float(*reinterpret_cast(b)); + if constexpr (Type == 16) + mult0 = mult1 = 1 + 2 * (meta >> 28); + else { + mult0 = 1 + 2 * (b[66 + sub] & 15); + mult1 = 1 + 2 * (b[66 + sub] >> 4); + } + } + const float scale = __fmul_rn(d, __low2float(x.ds)); + return __fmul_rn(__fmul_rn(scale, float(sums[0] * mult0 + sums[1] * mult1)), + 0.125f); + } +}; + +template +struct LatticeDot<16, BankAware> : V41LatticeDot<16, BankAware> {}; +template +struct LatticeDot<17, BankAware> : V41LatticeDot<17, BankAware> {}; +template +struct LatticeDot<29, BankAware> : V41LatticeDot<29, BankAware> {}; + +// The routed rows can share one expert. Decode the original integer words +// once, then apply each token's Q8_1 scale and exactly the original dot order. +template +struct V41IntegerGroup { + uint32_t words[8]; + float d; + int mult0, mult1; + + __device__ V41IntegerGroup(const uint8_t* row, int group, + const uint32_t* book, const uint32_t* masks) { + constexpr int bytes = V41LatticeDot::kBlockBytes; + const uint8_t* b = row + (group / 8) * bytes; + const int sub = group % 8; + const uint32_t meta = Type == 16 ? load_u32_2(b + 6 + sub * 8) : 0; +#pragma unroll + for (int octet = 0; octet < 4; ++octet) { + int index, signs = 0, delta = 0; + if constexpr (Type == 16) { + index = b[2 + sub * 8 + octet]; + signs = (meta >> (7 * octet)) & 127; + } else if constexpr (Type == 17) { + const auto code = + *reinterpret_cast(b + 2 + (sub * 4 + octet) * 2); + index = code & 511; + signs = code >> 9; + } else { + const int t = sub * 4 + octet; + const int high = (b[32 + t / 2] >> (4 * (t % 2))) & 15; + index = b[t] | ((high & 7) << 8); + delta = high >> 3; + } + uint32_t first = book[index * 2], second = book[index * 2 + 1]; + if constexpr (Type == 29) { + const uint32_t bias = delta ? 0xffffffffU : 0x01010101U; + first = __vadd4(first, bias); + second = __vadd4(second, bias); + } else { + signs |= (__popc(signs) & 1) << 7; + const auto low = masks[signs & 15], high = masks[signs >> 4]; + first = __vsub4(first ^ low, low); + second = __vsub4(second ^ high, high); + } + words[2 * octet] = first; + words[2 * octet + 1] = second; + } + if constexpr (Type == 29) { + const auto* scales = reinterpret_cast(b + 48); + const uint16_t bits = (scales[0] >> 12) | ((scales[1] >> 8) & 0xf0) | + ((scales[2] >> 4) & 0xf00) | (scales[3] & 0xf000); + d = __half2float(__ushort_as_half(bits)); + const int first = sub * 2, second = first + 1; + mult0 = 2 * ((scales[first / 4] >> (3 * (first % 4))) & 7) + 1; + mult1 = 2 * ((scales[second / 4] >> (3 * (second % 4))) & 7) + 1; + } else { + d = __half2float(*reinterpret_cast(b)); + if constexpr (Type == 16) + mult0 = mult1 = 1 + 2 * (meta >> 28); + else { + mult0 = 1 + 2 * (b[66 + sub] & 15); + mult1 = 1 + 2 * (b[66 + sub] >> 4); + } + } + } + + __device__ float apply(const Q8_1& x) const { + const auto* a = reinterpret_cast(x.qs); + int sums[2] = {0, 0}; +#pragma unroll + for (int i = 0; i < 8; ++i) + sums[i / 4] = __dp4a(int(words[i]), a[i], sums[i / 4]); + const float scale = __fmul_rn(d, __low2float(x.ds)); + return __fmul_rn(__fmul_rn(scale, float(sums[0] * mult0 + sums[1] * mult1)), + 0.125f); + } +}; + // One IQ3_S weight group decoded for several activations. apply() repeats // dot()'s integer dot and float boundary exactly, so a route's result does // not depend on how many routes share the decoded words. diff --git a/csrc/sm70_turbomind/ops/gguf_dp4a_sm70.cu b/csrc/sm70_turbomind/ops/gguf_dp4a_sm70.cu index a66dae5532..daf0385c45 100644 --- a/csrc/sm70_turbomind/ops/gguf_dp4a_sm70.cu +++ b/csrc/sm70_turbomind/ops/gguf_dp4a_sm70.cu @@ -21,14 +21,48 @@ __global__ void quantize_q8(Q8_1* out, const half* input, int k) { vllm::sm70_gguf::quantize_q8_1_warp(out + group, value); } +__device__ float v41_swiglu(float gate, float up, float probability) { + gate = gate > 10.f ? 10.f : gate; + up = up > 10.f ? 10.f : up; + up = up < -10.f ? -10.f : up; + const float activated = __fdiv_rn(gate, __fadd_rn(1.f, expf(-gate))); + return __fmul_rn(__fmul_rn(activated, up), probability); +} + +__global__ void v41_swiglu_q8(Q8_1* out, const half* projection, + const float* probabilities) { + const int col = blockIdx.x * blockDim.x + threadIdx.x; + if (col >= 576) + return; // Entire inactive warps; the 32-lane Q8 group stays full. + const int route = blockIdx.y; + const float gate = __half2float(projection[int64_t(route) * 1152 + col]); + const float up = __half2float(projection[int64_t(route) * 1152 + 576 + col]); + // Retain the FP16 down-input boundary before computing Q8_1 bytes. + const float value = + __half2float(__float2half_rn(v41_swiglu(gate, up, probabilities[route]))); + vllm::sm70_gguf::quantize_q8_1_warp(out + route * 18 + col / 32, value); +} + +__global__ void v41_shared_swiglu(half* out, const half* gate, const half* up, + int64_t count) { + const int64_t index = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + if (index < count) + out[index] = __float2half_rn( + v41_swiglu(__half2float(gate[index]), __half2float(up[index]), 1.f)); +} + template + bool ScalarLut = false, bool BankAware = false, int FixedGroups = 0, + bool RouteGuard = false> __global__ void gate_up(void* output, const Q8_1* activation, const Index* ids, const uint8_t* gate, const uint8_t* up, int n, int k, int stride, int top_k, const StridedPtr* gate_stats = nullptr, - const StridedPtr* up_stats = nullptr) { + const StridedPtr* up_stats = nullptr, + const int* reuse_mode = nullptr) { + if constexpr (RouteGuard) + if (*reuse_mode) return; using Dot = std::conditional_t< Canonical, vllm::sm70_gguf::CanonicalIntegerDot, std::conditional_t, @@ -78,8 +112,13 @@ __global__ void gate_up(void* output, const Q8_1* activation, const Index* ids, us += Dot::dot(up_ptrs[expert].ptr, up_stats[expert].ptr, n, k, row, group, shared_x[group]); } else { - gs += Dot::dot(g, group, shared_x[group], book, masks); - us += Dot::dot(u, group, shared_x[group], book, masks); + if constexpr (Type == 16 || Type == 17 || Type == 29) { + gs = __fadd_rn(gs, Dot::dot(g, group, shared_x[group], book, masks)); + us = __fadd_rn(us, Dot::dot(u, group, shared_x[group], book, masks)); + } else { + gs += Dot::dot(g, group, shared_x[group], book, masks); + us += Dot::dot(u, group, shared_x[group], book, masks); + } } } } @@ -170,39 +209,49 @@ template void launch_gate_up(torch::Tensor out, torch::Tensor activation, torch::Tensor ids, torch::Tensor gate, torch::Tensor up, bool activated, int lanes) { - if (out.scalar_type() == torch::kUInt8) { - if (lanes == 0) // K2560, 16 lanes, unrolled bank-aware schedule - launch_quantized_gate_up(out, activation, ids, - gate, up); - else if (lanes == 4) - launch_quantized_gate_up(out, activation, ids, - gate, up); - else if (lanes == 8) - launch_quantized_gate_up(out, activation, ids, - gate, up); + if constexpr (Type == 16 || Type == 17 || Type == 29) { + const int n = gate.size(1); + gate_up<<>>( + out.data_ptr(), reinterpret_cast(activation.data_ptr()), + ids.data_ptr(), gate.data_ptr(), up.data_ptr(), + n, activation.size(1) * 32, gate.size(2), ids.size(1)); + } else { + if (out.scalar_type() == torch::kUInt8) { + if (lanes == 0) // K2560, 16 lanes, unrolled bank-aware schedule + launch_quantized_gate_up(out, activation, + ids, gate, up); + else if (lanes == 4) + launch_quantized_gate_up(out, activation, + ids, gate, up); + else if (lanes == 8) + launch_quantized_gate_up(out, activation, + ids, gate, up); + else + launch_quantized_gate_up(out, activation, + ids, gate, up); + return; + } + const int n = gate.size(1), top_k = ids.size(1); + const dim3 grid((n + 7) / 8, activation.size(0) * top_k); + const auto stream = at::cuda::getCurrentCUDAStream(); + const size_t shared = activation.size(1) * sizeof(Q8_1); + const auto output = reinterpret_cast(out.data_ptr()); + const auto x = reinterpret_cast(activation.data_ptr()); + if (activated) + gate_up + <<>>( + output, x, ids.data_ptr(), gate.data_ptr(), + up.data_ptr(), n, activation.size(1) * 32, gate.size(2), + top_k); else - launch_quantized_gate_up(out, activation, ids, - gate, up); - return; + gate_up + <<>>( + output, x, ids.data_ptr(), gate.data_ptr(), + up.data_ptr(), n, activation.size(1) * 32, gate.size(2), + top_k); } - const int n = gate.size(1), top_k = ids.size(1); - const dim3 grid((n + 7) / 8, activation.size(0) * top_k); - const auto stream = at::cuda::getCurrentCUDAStream(); - const size_t shared = activation.size(1) * sizeof(Q8_1); - const auto output = reinterpret_cast(out.data_ptr()); - const auto x = reinterpret_cast(activation.data_ptr()); - if (activated) - gate_up - <<>>( - output, x, ids.data_ptr(), gate.data_ptr(), - up.data_ptr(), n, activation.size(1) * 32, gate.size(2), - top_k); - else - gate_up - <<>>( - output, x, ids.data_ptr(), gate.data_ptr(), - up.data_ptr(), n, activation.size(1) * 32, gate.size(2), - top_k); } template @@ -353,9 +402,12 @@ void require_sm70() { void gguf_quantize_q8_1_sm70_out(torch::Tensor out, torch::Tensor input) { TORCH_CHECK(input.is_cuda() && input.dim() == 2 && input.is_contiguous() && - input.scalar_type() == torch::kFloat16 && - input.size(0) <= 20 && input.size(1) > 0 && - input.size(1) % 256 == 0 && out.device() == input.device() && + input.scalar_type() == torch::kFloat16 && input.size(1) > 0 && + ((input.size(0) <= 20 && input.size(1) % 256 == 0) || + ((input.size(0) > 0 && input.size(0) <= 36 && + input.size(0) % 6 == 0) && + input.size(1) == 576)) && + out.device() == input.device() && out.scalar_type() == torch::kUInt8 && out.is_contiguous() && out.dim() == 3 && out.size(0) == input.size(0) && out.size(1) == input.size(1) / 32 && @@ -371,15 +423,68 @@ void gguf_quantize_q8_1_sm70_out(torch::Tensor out, torch::Tensor input) { C10_CUDA_KERNEL_LAUNCH_CHECK(); } +void gguf_v41_swiglu_q8_sm70_out(torch::Tensor out, torch::Tensor projection, + torch::Tensor probabilities) { + TORCH_CHECK(projection.is_cuda() && projection.is_contiguous() && + projection.scalar_type() == torch::kFloat16 && + projection.dim() == 4 && projection.size(0) > 0 && + projection.size(0) <= 6 && projection.size(1) == 6 && + projection.size(2) == 2 && projection.size(3) == 576 && + out.device() == projection.device() && out.is_contiguous() && + out.scalar_type() == torch::kUInt8 && out.dim() == 4 && + out.size(0) == projection.size(0) && out.size(1) == 6 && + out.size(2) == 18 && out.size(3) == sizeof(Q8_1) && + probabilities.device() == projection.device() && + probabilities.is_contiguous() && + probabilities.scalar_type() == torch::kFloat32 && + probabilities.dim() == 2 && + probabilities.size(0) == projection.size(0) && + probabilities.size(1) == 6, + "V4.1 SwiGLU Q8 requires FP16 M1..6 projections and FP32 routes"); + const c10::cuda::CUDAGuard guard(projection.device()); + require_sm70(); + v41_swiglu_q8<<>>( + reinterpret_cast(out.data_ptr()), + reinterpret_cast(projection.data_ptr()), + probabilities.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void deepseek_v41_swiglu_sm70_out(torch::Tensor out, torch::Tensor gate, + torch::Tensor up) { + TORCH_CHECK(gate.is_cuda() && gate.is_contiguous() && gate.dim() == 2 && + gate.scalar_type() == torch::kFloat16 && + up.device() == gate.device() && up.is_contiguous() && + up.scalar_type() == torch::kFloat16 && + up.sizes() == gate.sizes() && out.device() == gate.device() && + out.is_contiguous() && out.scalar_type() == torch::kFloat16 && + out.sizes() == gate.sizes(), + "V4.1 shared SwiGLU requires matching contiguous FP16 matrices"); + const c10::cuda::CUDAGuard guard(gate.device()); + require_sm70(); + if (!gate.numel()) return; + v41_shared_swiglu<<<(gate.numel() + 255) / 256, 256, 0, + at::cuda::getCurrentCUDAStream()>>>( + reinterpret_cast(out.data_ptr()), + reinterpret_cast(gate.data_ptr()), + reinterpret_cast(up.data_ptr()), gate.numel()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + void gguf_dp4a_gate_up_sm70_out(torch::Tensor out, torch::Tensor activation, torch::Tensor ids, torch::Tensor gate, torch::Tensor up, int64_t source_type, bool activated, int64_t lanes_per_row, bool bank_aware) { TORCH_CHECK(source_type == 18 || source_type == 20 || source_type == 21 || - source_type == 22 || source_type == 23, + source_type == 22 || source_type == 23 || source_type == 16 || + source_type == 17 || source_type == 29, "Unsupported GGUF lattice dp4a reader"); - const int block_bytes = source_type == 18 ? 98 + const int block_bytes = source_type == 16 ? 66 + : source_type == 17 ? 74 + : source_type == 29 ? 56 + : source_type == 18 ? 98 : source_type == 21 ? 110 : source_type == 23 ? 136 : 82; @@ -410,6 +515,10 @@ void gguf_dp4a_gate_up_sm70_out(torch::Tensor out, torch::Tensor activation, "Expected aligned original GGUF lattice expert rows"); const int n = gate.size(1), top_k = ids.size(1); const bool quantized = out.scalar_type() == torch::kUInt8; + const bool v41 = source_type == 16 || source_type == 17 || source_type == 29; + TORCH_CHECK(!v41 || (!activated && !quantized && !bank_aware && m > 0 && + m <= 6 && k == 5120 && n == 576 && top_k == 6), + "V4.1 gate/up requires unactivated FP16 M1..6, K5120 N576 top-6"); // lanes_per_row 0: K2560 16-lane schedule with a fully unrolled group loop. TORCH_CHECK(lanes_per_row == 16 || (quantized && (lanes_per_row == 4 || lanes_per_row == 8 || @@ -433,7 +542,13 @@ void gguf_dp4a_gate_up_sm70_out(torch::Tensor out, torch::Tensor activation, else \ dispatch_gate_up(out, activation, ids, gate, up, activated, \ lanes_per_row) - if (source_type == 18) { + if (source_type == 16) { + DISPATCH_LATTICE(16); + } else if (source_type == 17) { + DISPATCH_LATTICE(17); + } else if (source_type == 29) { + DISPATCH_LATTICE(29); + } else if (source_type == 18) { DISPATCH_LATTICE(18); } else if (source_type == 20) { DISPATCH_LATTICE(20); @@ -458,6 +573,125 @@ namespace { // Lanes, true> (IQ3_S: 16 lanes per row; IQ4_NL/IQ4_XS: 4). constexpr int kGroupedChunk = 8; +template +__global__ void v41_gate_reuse(half* output, const Q8_1* activation, + const Index* ids, const uint8_t* gate, + const uint8_t* up, int stride, + const int* enabled) { + if constexpr (Guarded) + if (!*enabled) return; + const Index expert = ids[blockIdx.y]; + // Every CTA of a repeated route exits before allocating/initializing its + // codebook. The first occurrence computes all six tokens selecting it. + for (int previous = 0; previous < blockIdx.y; ++previous) + if (ids[previous] == expert) return; + int members[6] = {}; +#pragma unroll + for (int token = 0; token < 6; ++token) +#pragma unroll + for (int route = 0; route < 6; ++route) + members[token] |= int(ids[token * 6 + route] == expert) << route; + using Dot = LatticeDot; + __shared__ uint32_t book[Dot::kBookWords], masks[16]; + Dot::initialize(book, masks); + const int lane = threadIdx.x % 32, warp = threadIdx.x / 32; + const int row = blockIdx.x * 8 + warp * 2 + lane / 16; + const auto* g = gate + (int64_t(expert) * 576 + row) * stride; + const auto* u = up + (int64_t(expert) * 576 + row) * stride; + float gs[6] = {}, us[6] = {}; + for (int group = lane % 16; group < 160; group += 16) { + const vllm::sm70_gguf::V41IntegerGroup dg(g, group, book, masks); + const vllm::sm70_gguf::V41IntegerGroup du(u, group, book, masks); +#pragma unroll + for (int token = 0; token < 6; ++token) + if (members[token]) { + const auto& x = activation[token * 160 + group]; + gs[token] = __fadd_rn(gs[token], dg.apply(x)); + us[token] = __fadd_rn(us[token], du.apply(x)); + } + } +#pragma unroll + for (int token = 0; token < 6; ++token) { +#pragma unroll + for (int offset = 8; offset; offset >>= 1) { + gs[token] += __shfl_down_sync(0xffffffffU, gs[token], offset, 16); + us[token] += __shfl_down_sync(0xffffffffU, us[token], offset, 16); + } + if (lane % 16 == 0) { +#pragma unroll + for (int route = 0; route < 6; ++route) + if (members[token] & (1 << route)) { + output[(int64_t(token) * 6 + route) * 1152 + row] = + __float2half_rn(gs[token]); + output[(int64_t(token) * 6 + route) * 1152 + 576 + row] = + __float2half_rn(us[token]); + } + } + } +} + +template +void launch_v41_gate_reuse(torch::Tensor out, torch::Tensor activation, + torch::Tensor ids, torch::Tensor gate, + torch::Tensor up, const int* enabled = nullptr) { + if (enabled) + v41_gate_reuse + <<>>( + reinterpret_cast(out.data_ptr()), + reinterpret_cast(activation.data_ptr()), + ids.data_ptr(), gate.data_ptr(), + up.data_ptr(), gate.size(2), enabled); + else + v41_gate_reuse + <<>>( + reinterpret_cast(out.data_ptr()), + reinterpret_cast(activation.data_ptr()), + ids.data_ptr(), gate.data_ptr(), + up.data_ptr(), gate.size(2), nullptr); +} + +template +__global__ void v41_reuse_plan(const Index* ids, int32_t* compact, int* enabled, + int threshold, int slots) { + __shared__ int values[36], counts[2]; + const int i = threadIdx.x; + if (i < 36) { + const auto value = ids[i]; + if (value < 0 || value >= slots) asm volatile("trap;"); + compact[i] = values[i] = int(value); + } + __syncthreads(); + bool first = i < 36; + if (first) + for (int j = 0; j < i; ++j) + if (values[j] == values[i]) first = false; + const unsigned mask = __ballot_sync(0xffffffffU, first); + if (i % 32 == 0) counts[i / 32] = __popc(mask); + __syncthreads(); + if (!i) *enabled = counts[0] + counts[1] <= threshold; +} + +template +void launch_v41_gate_adaptive(torch::Tensor out, torch::Tensor activation, + torch::Tensor ids, torch::Tensor compact, + torch::Tensor mode, torch::Tensor gate, + torch::Tensor up) { + const auto stream = at::cuda::getCurrentCUDAStream(); + v41_reuse_plan<<<1, 64, 0, stream>>>( + ids.data_ptr(), compact.data_ptr(), mode.data_ptr(), + Type == 29 ? 32 : 24, gate.size(0)); + // Both branches are captured. The device decision executes exactly one; + // the fallback retains the caller's original index width and dot order. + gate_up + <<>>( + out.data_ptr(), reinterpret_cast(activation.data_ptr()), + ids.data_ptr(), gate.data_ptr(), + up.data_ptr(), 576, 5120, gate.size(2), 6, nullptr, nullptr, + mode.data_ptr()); + launch_v41_gate_reuse(out, activation, compact, gate, up, + mode.data_ptr()); +} + template __global__ __launch_bounds__(32 * Lanes) void gate_up_grouped( Q8_1* output, const Q8_1* activation, const Index* ids, const uint8_t* gate, @@ -591,6 +825,83 @@ void dispatch_gate_up_grouped(torch::Tensor out, torch::Tensor activation, } } // namespace +static void v41_gate_reuse_out(torch::Tensor out, torch::Tensor activation, + torch::Tensor ids, torch::Tensor gate, + torch::Tensor up, int64_t source_type, + bool adaptive) { + TORCH_CHECK(source_type == 16 || source_type == 17 || source_type == 29, + "V4.1 integer reuse requires IQ1_M/IQ2 original blocks"); + const int bytes = source_type == 16 ? 66 : source_type == 17 ? 74 : 56; + TORCH_CHECK(activation.is_cuda() && activation.is_contiguous() && + activation.scalar_type() == torch::kUInt8 && + activation.dim() == 3 && activation.size(0) == 6 && + activation.size(1) == 160 && + activation.size(2) == sizeof(Q8_1), + "V4.1 integer reuse requires M6 K5120 Q8_1"); + TORCH_CHECK(ids.device() == activation.device() && ids.is_contiguous() && + ids.dim() == 2 && ids.size(0) == 6 && ids.size(1) == 6 && + (ids.scalar_type() == torch::kInt32 || + ids.scalar_type() == torch::kInt64), + "V4.1 integer reuse requires top-6 indices"); + TORCH_CHECK(gate.device() == activation.device() && + gate.scalar_type() == torch::kUInt8 && gate.is_contiguous() && + gate.dim() == 3 && gate.size(0) > 0 && + gate.size(0) <= 65535 && gate.size(1) == 576 && + (gate.size(2) == 20 * bytes || + gate.size(2) == ((20 * bytes + 7) / 8 * 8)) && + up.device() == gate.device() && + up.scalar_type() == gate.scalar_type() && + up.is_contiguous() && up.sizes() == gate.sizes(), + "V4.1 integer reuse requires aligned raw expert carriers"); + TORCH_CHECK(out.device() == gate.device() && out.is_contiguous() && + out.scalar_type() == torch::kFloat16 && out.dim() == 4 && + out.size(0) == 6 && out.size(1) == 6 && out.size(2) == 2 && + out.size(3) == 576, + "V4.1 integer reuse requires FP16 M6 gate/up outputs"); + const c10::cuda::CUDAGuard guard(activation.device()); + require_sm70(); +#define V41_REUSE(TYPE) \ + if (adaptive && ids.scalar_type() == torch::kInt32) \ + launch_v41_gate_adaptive(out, activation, ids, compact, \ + mode, gate, up); \ + else if (adaptive) \ + launch_v41_gate_adaptive(out, activation, ids, compact, \ + mode, gate, up); \ + else if (ids.scalar_type() == torch::kInt32) \ + launch_v41_gate_reuse(out, activation, ids, gate, up); \ + else \ + launch_v41_gate_reuse(out, activation, ids, gate, up) + torch::Tensor compact, mode; + if (adaptive) { + compact = torch::empty({6, 6}, ids.options().dtype(torch::kInt32)); + mode = torch::empty({1}, ids.options().dtype(torch::kInt32)); + } + if (source_type == 16) { + V41_REUSE(16); + } else if (source_type == 17) { + V41_REUSE(17); + } else { + V41_REUSE(29); + } +#undef V41_REUSE + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void gguf_v41_dp4a_gate_reuse_sm70_out(torch::Tensor out, + torch::Tensor activation, + torch::Tensor ids, torch::Tensor gate, + torch::Tensor up, int64_t source_type) { + v41_gate_reuse_out(out, activation, ids, gate, up, source_type, false); +} + +void gguf_v41_dp4a_gate_adaptive_sm70_out(torch::Tensor out, + torch::Tensor activation, + torch::Tensor ids, torch::Tensor gate, + torch::Tensor up, + int64_t source_type) { + v41_gate_reuse_out(out, activation, ids, gate, up, source_type, true); +} + // Routed gate/up with SiLU and Q8_1 output, decoding each routed expert once // per `per_cta` of its routes. Results equal gguf_dp4a_gate_up_sm70_out(..., // activated=True, lanes_per_row=16 for IQ3_S or 4 for IQ4_NL/IQ4_XS) bit for @@ -737,15 +1048,23 @@ void gguf_dp4a_down_unroute_sm70_out(torch::Tensor out, torch::Tensor input, } namespace { -template +template __global__ void raw_down_unroute(half* output, const void* input, const Index* ids, const float* route_weights, const uint8_t* weights, int n, int k, int row_stride, int left_groups, int top_k) { extern __shared__ Q8_1 x[]; - __shared__ float partial[4][32]; + constexpr bool lattice = Type == 16 || Type == 17 || Type == 29; + constexpr int book_words = [] { + if constexpr (lattice) + return LatticeDot::kBookWords; + else + return 1; + }(); + __shared__ uint32_t book[book_words], masks[lattice ? 16 : 1]; + __shared__ float partial[4][32 / Lanes]; const int lane = threadIdx.x % 32, warp = threadIdx.x / 32; - const int groups = k / 32, token = blockIdx.y; + const int groups = lattice ? 18 : k / 32, token = blockIdx.y; if constexpr (Quantized) { const auto* source = static_cast(input) + int64_t{token} * top_k * groups * 9; @@ -758,35 +1077,73 @@ __global__ void raw_down_unroute(half* output, const void* input, vllm::sm70_gguf::quantize_q8_1_warp(x + i, __half2float(value)); } } - __syncthreads(); - const int col = blockIdx.x * 32 + lane; - float total = 0.f; - for (int route = warp; route < top_k; route += 4) { - const int slot = token * top_k + route; - const uint8_t* row = weights + (int64_t{ids[slot]} * n + col) * row_stride; - float dot = 0.f; - for (int group = 0; group < groups; ++group) - dot += vllm::sm70_gguf::OriginalIntegerDot::dot( - row, group + left_groups, x[route * groups + group]); - total += __half2float(__float2half_rn(dot)) * route_weights[slot]; - } - partial[warp][lane] = total; - __syncthreads(); - if (!warp) { - float sum = 0.f; + if constexpr (lattice) + LatticeDot::initialize(book, masks); + else + __syncthreads(); + // Reuse the activation and integer codebook across adjacent output tiles. + // Each tile retains exactly the original K-group and route additions. + for (int tile = 0; tile < Tiles; ++tile) { + const int col = (blockIdx.x * Tiles + tile) * (32 / Lanes) + lane / Lanes; + float total = 0.f; + for (int route = warp; route < top_k; route += 4) { + const int slot = token * top_k + route; + const uint8_t* row = + weights + (int64_t{ids[slot]} * n + col) * row_stride; + float dot = 0.f; + if constexpr (Lanes > 1) { + static_assert(lattice && Quantized); + // Neighboring lanes fetch neighboring K groups of the same packed + // row. Shuffles then restore the original ascending FP32 additions; + // a parallel reduction would change the numerical boundary. + for (int first = 0; first < groups; first += Lanes) { + const int group = first + lane % Lanes; + const float value = + group < groups ? LatticeDot::dot(row, group + left_groups, + x[route * groups + group], + book, masks) + : 0.f; #pragma unroll - for (int i = 0; i < 4; ++i) sum += partial[i][lane]; - output[int64_t{token} * n + col] = __float2half_rn(sum); + for (int i = 0; i < Lanes; ++i) { + const float ordered = + __shfl_sync(0xffffffff, value, (lane / Lanes) * Lanes + i); + if (first + i < groups) dot = __fadd_rn(dot, ordered); + } + } + } else { + for (int group = 0; group < groups; ++group) { + if constexpr (lattice) + dot = __fadd_rn(dot, LatticeDot::dot( + row, group + left_groups, + x[route * groups + group], book, masks)); + else + dot += vllm::sm70_gguf::OriginalIntegerDot::dot( + row, group + left_groups, x[route * groups + group]); + } + } + total += __half2float(__float2half_rn(dot)) * route_weights[slot]; + } + if (lane % Lanes == 0) partial[warp][lane / Lanes] = total; + __syncthreads(); + if (!warp && lane % Lanes == 0) { + float sum = 0.f; +#pragma unroll + for (int i = 0; i < 4; ++i) sum += partial[i][lane / Lanes]; + output[int64_t{token} * n + col] = __float2half_rn(sum); + } + if constexpr (Tiles > 1) __syncthreads(); } } -template +template void launch_raw_down(torch::Tensor out, torch::Tensor input, torch::Tensor ids, torch::Tensor probabilities, torch::Tensor weights, int left) { const int k = input.size(2) * (Quantized ? 32 : 1); - raw_down_unroute - << + <<>>( reinterpret_cast(out.data_ptr()), input.data_ptr(), @@ -798,62 +1155,94 @@ void launch_raw_down(torch::Tensor out, torch::Tensor input, torch::Tensor ids, template void dispatch_raw_down(torch::Tensor out, torch::Tensor input, torch::Tensor ids, torch::Tensor probabilities, - torch::Tensor weights, int left) { - if (input.scalar_type() == torch::kUInt8) - launch_raw_down(out, input, ids, probabilities, weights, - left); - else + torch::Tensor weights, int left, int tiles) { + if (input.scalar_type() == torch::kUInt8) { + if constexpr (Type == 16 || Type == 17 || Type == 29) { + if (tiles == 2) + launch_raw_down(out, input, ids, probabilities, + weights, left); + else if (tiles == 4) + launch_raw_down(out, input, ids, probabilities, + weights, left); + else if (tiles == 8) + launch_raw_down(out, input, ids, probabilities, + weights, left); + else + launch_raw_down(out, input, ids, probabilities, + weights, left); + } else { + launch_raw_down(out, input, ids, probabilities, + weights, left); + } + } else launch_raw_down(out, input, ids, probabilities, weights, left); } } // namespace -void gguf_dp4a_raw_down_unroute_sm70_out(torch::Tensor out, torch::Tensor input, - torch::Tensor ids, - torch::Tensor probabilities, - torch::Tensor weights, - int64_t source_type, int64_t left) { - TORCH_CHECK(source_type == 20 || source_type == 42, - "Original down requires IQ4_NL or Q2_0"); +static void raw_down_out(torch::Tensor out, torch::Tensor input, + torch::Tensor ids, torch::Tensor probabilities, + torch::Tensor weights, int64_t source_type, + int64_t left, int tiles) { + const bool v41 = source_type == 16 || source_type == 17 || source_type == 29; + TORCH_CHECK(source_type == 20 || source_type == 42 || v41, + "Unsupported original integer down source"); const bool quantized = input.scalar_type() == torch::kUInt8; + TORCH_CHECK(tiles == 1 || (v41 && quantized && + (tiles == 2 || tiles == 4 || tiles == 8)), + "Tiled original down requires V4.1 Q8_1 input"); TORCH_CHECK(input.is_cuda() && input.is_contiguous() && (quantized ? input.dim() == 4 && input.size(3) == sizeof(Q8_1) : input.dim() == 3 && input.scalar_type() == torch::kFloat16), "Expected routed FP16 or Q8_1 input"); const int k = input.size(2) * (quantized ? 32 : 1); - const int block = source_type == 20 ? 32 : 64; + const int block = v41 ? 256 : source_type == 20 ? 32 : 64; + const int block_bytes = source_type == 16 ? 66 + : source_type == 17 ? 74 + : source_type == 29 ? 56 + : 18; TORCH_CHECK(input.size(0) > 0 && input.size(0) <= 20 && input.size(1) > 0 && - input.size(1) <= 16 && k > 0 && k <= 256 && k % 32 == 0 && - left >= 0 && left < block && left % 32 == 0, + input.size(1) <= 16 && k > 0 && (v41 ? k == 576 : k <= 256) && + k % 32 == 0 && left >= 0 && left < block && left % 32 == 0, "Unsupported original down geometry"); + TORCH_CHECK(!v41 || ((input.size(0) > 0 && input.size(0) <= 6) && + input.size(1) == 6 && out.dim() == 2 && + out.size(1) == 5120 && left % 64 == 0), + "V4.1 down requires FP16 or Q8_1 M1..6 top-6"); for (const auto& tensor : {out, ids, probabilities, weights}) TORCH_CHECK(tensor.device() == input.device() && tensor.is_contiguous(), "Original down tensors must share a CUDA device"); - TORCH_CHECK(out.scalar_type() == torch::kFloat16 && out.dim() == 2 && - out.size(0) == input.size(0) && out.size(1) > 0 && - out.size(1) % 32 == 0 && ids.dim() == 2 && - ids.size(0) == input.size(0) && - ids.size(1) == input.size(1) && - (ids.scalar_type() == torch::kInt32 || - ids.scalar_type() == torch::kInt64) && - probabilities.scalar_type() == torch::kFloat32 && - probabilities.sizes() == ids.sizes() && - weights.scalar_type() == torch::kUInt8 && - weights.dim() == 3 && weights.size(0) > 0 && - weights.size(0) <= 65535 && weights.size(1) == out.size(1) && - weights.size(2) == ((left + k + block - 1) / block) * 18, - "Invalid original down storage or routing"); + TORCH_CHECK( + out.scalar_type() == torch::kFloat16 && out.dim() == 2 && + out.size(0) == input.size(0) && out.size(1) > 0 && + out.size(1) % 32 == 0 && ids.dim() == 2 && + ids.size(0) == input.size(0) && ids.size(1) == input.size(1) && + (ids.scalar_type() == torch::kInt32 || + ids.scalar_type() == torch::kInt64) && + probabilities.scalar_type() == torch::kFloat32 && + probabilities.sizes() == ids.sizes() && + weights.scalar_type() == torch::kUInt8 && weights.dim() == 3 && + weights.size(0) > 0 && weights.size(0) <= 65535 && + weights.size(1) == out.size(1) && + weights.size(2) == ((left + k + block - 1) / block) * block_bytes, + "Invalid original down storage or routing"); const c10::cuda::CUDAGuard guard(input.device()); require_sm70(); #define DISPATCH_RAW(TYPE) \ if (ids.scalar_type() == torch::kInt32) \ dispatch_raw_down(out, input, ids, probabilities, weights, \ - left); \ + left, tiles); \ else \ dispatch_raw_down(out, input, ids, probabilities, weights, \ - left) - if (source_type == 20) { + left, tiles) + if (source_type == 16) { + DISPATCH_RAW(16); + } else if (source_type == 17) { + DISPATCH_RAW(17); + } else if (source_type == 29) { + DISPATCH_RAW(29); + } else if (source_type == 20) { DISPATCH_RAW(20); } else { DISPATCH_RAW(42); @@ -862,6 +1251,28 @@ void gguf_dp4a_raw_down_unroute_sm70_out(torch::Tensor out, torch::Tensor input, C10_CUDA_KERNEL_LAUNCH_CHECK(); } +void gguf_dp4a_raw_down_unroute_sm70_out(torch::Tensor out, torch::Tensor input, + torch::Tensor ids, + torch::Tensor probabilities, + torch::Tensor weights, + int64_t source_type, int64_t left) { + raw_down_out(out, input, ids, probabilities, weights, source_type, left, 1); +} + +void gguf_v41_dp4a_q8_down_tiled_sm70_out( + torch::Tensor out, torch::Tensor input, torch::Tensor ids, + torch::Tensor probabilities, torch::Tensor weights, int64_t source_type, + int64_t left, int64_t tiles) { + TORCH_CHECK( + (source_type == 16 || source_type == 17 || source_type == 29) && + input.scalar_type() == torch::kUInt8 && + (tiles == 0 || tiles == 1 || tiles == 2 || tiles == 4 || tiles == 8), + "V4.1 tiled down requires original lattice rows and Q8_1 input"); + raw_down_out( + out, input, ids, probabilities, weights, source_type, left, + tiles ? tiles : (source_type == 29 ? (input.size(0) == 1 ? 2 : 4) : 1)); +} + void gguf_dp4a_scalar_lut_gate_up_sm70_out( torch::Tensor out, torch::Tensor activation, torch::Tensor ids, torch::Tensor gate, torch::Tensor up, int64_t source_type, bool activated) { diff --git a/csrc/sm70_turbomind/ops/gguf_prefill_transcode.cuh b/csrc/sm70_turbomind/ops/gguf_prefill_transcode.cuh index 8c0500e1b8..8de5c855ca 100644 --- a/csrc/sm70_turbomind/ops/gguf_prefill_transcode.cuh +++ b/csrc/sm70_turbomind/ops/gguf_prefill_transcode.cuh @@ -34,6 +34,21 @@ struct Format; // Lattice formats: 256-value blocks; U2 carrier, 16-bit packet per 8 values. template <> +struct Format<16> { // IQ2_XXS + static constexpr int kBlock = 256, kBlockBytes = 66, kGroup = 32; + using Stat = uint32_t; +}; +template <> +struct Format<17> { // IQ2_XS + static constexpr int kBlock = 256, kBlockBytes = 74, kGroup = 16; + using Stat = uint32_t; +}; +template <> +struct Format<29> { // IQ1_M + static constexpr int kBlock = 256, kBlockBytes = 56, kGroup = 16; + using Stat = uint16_t; +}; +template <> struct Format<21> { // IQ3_S static constexpr int kBlock = 256, kBlockBytes = 110, kGroup = 32; using Stat = uint64_t; @@ -90,6 +105,77 @@ template using PacketOf = std::conditional_t; +__device__ __forceinline__ uint16_t u16(const uint8_t* p) { + return uint16_t(p[0]) | (uint16_t(p[1]) << 8); +} + +__device__ __forceinline__ uint32_t u32(const uint8_t* p) { + return uint32_t(u16(p)) | (uint32_t(u16(p + 2)) << 16); +} + +__device__ __forceinline__ uint16_t iq2_packet(int index, uint32_t signs) { + signs |= uint32_t(__popc(signs) & 1) << 7; + return uint16_t((index & 255) | (signs << 8)); +} + +template <> +struct Group<16> { + __device__ static uint32_t decode(const uint8_t* row, int k0, + uint16_t (&packets)[4], int* error) { + const uint8_t* b = row + (k0 / 256) * 66; + const int sub = (k0 % 256) / 32; + const uint32_t indices = u32(b + 2 + sub * 8); + const uint32_t meta = u32(b + 6 + sub * 8); +#pragma unroll + for (int q = 0; q < 4; ++q) + packets[q] = + iq2_packet((indices >> (8 * q)) & 255, (meta >> (7 * q)) & 127); + const float scale = f16(b) * (0.5f + float(meta >> 28)) * 0.25f; + return to_f16(scale, error); + } +}; + +template <> +struct Group<17> { + __device__ static uint32_t decode(const uint8_t* row, int k0, + uint16_t (&packets)[2], int* error) { + const uint8_t* b = row + (k0 / 256) * 74; + const int sub = (k0 % 256) / 16; + uint32_t high = 0; +#pragma unroll + for (int q = 0; q < 2; ++q) { + const uint16_t code = u16(b + 2 + (2 * sub + q) * 2); + packets[q] = iq2_packet(code & 255, code >> 9); + high |= uint32_t((code >> 8) & 1) << (2 * q); + } + const int nib = (b[66 + sub / 2] >> (4 * (sub % 2))) & 15; + const float scale = f16(b) * (0.5f + float(nib)) * 0.25f; + return uint32_t(to_f16(scale, error)) | (high << 16); + } +}; + +template <> +struct Group<29> { + __device__ static uint16_t decode(const uint8_t* row, int k0, + uint16_t (&packets)[2], int* error) { + const uint8_t* b = row + (k0 / 256) * 56; + const int sub = (k0 % 256) / 16; + const uint16_t bits = (u16(b + 48) >> 12) | ((u16(b + 50) >> 8) & 0x00f0) | + ((u16(b + 52) >> 4) & 0x0f00) | + (u16(b + 54) & 0xf000); + const int local = (u16(b + 48 + 2 * (sub / 4)) >> (3 * (sub % 4))) & 7; +#pragma unroll + for (int q = 0; q < 2; ++q) { + const int t = sub * 2 + q; + const int high = (b[32 + t / 2] >> (4 * (t % 2))) & 15; + packets[q] = uint16_t(b[t] | ((high & 7) << 8) | ((high >> 3) << 15)); + } + const float scale = + __half2float(__ushort_as_half(bits)) * float(2 * local + 1); + return to_f16(scale, error); + } +}; + template <> struct Group<21> { __device__ static uint64_t decode(const uint8_t* row, int k0, diff --git a/csrc/sm70_turbomind/ops/gguf_prefill_transcode_sm70.cu b/csrc/sm70_turbomind/ops/gguf_prefill_transcode_sm70.cu index fccb05946d..060188bc84 100644 --- a/csrc/sm70_turbomind/ops/gguf_prefill_transcode_sm70.cu +++ b/csrc/sm70_turbomind/ops/gguf_prefill_transcode_sm70.cu @@ -82,6 +82,9 @@ void gguf_prefill_transcode_sm70(torch::Tensor dst_w, torch::Tensor dst_s, ctas, n_out, k, k_offset, src_stride, error); \ break; switch (source_type) { + CASE(16, 256, 66) + CASE(17, 256, 74) + CASE(29, 256, 56) CASE(18, 256, 98) CASE(21, 256, 110) CASE(22, 256, 82) diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index e5b33a64d4..5751c1eba8 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -5,6 +5,24 @@ #include #include "cuda_utils.h" #include "ops.h" + +void gguf_v41_swiglu_q8_sm70_out(torch::Tensor out, torch::Tensor projection, + torch::Tensor probabilities); +void deepseek_v41_swiglu_sm70_out(torch::Tensor out, torch::Tensor gate, + torch::Tensor up); +void gguf_v41_dp4a_gate_reuse_sm70_out(torch::Tensor out, + torch::Tensor activation, + torch::Tensor ids, torch::Tensor gate, + torch::Tensor up, int64_t source_type); +void gguf_v41_dp4a_gate_adaptive_sm70_out(torch::Tensor out, + torch::Tensor activation, + torch::Tensor ids, torch::Tensor gate, + torch::Tensor up, + int64_t source_type); +void gguf_v41_dp4a_q8_down_tiled_sm70_out( + torch::Tensor out, torch::Tensor input, torch::Tensor ids, + torch::Tensor probabilities, torch::Tensor weights, int64_t source_type, + int64_t left, int64_t tiles); #include "core/registration.h" #include #include @@ -390,6 +408,36 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "int left=0) -> ()"); ops.impl("gguf_dp4a_raw_down_unroute_sm70_out", torch::kCUDA, &gguf_dp4a_raw_down_unroute_sm70_out); + // Distinct names let V4.1 admission reject older prebuilt extensions. + ops.def( + "gguf_v41_dp4a_down_sm70_out(Tensor(a!) out, Tensor input, Tensor ids, " + "Tensor probabilities, Tensor weights, int source_type, int left=0) -> " + "()"); + ops.impl("gguf_v41_dp4a_down_sm70_out", torch::kCUDA, + &gguf_dp4a_raw_down_unroute_sm70_out); + ops.def( + "gguf_v41_dp4a_q8_down_sm70_out(Tensor(a!) out, Tensor input, Tensor " + "ids, " + "Tensor probabilities, Tensor weights, int source_type, int left=0) -> " + "()"); + ops.impl("gguf_v41_dp4a_q8_down_sm70_out", torch::kCUDA, + &gguf_dp4a_raw_down_unroute_sm70_out); + ops.def( + "gguf_v41_dp4a_q8_down_tiled_sm70_out(Tensor(a!) out, Tensor input, " + "Tensor ids, Tensor probabilities, Tensor weights, int source_type, " + "int left, int tiles=0) -> ()"); + ops.impl("gguf_v41_dp4a_q8_down_tiled_sm70_out", torch::kCUDA, + &gguf_v41_dp4a_q8_down_tiled_sm70_out); + ops.def( + "gguf_v41_swiglu_q8_sm70_out(Tensor(a!) out, Tensor projection, " + "Tensor probabilities) -> ()"); + ops.impl("gguf_v41_swiglu_q8_sm70_out", torch::kCUDA, + &gguf_v41_swiglu_q8_sm70_out); + ops.def( + "deepseek_v41_swiglu_sm70_out(Tensor(a!) out, Tensor gate, Tensor up) -> " + "()"); + ops.impl("deepseek_v41_swiglu_sm70_out", torch::kCUDA, + &deepseek_v41_swiglu_sm70_out); ops.def( "gguf_dense_restore_canonical_sm70_out(Tensor(a!) weight, Tensor(b!) " "stats, " @@ -438,6 +486,23 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "int lanes_per_row=16, bool bank_aware=False) -> ()"); ops.impl("gguf_dp4a_gate_up_sm70_out", torch::kCUDA, &gguf_dp4a_gate_up_sm70_out); + ops.def( + "gguf_v41_dp4a_gate_up_sm70_out(Tensor(a!) out, Tensor activation, " + "Tensor " + "ids, Tensor gate, Tensor up, int source_type, bool activated, " + "int lanes_per_row=16, bool bank_aware=False) -> ()"); + ops.impl("gguf_v41_dp4a_gate_up_sm70_out", torch::kCUDA, + &gguf_dp4a_gate_up_sm70_out); + ops.def( + "gguf_v41_dp4a_gate_reuse_sm70_out(Tensor(a!) out, Tensor activation, " + "Tensor ids, Tensor gate, Tensor up, int source_type) -> ()"); + ops.impl("gguf_v41_dp4a_gate_reuse_sm70_out", torch::kCUDA, + &gguf_v41_dp4a_gate_reuse_sm70_out); + ops.def( + "gguf_v41_dp4a_gate_adaptive_sm70_out(Tensor(a!) out, Tensor activation, " + "Tensor ids, Tensor gate, Tensor up, int source_type) -> ()"); + ops.impl("gguf_v41_dp4a_gate_adaptive_sm70_out", torch::kCUDA, + &gguf_v41_dp4a_gate_adaptive_sm70_out); ops.def( "gguf_dp4a_gate_up_grouped_sm70_out(Tensor(a!) out, Tensor activation, " "Tensor ids, Tensor gate, Tensor up, int source_type, " diff --git a/docs/design/architecture/runtime_reference.md b/docs/design/architecture/runtime_reference.md index 874754bb6e..8546b191fa 100644 --- a/docs/design/architecture/runtime_reference.md +++ b/docs/design/architecture/runtime_reference.md @@ -38,12 +38,12 @@ identify initialization/hash code; they do not evaluate its semantics. | `vllm.config.gdn_state.GdnStateConfig` | kernel_config.gdn.state | resolve: [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L39); graph_options: [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L52) | [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L23) | | `vllm.config.gdn_state.GdnStateTraceConfig` | observability_config.gdn_state | resolve: [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L99); compile_ignored_aliases: [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L142) | [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L59) | | `vllm.config.kernel.IrOpPriorityConfig` | kernel_config.ir_op_priority | compute_hash: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L84) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L67) | -| `vllm.config.kernel.KernelConfig` | kernel_config | compute_hash: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L1095) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L632) | +| `vllm.config.kernel.KernelConfig` | kernel_config | compute_hash: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L1129) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L662) | | `vllm.config.kernel.Sm70AwqConfig` | kernel_config.sm70_awq | resolve: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L338); capture_inputs: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L216) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L317) | | `vllm.config.kernel.Sm70Fp8Config` | kernel_config.sm70_fp8 | resolve: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L433); capture_inputs: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L216) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L369) | | `vllm.config.kernel.Sm70GgufConfig` | kernel_config.sm70_gguf | | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L498) | | `vllm.config.kernel.Sm70NvFp4Config` | kernel_config.sm70_nvfp4 | resolve: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L277); capture_inputs: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L216) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L236) | -| `vllm.config.kernel.Sm70RingConfig` | kernel_config.sm70_ring | | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L621) | +| `vllm.config.kernel.Sm70RingConfig` | kernel_config.sm70_ring | | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L651) | | `vllm.config.kv_events.KVEventsConfig` | kv_events_config | __post_init__: [vllm/config/kv_events.py](../../../vllm/config/kv_events.py#L50) | [vllm/config/kv_events.py](../../../vllm/config/kv_events.py#L11) | | `vllm.config.kv_transfer.KVTransferConfig` | kv_transfer_config | compute_hash: [vllm/config/kv_transfer.py](../../../vllm/config/kv_transfer.py#L75); __post_init__: [vllm/config/kv_transfer.py](../../../vllm/config/kv_transfer.py#L93) | [vllm/config/kv_transfer.py](../../../vllm/config/kv_transfer.py#L23) | | `vllm.config.legacy_inputs.LegacyInputs` | kernel_config.sm70_awq.legacy, kernel_config.sm70_fp8.legacy, kernel_config.sm70_moe.awq.legacy, kernel_config.sm70_moe.fp8.legacy, kernel_config.sm70_moe.mxfp4.legacy, kernel_config.sm70_moe.nvfp4.legacy, kernel_config.sm70_nvfp4.legacy, scheduler_config.sm70_inputs | value: [vllm/config/legacy_inputs.py](../../../vllm/config/legacy_inputs.py#L39) | [vllm/config/legacy_inputs.py](../../../vllm/config/legacy_inputs.py#L15) | @@ -51,7 +51,7 @@ identify initialization/hash code; they do not evaluate its semantics. | `vllm.config.lora.LoRAConfig` | lora_config | compute_hash: [vllm/config/lora.py](../../../vllm/config/lora.py#L81) | [vllm/config/lora.py](../../../vllm/config/lora.py#L31) | | `vllm.config.mamba.MambaBackendEnum` | mamba_config.backend | | [vllm/config/mamba.py](../../../vllm/config/mamba.py#L25) | | `vllm.config.mamba.MambaConfig` | mamba_config | __post_init__: [vllm/config/mamba.py](../../../vllm/config/mamba.py#L56) | [vllm/config/mamba.py](../../../vllm/config/mamba.py#L33) | -| `vllm.config.model.ModelConfig` | model_config, speculative_config.draft_model_config, speculative_config.target_model_config | compute_hash: [vllm/config/model.py](../../../vllm/config/model.py#L373); __post_init__: [vllm/config/model.py](../../../vllm/config/model.py#L477) | [vllm/config/model.py](../../../vllm/config/model.py#L117) | +| `vllm.config.model.ModelConfig` | model_config, speculative_config.draft_model_config, speculative_config.target_model_config | compute_hash: [vllm/config/model.py](../../../vllm/config/model.py#L376); __post_init__: [vllm/config/model.py](../../../vllm/config/model.py#L480) | [vllm/config/model.py](../../../vllm/config/model.py#L119) | | `vllm.config.multimodal.MultiModalConfig` | model_config.multimodal_config, speculative_config.draft_model_config.multimodal_config, speculative_config.target_model_config.multimodal_config | compute_hash: [vllm/config/multimodal.py](../../../vllm/config/multimodal.py#L292) | [vllm/config/multimodal.py](../../../vllm/config/multimodal.py#L74) | | `vllm.config.observability.ObservabilityConfig` | observability_config | compute_hash: [vllm/config/observability.py](../../../vllm/config/observability.py#L119) | [vllm/config/observability.py](../../../vllm/config/observability.py#L25) | | `vllm.config.offload.OffloadConfig` | offload_config | compute_hash: [vllm/config/offload.py](../../../vllm/config/offload.py#L146) | [vllm/config/offload.py](../../../vllm/config/offload.py#L84) | diff --git a/pyproject.toml b/pyproject.toml index 2b4c7c3f5e..1abd91e3a2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,6 +59,8 @@ include = ["vllm*", "flash_qla*", "flash_attn_v100*"] "vllm/third_party/**" = ["ALL"] "vllm/version.py" = ["F401"] "vllm/_version.py" = ["ALL"] +# DS-V4.1 bring-up scripts (run by hand on the 4x V100 host). +"benchmarks/deepseek_v41/**" = ["E501", "E402", "E741", "B023", "UP031", "SIM115"] [tool.ruff.lint] select = [ diff --git a/tests/benchmarks/test_deepseek_v41_acceptance.py b/tests/benchmarks/test_deepseek_v41_acceptance.py new file mode 100644 index 0000000000..da64288b35 --- /dev/null +++ b/tests/benchmarks/test_deepseek_v41_acceptance.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Request-end evidence must agree with all private state boundaries.""" + +import copy + +import pytest + +from benchmarks.deepseek_v41.acceptance import validate_commits + + +def _records(): + return [ + dict( + rank=rank, + commits=[ + dict( + rows=128, + accepted_inputs=1, + prefill=True, + start=0, + hash_history=[2, 3, 4, 128], + hash_metadata=[0, 128, 0], + c2_positions=[128] * 3, + ), + dict( + rows=6, + accepted_inputs=3, + prefill=False, + start=128, + hash_history=[5, 6, 7, 131], + hash_metadata=[128, 6, 0], + c2_positions=[131] * 3, + ), + ], + ) + for rank in range(4) + ] + + +def test_final_state_matches_emitted_prefix_including_eos_or_length_stop(): + records = _records() + validate_commits(records, expected_end=128 + 4 - 1) + # A self-consistent state can still commit an unaccepted candidate after + # the frontend stopped. Check against the observed emitted-token count. + with pytest.raises(AssertionError): + validate_commits(records, expected_end=130) + + +@pytest.mark.parametrize("failure", ["hash", "c2", "pending", "count", "missing"]) +def test_one_rank_cannot_report_a_different_private_boundary(failure): + records = copy.deepcopy(_records()) + row = records[2]["commits"][-1] + if failure == "hash": + row["hash_history"][3] += 1 + elif failure == "c2": + row["c2_positions"][1] -= 1 + elif failure == "pending": + row["hash_metadata"][2] = 1 + elif failure == "count": + row["accepted_inputs"] = 7 + else: + records[2]["commits"] = [] + with pytest.raises(AssertionError): + validate_commits(records, expected_end=131) diff --git a/tests/benchmarks/test_deepseek_v41_checkpoint_audit.py b/tests/benchmarks/test_deepseek_v41_checkpoint_audit.py new file mode 100644 index 0000000000..645cab5dd3 --- /dev/null +++ b/tests/benchmarks/test_deepseek_v41_checkpoint_audit.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +import math +import struct +from pathlib import Path + +import gguf +import pytest + +from benchmarks.deepseek_v41.checkpoint_audit import ( + Tensor, + packed_bank_bytes, + safetensors_names, + shard_header, + verified_paths, +) + + +def expert(role: str, kind: int, shape: tuple[int, int, int]) -> Tensor: + block, unit = gguf.GGML_QUANT_SIZES[gguf.GGMLQuantizationType(kind)] + return Tensor( + f"blk.7.ffn_{role}_exps.weight", + shape, + kind, + math.prod(shape) // block * unit, + ) + + +@pytest.mark.parametrize( + ("kind", "expected"), [(29, 330301496), (16, 389283906), (17, 436469834)] +) +def test_tp4_down_keeps_three_blocks_and_final_row_guard(kind, expected): + # A 576-column logical slice needs three 256-value blocks, not 2.25. + tensor = expert("down", kind, (2304, 5120, 384)) + result = packed_bank_bytes(tensor, 4) + assert result == [expected] * 4 + assert sum(result) > tensor.nbytes + + +def test_gate_tp_splits_rows_without_expanding_source_blocks(): + tensor = expert("gate", 16, (5120, 2304, 384)) + assert packed_bank_bytes(tensor, 4) == [tensor.nbytes // 4] * 4 + + +@pytest.mark.parametrize( + ("role", "shape"), + [("down", (2305, 5120, 384)), ("gate", (5120, 2303, 384))], +) +def test_incomplete_blocks_or_unequal_tp_slices_are_rejected(role, shape): + with pytest.raises(ValueError): + packed_bank_bytes(expert(role, 16, shape), 4) + + +def test_declared_byte_count_cannot_silently_change_the_memory_budget(): + tensor = Tensor("blk.7.ffn_down_exps.weight", (2304, 5120, 384), 16, 1) + with pytest.raises(ValueError, match="byte count"): + packed_bank_bytes(tensor, 4) + + +@pytest.mark.parametrize("success", ["OK", "成功"]) +def test_existing_checksum_results_accept_both_locales(success): + manifest = ( + f"{'a' * 64} gguf/model.gguf\n{'b' * 64} base/dspark/draft.safetensors\n" + ) + log = ( + f"gguf/model.gguf: {success}\n" + f"base/dspark/draft.safetensors: {success}\nEXIT=0\nEXIT=0\n" + ) + assert verified_paths(manifest, log) == { + "gguf/model.gguf", + "base/dspark/draft.safetensors", + } + + +@pytest.mark.parametrize( + "log", + ["EXIT=0\n", "gguf/model.gguf: OK\nEXIT=1\n", "gguf/model.gguf: FAILED\nEXIT=0\n"], +) +def test_exit_zero_alone_is_not_checksum_success(log): + with pytest.raises(ValueError, match="incomplete or failed"): + verified_paths(f"{'a' * 64} gguf/model.gguf\n", log) + + +def write_safetensors(path: Path, offsets=(8, 12)): + header = json.dumps( + { + "mtp.0.weight": {"dtype": "F32", "shape": [2], "data_offsets": [0, 8]}, + "mtp.1.weight": { + "dtype": "I8", + "shape": [4], + "data_offsets": list(offsets), + }, + } + ).encode() + header += b" " * (-len(header) % 8) + path.write_bytes(struct.pack("