diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index 7fa060a..059817d 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/qwen3_6_scripts/qwen3_5.py @@ -47,6 +47,8 @@ logger = init_logger(__name__) _ENGINEX_PROFILE_ENABLED = os.getenv("ENGINEX_PROFILE_DECODE", "0") == "1" _ENGINEX_PROFILE_EVERY = int(os.getenv("ENGINEX_PROFILE_EVERY", "32")) _ENGINEX_PROFILE_SYNC = os.getenv("ENGINEX_PROFILE_SYNC", "1") != "0" +_ENGINEX_MOE_TINY_IMPL = os.getenv("ENGINEX_MOE_TINY_IMPL", "tokenwise") +_ENGINEX_MOE_TINY_MAX = int(os.getenv("ENGINEX_MOE_TINY_MAX", "4")) _enginex_profile_stats: Dict[str, List[float]] = {} _enginex_profile_steps = 0 _enginex_profile_mode = "unknown" @@ -853,11 +855,35 @@ class Qwen3_5MoeSparseBlock(nn.Module): out = (expert_out * ws.unsqueeze(-1)).sum(0, keepdim=True).to( hidden_states.dtype) # (1, H) - elif T <= 2: - # Fast path: tiny decode batch. With max_num_seqs=2, normal decode - # often has T=2 and should not fall back to the Python per-expert - # loop used for long prefill chunks. - with _enginex_profile("moe.routed_tiny_batch_experts"): + elif T <= _ENGINEX_MOE_TINY_MAX and _ENGINEX_MOE_TINY_IMPL == "tokenwise": + # Fast path: tiny decode batch. Compute each token with the proven + # T==1 large-F.linear path, avoiding the per-expert Python loop. + # This usually beats batched bmm for very small T because the first + # projection becomes T larger GEMMs instead of T*K tiny GEMMs. + with _enginex_profile("moe.routed_tiny_tokenwise_experts"): + H = hidden_states.shape[-1] + pieces = [] + for t in range(T): + eids = topk_ids[t] + ws = topk_weights[t].to(hidden_states.dtype) + w13_sel = w13[eids] # (K, 2*I, H) + w2_sel = w2[eids] # (K, H, I) + + gate_up = F.linear( + hidden_states[t:t + 1], + w13_sel.reshape(-1, H), + ).view(self.top_k, -1) # (K, 2*I) + gate, up = gate_up.chunk(2, dim=-1) + act = F.silu(gate) * up + expert_out = torch.bmm( + w2_sel, act.unsqueeze(-1)).squeeze(-1) # (K, H) + pieces.append((expert_out * ws.unsqueeze(-1)).sum(0)) + + out = torch.stack(pieces, dim=0).to(hidden_states.dtype) + elif T <= _ENGINEX_MOE_TINY_MAX: + # Alternative tiny decode batch implementation. Kept for A/B + # testing via ENGINEX_MOE_TINY_IMPL=bmm. + with _enginex_profile("moe.routed_tiny_bmm_experts"): H = hidden_states.shape[-1] flat_eids = topk_ids.reshape(-1) # (T*K,) flat_ws = topk_weights.reshape(-1).to(hidden_states.dtype) diff --git a/worklogs/decode_analysis_report_2026-07-14.md b/worklogs/decode_analysis_report_2026-07-14.md index c4a863b..e88cb5f 100644 --- a/worklogs/decode_analysis_report_2026-07-14.md +++ b/worklogs/decode_analysis_report_2026-07-14.md @@ -667,3 +667,67 @@ decode 聚合吞吐提升约 **31%**。TTFT 变差可能来自冷缓存/加载 4. **TP all-reduce 合并** - 当前 MoE 每层 routed+shared 后做一次 all-reduce。 - 如果后续能将若干小通信或 residual 路径合并，可能进一步改善 decode 抖动，但优先级低于 MoE/linear_attention 计算本体。 + +## 2026-07-14：MoE decode fast path v2（tokenwise） + +### 背景 + +上一轮 `T <= 2` tiny-batch fast path 使用 batched `torch.bmm`，将 `T * top_k` 个 expert 选择展平后批量计算。它已经将 c2 aggregate Output TPS 从 8.07 提升到 10.60。 + +继续分析后发现，对于极小 batch（`T=2`）： + +- 第一段 gate/up projection 用 `torch.bmm` 实际是 `T * top_k` 个 `1 x H` 小矩阵乘。 +- 原 `T == 1` 路径用的是 `F.linear(hidden, w13_sel.reshape(-1, H))`，会形成一个更大的 GEMM，通常更适合 GPU。 + +因此新增 v2：`tokenwise` tiny-batch 实现。 + +### 实现 + +新增环境变量： + +- `ENGINEX_MOE_TINY_IMPL=tokenwise|bmm` +- `ENGINEX_MOE_TINY_MAX=4` + +默认： + +```text +ENGINEX_MOE_TINY_IMPL=tokenwise +ENGINEX_MOE_TINY_MAX=4 +``` + +核心逻辑： + +- 对 `T <= 4` 的小批量 decode，逐 token 走已经验证过的 `T == 1` 大 `F.linear` 路径。 +- 每个 token 内仍然把 top-k expert 的 `w13` 拼成一个大权重，减少第一段 projection 的小 GEMM。 +- 保留上一轮 batched bmm 实现，可通过 `ENGINEX_MOE_TINY_IMPL=bmm` 回退做 A/B 测试。 + +代码位置： + +- `qwen3_6_scripts/qwen3_5.py` +- `Qwen3_5MoeSparseBlock._pure_pytorch_experts` + +### 结果 + +同口径 `short c2, max_tokens=128, requests=4`： + +| 版本 | 成功率 | TTFT P90 | per-request Output TPS P10 | Aggregate Output TPS | +| --- | ---: | ---: | ---: | ---: | +| custom all-reduce on，优化前 | 100% | 1.46s | 4.22 | 8.07 | +| tiny-batch bmm | 100% | 4.42s | 5.95 | 10.60 | +| tiny-batch tokenwise v2 | 100% | 4.49s | 6.43 | 11.28 | + +相对上一版 bmm，aggregate Output TPS 又提升约 **6.4%**；相对优化前提升约 **39.8%**。 + +本地归档： + +- `worklogs/remote_results/2026-07-14-code-profile/decode_moe_tokenwise_short_c2_t128_r4.json` +- `worklogs/remote_results/2026-07-14-code-profile/server_moe_tokenwise_eager_custom_ar_seq2_b8192.log` + +### 结论 + +这个结果说明：对当前 BI-V100 + PyTorch fallback MoE 路径，**将小 GEMM 尽量合并为较大的 per-token GEMM** 比把所有 token/expert 都塞进 batched bmm 更好。瓶颈确实集中在 decode MoE routed expert 的小矩阵计算与调度开销。 + +下一步可继续沿两个方向走： + +1. 测试 `max_num_seqs=4`，利用 `T <= 4` tokenwise fast path，看 aggregate Output TPS 是否继续上涨。 +2. 继续优化 shared expert / linear_attention，因为 routed expert 已经明显改善，decode 下一个大头会逐渐转向 `linear_attention` 和 shared expert。 diff --git a/worklogs/remote_results/2026-07-14-code-profile/decode_moe_tokenwise_short_c2_t128_r4.json b/worklogs/remote_results/2026-07-14-code-profile/decode_moe_tokenwise_short_c2_t128_r4.json new file mode 100644 index 0000000..b4cfacc --- /dev/null +++ b/worklogs/remote_results/2026-07-14-code-profile/decode_moe_tokenwise_short_c2_t128_r4.json @@ -0,0 +1,75 @@ +{ + "created_at": "2026-07-14T13:34:44", + "label": "moe_tokenwise_short_c2_t128_r4", + "url": "http://127.0.0.1:1111", + "model": "llm", + "prompt_mode": "short", + "with_tools": false, + "tool_count": 0, + "concurrency": 2, + "requests": 4, + "max_tokens": 128, + "wall_sec": 45.37405508942902, + "success_rate": 1.0, + "ttft_p50_sec": 2.9874187149107456, + "ttft_p90_sec": 4.485993221774697, + "output_tps_p10_per_request": 6.428466023017078, + "output_tps_p50_per_request": 6.498744495352444, + "aggregate_output_tps": 11.283981539469739, + "prompt_tokens": 156, + "cached_tokens": 64, + "completion_tokens": 512, + "reasoning_tokens": 512, + "chars": 1841, + "monitor": null, + "results": [ + { + "ok": true, + "elapsed_sec": 23.971454864367843, + "ttft_sec": 4.486160388216376, + "completion_tokens": 128, + "prompt_tokens": 39, + "cached_tokens": 0, + "reasoning_tokens": 128, + "output_tps": 6.569056482911376, + "chars": 449, + "error": null + }, + { + "ok": true, + "elapsed_sec": 23.97104039043188, + "ttft_sec": 4.485603166744113, + "completion_tokens": 128, + "prompt_tokens": 39, + "cached_tokens": 0, + "reasoning_tokens": 128, + "output_tps": 6.569008358939714, + "chars": 449, + "error": null + }, + { + "ok": true, + "elapsed_sec": 21.4003098718822, + "ttft_sec": 1.4888528920710087, + "completion_tokens": 128, + "prompt_tokens": 39, + "cached_tokens": 32, + "reasoning_tokens": 128, + "output_tps": 6.428459762125039, + "chars": 457, + "error": null + }, + { + "ok": true, + "elapsed_sec": 21.40062660165131, + "ttft_sec": 1.4892342630773783, + "completion_tokens": 128, + "prompt_tokens": 39, + "cached_tokens": 32, + "reasoning_tokens": 128, + "output_tps": 6.428480631765174, + "chars": 486, + "error": null + } + ] +} \ No newline at end of file diff --git a/worklogs/remote_results/2026-07-14-code-profile/server_moe_tokenwise_eager_custom_ar_seq2_b8192.log b/worklogs/remote_results/2026-07-14-code-profile/server_moe_tokenwise_eager_custom_ar_seq2_b8192.log new file mode 100644 index 0000000..740e9a5 --- /dev/null +++ b/worklogs/remote_results/2026-07-14-code-profile/server_moe_tokenwise_eager_custom_ar_seq2_b8192.log @@ -0,0 +1,301 @@ +/usr/local/corex/lib/python3/dist-packages/torch/cuda/__init__.py:51: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you. + import pynvml # type: ignore[import] +INFO 07-14 13:29:12 importing.py:10] Triton not installed; certain GPU-related functions will not be available. +2026-07-14 13:29:13.750824: I tensorflow/core/util/port.cc:110] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`. +2026-07-14 13:29:13.802713: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations. +To enable the following instructions: SSE3 SSE4.1 SSE4.2 AVX AVX2 AVX512F AVX512_VNNI AVX512_BF16 AVX_VNNI AMX_TILE AMX_INT8 AMX_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags. +WARNING:tensorflow:Deprecation warnings have been disabled. Set TF_ENABLE_DEPRECATION_WARNINGS=1 to re-enable them. +INFO 07-14 13:29:19 api_server.py:530] vLLM API server version 0.6.3 +INFO 07-14 13:29:19 api_server.py:531] args: Namespace(host='0.0.0.0', port=1111, uvicorn_log_level='info', allow_credentials=False, allowed_origins=['*'], allowed_methods=['*'], allowed_headers=['*'], api_key=None, lora_modules=None, prompt_adapters=None, chat_template=None, response_role='assistant', ssl_keyfile=None, ssl_certfile=None, ssl_ca_certs=None, ssl_cert_reqs=0, root_path=None, middleware=[], return_tokens_as_token_ids=False, disable_frontend_multiprocessing=True, enable_auto_tool_choice=True, tool_call_parser='qwen3_coder', tool_parser_plugin='', reasoning_parser='qwen3', model='/root/public-storage/models/Qwen/Qwen3.6-35B-A3B', tokenizer=None, skip_tokenizer_init=False, revision=None, code_revision=None, tokenizer_revision=None, tokenizer_mode='auto', trust_remote_code=True, download_dir=None, load_format='auto', config_format='auto', dtype='auto', kv_cache_dtype='auto', quantization_param_path=None, max_model_len=100000, guided_decoding_backend='outlines', distributed_executor_backend=None, worker_use_ray=False, pipeline_parallel_size=1, tensor_parallel_size=4, max_parallel_loading_workers=None, ray_workers_use_nsight=False, block_size=16, enable_prefix_caching=True, disable_sliding_window=False, use_v2_block_manager=True, num_lookahead_slots=0, seed=0, swap_space=4, cpu_offload_gb=0, gpu_memory_utilization=0.95, num_gpu_blocks_override=None, max_num_batched_tokens=8192, max_num_seqs=2, max_logprobs=20, disable_log_stats=False, quantization=None, rope_scaling=None, rope_theta=None, enforce_eager=True, max_context_len_to_capture=None, max_seq_len_to_capture=32768, disable_custom_all_reduce=False, tokenizer_pool_size=0, tokenizer_pool_type='ray', tokenizer_pool_extra_config=None, limit_mm_per_prompt=None, mm_processor_kwargs=None, enable_lora=False, max_loras=1, max_lora_rank=16, lora_extra_vocab_size=256, lora_dtype='auto', long_lora_scaling_factors=None, max_cpu_loras=None, fully_sharded_loras=False, enable_prompt_adapter=False, max_prompt_adapters=1, max_prompt_adapter_token=0, device='auto', num_scheduler_steps=1, multi_step_stream_outputs=True, scheduler_delay_factor=0.0, enable_chunked_prefill=True, speculative_model=None, speculative_model_quantization=None, num_speculative_tokens=None, speculative_disable_mqa_scorer=False, speculative_draft_tensor_parallel_size=None, speculative_max_model_len=None, speculative_disable_by_batch_size=None, ngram_prompt_lookup_max=None, ngram_prompt_lookup_min=None, spec_decoding_acceptance_method='rejection_sampler', typical_acceptance_sampler_posterior_threshold=None, typical_acceptance_sampler_posterior_alpha=None, disable_logprobs_during_spec_decoding=None, model_loader_extra_config=None, ignore_patterns=[], preemption_mode=None, served_model_name=['llm'], qlora_adapter_name_or_path=None, otlp_traces_endpoint=None, collect_detailed_traces=None, disable_async_output_proc=False, override_neuron_config=None, scheduling_policy='fcfs', disable_log_requests=True, max_log_len=None, disable_fastapi_docs=False) +INFO 07-14 13:29:19 config.py:1670] Downcasting torch.float32 to torch.float16. +INFO 07-14 13:29:30 config.py:887] Defaulting to use mp for distributed inference +INFO 07-14 13:29:30 config.py:1005] Chunked prefill is enabled with max_num_batched_tokens=8192. +WARNING 07-14 13:29:30 config.py:380] To see benefits of async output processing, enable CUDA graph. Since, enforce-eager is enabled, async output processor cannot be used +INFO 07-14 13:29:30 llm_engine.py:237] Initializing an LLM engine (v0.6.3) with config: model='/root/public-storage/models/Qwen/Qwen3.6-35B-A3B', speculative_config=None, tokenizer='/root/public-storage/models/Qwen/Qwen3.6-35B-A3B', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, override_neuron_config=None, rope_scaling=None, rope_theta=None, tokenizer_revision=None, trust_remote_code=True, dtype=torch.float16, max_seq_len=100000, download_dir=None, load_format=LoadFormat.AUTO, tensor_parallel_size=4, pipeline_parallel_size=1, disable_custom_all_reduce=False, quantization=None, enforce_eager=True, kv_cache_dtype=auto, quantization_param_path=None, device_config=cuda, decoding_config=DecodingConfig(guided_decoding_backend='outlines'), observability_config=ObservabilityConfig(otlp_traces_endpoint=None, collect_model_forward_time=False, collect_model_execute_time=False), seed=0, served_model_name=llm, use_v2_block_manager=True, num_scheduler_steps=1, chunked_prefill_enabled=True multi_step_stream_outputs=True, enable_prefix_caching=True, use_async_output_proc=False, use_cached_outputs=False, mm_processor_kwargs=None) +WARNING 07-14 13:29:30 multiproc_gpu_executor.py:53] Reducing Torch parallelism from 64 threads to 1 to avoid unnecessary CPU contention. Set OMP_NUM_THREADS in the external environment to tune this value as needed. +INFO 07-14 13:29:30 custom_cache_manager.py:17] Setting Triton cache manager to: vllm.triton_utils.custom_cache_manager:CustomCacheManager +INFO 07-14 13:29:30 selector.py:266] Cannot use FlashAttention-2 backend because the vllm.vllm_flash_attn package is not found. Make sure that vllm_flash_attn was built and installed (on by default). +INFO 07-14 13:29:30 selector.py:115] Using XFormers backend. +/usr/local/corex/lib/python3/dist-packages/torch/cuda/__init__.py:51: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you. + import pynvml # type: ignore[import] +/usr/local/corex/lib/python3/dist-packages/torch/cuda/__init__.py:51: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you. + import pynvml # type: ignore[import] +/usr/local/corex/lib/python3/dist-packages/torch/cuda/__init__.py:51: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you. + import pynvml # type: ignore[import] +INFO 07-14 13:29:32 importing.py:10] Triton not installed; certain GPU-related functions will not be available. +INFO 07-14 13:29:32 importing.py:10] Triton not installed; certain GPU-related functions will not be available. +INFO 07-14 13:29:32 importing.py:10] Triton not installed; certain GPU-related functions will not be available. +WARNING:tensorflow:Deprecation warnings have been disabled. Set TF_ENABLE_DEPRECATION_WARNINGS=1 to re-enable them. +WARNING:tensorflow:Deprecation warnings have been disabled. Set TF_ENABLE_DEPRECATION_WARNINGS=1 to re-enable them. +WARNING:tensorflow:Deprecation warnings have been disabled. Set TF_ENABLE_DEPRECATION_WARNINGS=1 to re-enable them. +(VllmWorkerProcess pid=15585) INFO 07-14 13:29:39 selector.py:266] Cannot use FlashAttention-2 backend because the vllm.vllm_flash_attn package is not found. Make sure that vllm_flash_attn was built and installed (on by default). +(VllmWorkerProcess pid=15585) INFO 07-14 13:29:39 selector.py:115] Using XFormers backend. +(VllmWorkerProcess pid=15585) INFO 07-14 13:29:39 multiproc_worker_utils.py:216] Worker ready; awaiting tasks +(VllmWorkerProcess pid=15584) INFO 07-14 13:29:40 selector.py:266] Cannot use FlashAttention-2 backend because the vllm.vllm_flash_attn package is not found. Make sure that vllm_flash_attn was built and installed (on by default). +(VllmWorkerProcess pid=15584) INFO 07-14 13:29:40 selector.py:115] Using XFormers backend. +(VllmWorkerProcess pid=15584) INFO 07-14 13:29:40 multiproc_worker_utils.py:216] Worker ready; awaiting tasks +(VllmWorkerProcess pid=15583) INFO 07-14 13:29:40 selector.py:266] Cannot use FlashAttention-2 backend because the vllm.vllm_flash_attn package is not found. Make sure that vllm_flash_attn was built and installed (on by default). +(VllmWorkerProcess pid=15583) INFO 07-14 13:29:40 selector.py:115] Using XFormers backend. +(VllmWorkerProcess pid=15583) INFO 07-14 13:29:40 multiproc_worker_utils.py:216] Worker ready; awaiting tasks +INFO 07-14 13:29:40 shm_broadcast.py:242] vLLM message queue communication handle: Handle(connect_ip='127.0.0.1', local_reader_ranks=[1, 2, 3], buffer=, local_subscribe_port=42675, remote_subscribe_port=None) +INFO 07-14 13:29:40 model_runner.py:1112] Starting to load model /root/public-storage/models/Qwen/Qwen3.6-35B-A3B... +(VllmWorkerProcess pid=15583) INFO 07-14 13:29:40 model_runner.py:1112] Starting to load model /root/public-storage/models/Qwen/Qwen3.6-35B-A3B... +(VllmWorkerProcess pid=15584) INFO 07-14 13:29:40 model_runner.py:1112] Starting to load model /root/public-storage/models/Qwen/Qwen3.6-35B-A3B... +(VllmWorkerProcess pid=15585) INFO 07-14 13:29:40 model_runner.py:1112] Starting to load model /root/public-storage/models/Qwen/Qwen3.6-35B-A3B... +INFO 07-14 13:29:40 selector.py:266] Cannot use FlashAttention-2 backend because the vllm.vllm_flash_attn package is not found. Make sure that vllm_flash_attn was built and installed (on by default). +INFO 07-14 13:29:40 selector.py:115] Using XFormers backend. +(VllmWorkerProcess pid=15584) INFO 07-14 13:29:40 selector.py:266] Cannot use FlashAttention-2 backend because the vllm.vllm_flash_attn package is not found. Make sure that vllm_flash_attn was built and installed (on by default). +(VllmWorkerProcess pid=15584) INFO 07-14 13:29:40 selector.py:115] Using XFormers backend. +(VllmWorkerProcess pid=15585) INFO 07-14 13:29:40 selector.py:266] Cannot use FlashAttention-2 backend because the vllm.vllm_flash_attn package is not found. Make sure that vllm_flash_attn was built and installed (on by default). +(VllmWorkerProcess pid=15585) INFO 07-14 13:29:40 selector.py:115] Using XFormers backend. +(VllmWorkerProcess pid=15583) INFO 07-14 13:29:40 selector.py:266] Cannot use FlashAttention-2 backend because the vllm.vllm_flash_attn package is not found. Make sure that vllm_flash_attn was built and installed (on by default). +(VllmWorkerProcess pid=15583) INFO 07-14 13:29:40 selector.py:115] Using XFormers backend. + Loading safetensors checkpoint shards: 0% Completed | 0/26 [00:00