diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000000..40e0ac5d5c --- /dev/null +++ b/.dockerignore @@ -0,0 +1,46 @@ +.git +.git/** +.github +.idea +.vscode +.pytest_cache +.mypy_cache +.ruff_cache +**/__pycache__ +**/*.pyc +**/*.pyo + +# Local environments, credentials, and host configuration. +.env +.env.* +.aws +.sco +.netrc +*.key +*.pem +id_rsa* +venv +.venv + +# Models, caches, traces, and experiment output never belong in an image. +models +checkpoints +cache* +.cache +traces +trace-* +*.log +*.jsonl +*.pt +*.pth +*.safetensors + +# Runtime images only need the package and release Docker inputs. +assets +demos +docs +skills +test +unit_tests +tools +format_out diff --git a/GLM53_H100_DEPLOY.md b/GLM53_H100_DEPLOY.md new file mode 100644 index 0000000000..f4c9f089f6 --- /dev/null +++ b/GLM53_H100_DEPLOY.md @@ -0,0 +1,91 @@ +# GLM-5.3-Flash H100/H200 部署与验证 + +该配置使用 TP8、BF16 KV cache、32K batch token、1.12M token cache 和 +CUDA Graph 128，面向 16K 输入、256 输出的高并发文本服务。 + +## 镜像 + +私有仓库标签： + +```text +registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm:v1.5.0-glm53-16k256-kpool-996cef93 +``` + +不可变镜像： + +```text +registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm@sha256:0260e9884e46de899f4b845aa3796d0603b7b6ee7593c1fded35b7cee4462169 +``` + +本机镜像为 `lightllm-glm53:16k256-kpool`，镜像 ID 为 +`sha256:8fbb91ee50dde5af4289d9e3cc75dddf93a8dd6fc2a201222a28612a5679cf0c`。 +同一不可变镜像已拉取到 H100 节点。 + +## H100 部署命令 + +镜像已内置服务参数；启动时显式开启 K-pool decode 快路径： + +```bash +IMAGE="registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm@sha256:0260e9884e46de899f4b845aa3796d0603b7b6ee7593c1fded35b7cee4462169" + +sudo docker pull "$IMAGE" +sudo docker run -d \ + --name glm53-lightllm-16k256-kpool \ + --restart unless-stopped \ + --network host \ + --ipc host \ + --gpus all \ + -e LIGHTLLM_ENABLE_KPOOL_DECODE_FASTPATH=1 \ + -e PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \ + --ulimit memlock=-1 \ + --ulimit nofile=1048576:1048576 \ + -v /home/devsft/models/GLM-5.3-Flash:/model:ro \ + -v /home/devsft/cache-glm53-lightllm:/root/.cache \ + -v /home/devsft/cache-glm53-triton:/root/.triton \ + -v /home/devsft/cache-glm53-deep-gemm:/root/.deep_gemm \ + "$IMAGE" +``` + +接口为 `http://127.0.0.1:8002/v1`。该快路径要求请求 prompt 长度按 +16 token 对齐，并使用镜像内置的关闭 prompt cache、关闭 MTP 配置；其他负载可移除 +`LIGHTLLM_ENABLE_KPOOL_DECODE_FASTPATH`，使用保守路径。 + +本机 H200 可直接运行： + +```bash +LIGHTLLM_GLM53_MODEL_DIR=/nvme/sufubao/models/GLM-5.3-Flash \ +LIGHTLLM_GLM53_CACHE_DIR=/nvme/sufubao/m39-home/cache/glm53-lightllm-h200 \ +LIGHTLLM_GLM53_TRITON_CACHE_DIR=/nvme/sufubao/m39-home/cache/glm53-triton-h200 \ +LIGHTLLM_GLM53_DEEP_GEMM_CACHE_DIR=/nvme/sufubao/m39-home/cache/glm53-deep-gemm-h200 \ +tools/run_glm53_h100_container.sh +``` + +## 16K/256 性能 + +LightLLM 与 vLLM 在同一台空闲的 8×H100 80GB 节点顺序测试，使用同一 FP8 +checkpoint、TP8、BF16 KV cache、32K batch token、关闭 prompt/prefix cache、 +seed 42、temperature 0、无限请求速率和一次不计入结果的 warmup。每档请求数等于 +并发数，每个请求输入 16,384 token、输出 256 token。 + +| 并发 | LightLLM 总吞吐 tok/s | vLLM 总吞吐 tok/s | LightLLM / vLLM | +| ---: | ---: | ---: | ---: | +| 1 | 5,373.86 | 6,542.89 | 82.13% | +| 8 | 16,444.11 | 17,635.33 | 93.25% | +| 16 | 20,669.61 | 25,246.99 | 81.87% | +| 64 | 25,614.05 | 31,485.12 | 81.35% | +| 128 | 25,588.22 | 27,988.09 | 91.43% | +| 256 | 25,623.93 | 31,206.94 | 82.11% | + +结果文件位于 +`/nvme/sufubao/m39-home/results/glm53_h100_16k_256_kpool_decode_ab/` 和 +`/nvme/sufubao/m39-home/results/glm53_h100_16k_256_vllm_opt/`；所有启动、性能与 +精度实验均由 `exp` 归档到 `~/experiments/runs/`。 + +## 精度 + +| 验证 | 结果 | +| --- | --- | +| GSM8K 固定前 100 题、5-shot、greedy | **99/100**；100/100 完成；无截断 | +| 精确 16,384-token needle | **PASS**；API 与本地 tokenizer 均计数 16,384；找回 `ZEBRA-4821` | + +相关实验记录为 `260901-020939` 和 `260901-021013`。 diff --git a/docker/Dockerfile.glm53-h100 b/docker/Dockerfile.glm53-h100 new file mode 100644 index 0000000000..902897ed3b --- /dev/null +++ b/docker/Dockerfile.glm53-h100 @@ -0,0 +1,188 @@ +# This is the exact SGLang GLM-5.3 runtime used for the H200 validation. +ARG BASE_IMAGE=lmsysorg/sglang@sha256:92afb4c878eef9cbb17ca9a2c1d15d5cda58585f90bbf5915a79f0f6284aad10 +FROM ${BASE_IMAGE} AS prepared + +ARG PIP_INDEX_URL + +WORKDIR /opt/lightllm + +COPY docker/requirements-glm53-runtime.txt /tmp/requirements-glm53-runtime.txt +RUN --mount=type=cache,id=glm53-pip,target=/root/.cache/pip \ + python -m pip install --no-deps --require-hashes \ + -r /tmp/requirements-glm53-runtime.txt && \ + rm /tmp/requirements-glm53-runtime.txt + +COPY setup.py LICENSE README.md ./ +COPY lightllm ./lightllm +COPY docker/glm53-h100.openvex.json /usr/share/doc/lightllm/glm53-h100.openvex.json +RUN --mount=type=cache,id=glm53-pip,target=/root/.cache/pip \ + python -m pip install --no-deps . && \ + rm -rf \ + /etc/ssh/ssh_host_*_key* \ + /opt/nvidia/nsight-compute/2025.3.1/host/target-linux-x64/plugins/efa_metrics \ + /opt/nvidia/nsight-systems-cli/2026.4.1/target-linux-x64/plugins/efa_metrics \ + /root/.cargo/registry \ + /sgl-workspace/sglang/python/sglang/multimodal_gen \ + /sgl-workspace/sglang/python/sglang/srt/disaggregation \ + /sgl-workspace/sglang/scripts/playground/replay_request_dump.py \ + /opt/sglang/bin/pip \ + /opt/sglang/bin/pip3 \ + /opt/sglang/bin/pip3.12 \ + /opt/sglang/lib/python3.12/site-packages/pip \ + /opt/sglang/lib/python3.12/site-packages/pip-*.dist-info \ + /opt/sglang/lib/python3.12/site-packages/setuptools \ + /opt/sglang/lib/python3.12/site-packages/setuptools-*.dist-info \ + /opt/sglang/lib/python3.12/site-packages/wheel \ + /opt/sglang/lib/python3.12/site-packages/wheel-*.dist-info && \ + python -c "import frozendict, hypercorn, importlib.metadata as metadata, importlib.util, lightllm.server.api_start, msgpack, redis, rpyc, ujson; from sglang.srt.layers.moe.moe_runner.triton_utils import fused_moe; assert metadata.version('hypercorn') == '0.18.0'; assert msgpack.__version__ == '1.2.2'; assert redis.__version__ == '7.3.0'; assert rpyc.__version__ == '6.0.2'; assert importlib.util.find_spec('pip') is None; assert importlib.util.find_spec('setuptools') is None; print('Hardened LightLLM GLM-5.3 runtime import OK')" && \ + test ! -e /etc/ssh/ssh_host_rsa_key && \ + test ! -e /sgl-workspace/sglang/python/sglang/multimodal_gen && \ + test ! -e /sgl-workspace/sglang/python/sglang/srt/disaggregation + +# Flatten the prepared root filesystem so deleted base-image host keys and +# developer-only vulnerable components are absent from published layers, +# rather than merely hidden by an OCI whiteout. Size-bound the clean content +# layers because the private registry's cross-region gateway cannot reliably +# accept one multi-gigabyte upload stream; the merged filesystem is unchanged. +FROM prepared AS partitioned + +COPY docker/split_glm53_rootfs.py /tmp/split_glm53_rootfs.py +RUN /usr/bin/python3 /tmp/split_glm53_rootfs.py + +FROM prepared AS prepared_remainder + +RUN rm -rf \ + /opt \ + /root \ + /sgl-workspace \ + /usr/include \ + /usr/lib/x86_64-linux-gnu \ + /usr/libexec \ + /usr/local \ + /usr/share + +FROM scratch AS runtime + +ARG OCI_CREATED +ARG OCI_REVISION +ARG OCI_SOURCE=https://github.com/sufubao/LightLLM +ARG OCI_VERSION +ARG BASE_NAME=lmsysorg/sglang:glm-5.3-flash +ARG BASE_DIGEST=sha256:92afb4c878eef9cbb17ca9a2c1d15d5cda58585f90bbf5915a79f0f6284aad10 + +COPY --from=prepared_remainder / / +COPY --from=partitioned /__image_layers/00/ / +COPY --from=partitioned /__image_layers/01/ / +COPY --from=partitioned /__image_layers/02/ / +COPY --from=partitioned /__image_layers/03/ / +COPY --from=partitioned /__image_layers/04/ / +COPY --from=partitioned /__image_layers/05/ / +COPY --from=partitioned /__image_layers/06/ / +COPY --from=partitioned /__image_layers/07/ / +COPY --from=partitioned /__image_layers/08/ / +COPY --from=partitioned /__image_layers/09/ / +COPY --from=partitioned /__image_layers/10/ / +COPY --from=partitioned /__image_layers/11/ / +COPY --from=partitioned /__image_layers/12/ / +COPY --from=partitioned /__image_layers/13/ / +COPY --from=partitioned /__image_layers/14/ / +COPY --from=partitioned /__image_layers/15/ / +COPY --from=partitioned /__image_layers/16/ / +COPY --from=partitioned /__image_layers/17/ / +COPY --from=partitioned /__image_layers/18/ / +COPY --from=partitioned /__image_layers/19/ / +COPY --from=partitioned /__image_layers/20/ / +COPY --from=partitioned /__image_layers/21/ / +COPY --from=partitioned /__image_layers/22/ / +COPY --from=partitioned /__image_layers/23/ / +COPY --from=partitioned /__image_layers/24/ / +COPY --from=partitioned /__image_layers/25/ / +COPY --from=partitioned /__image_layers/26/ / +COPY --from=partitioned /__image_layers/27/ / +COPY --from=partitioned /__image_layers/28/ / +COPY --from=partitioned /__image_layers/29/ / +COPY --from=partitioned /__image_layers/30/ / +COPY --from=partitioned /__image_layers/31/ / + +LABEL org.opencontainers.image.created="${OCI_CREATED}" \ + org.opencontainers.image.revision="${OCI_REVISION}" \ + org.opencontainers.image.source="${OCI_SOURCE}" \ + org.opencontainers.image.version="${OCI_VERSION}" \ + org.opencontainers.image.title="LightLLM GLM-5.3-Flash H100/H200 TP8 optimized text" \ + org.opencontainers.image.base.name="${BASE_NAME}" \ + org.opencontainers.image.base.digest="${BASE_DIGEST}" \ + ai.lightllm.model="GLM-5.3-Flash" \ + ai.lightllm.accelerator="NVIDIA H100/H200" \ + ai.lightllm.tensor-parallel-size="8" \ + ai.lightllm.profile="text-16k256-tp8-kpool-c256-no-prompt-cache" \ + ai.lightllm.security-profile="flattened-no-sglang-server-components" + +ENV PATH=/opt/sglang/bin:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin \ + LD_LIBRARY_PATH=/usr/local/nvidia/lib:/usr/local/nvidia/lib64:/usr/local/cuda/lib64 \ + LIBRARY_PATH=/usr/local/cuda/lib64/stubs \ + CUDA_HOME=/usr/local/cuda \ + CUDA_VERSION=13.0.3 \ + NVARCH=x86_64 \ + NVIDIA_REQUIRE_CUDA="cuda>=13.0" \ + NVIDIA_VISIBLE_DEVICES=all \ + NVIDIA_DRIVER_CAPABILITIES=compute,utility \ + LANG=en_US.UTF-8 \ + LANGUAGE=en_US:en \ + LC_ALL=en_US.UTF-8 \ + PYTHONUNBUFFERED=1 \ + PYTHONDONTWRITEBYTECODE=1 \ + LOADWORKER=18 \ + LIGHTLLM_USE_SGLANG_TRITON_MOE=1 \ + LIGHTLLM_SGLANG_FUSED_MOE_SUM=1 \ + LIGHTLLM_VOCAB_PARALLEL_GREEDY=1 \ + LIGHTLLM_SYMM_MEM_OUT_OF_PLACE=1 \ + LIGHTLLM_LOG_LEVEL=warning \ + NCCL_CUMEM_ENABLE=1 \ + NCCL_NVLS_ENABLE=1 \ + CUDA_DEVICE_MAX_CONNECTIONS=8 \ + TRTLLM_ENABLE_PDL=1 \ + NCCL_GRAPH_MIXING_SUPPORT=0 + +WORKDIR /opt/lightllm + +RUN test "$(command -v python)" = /opt/sglang/bin/python && \ + python -c "import hypercorn, importlib.metadata as metadata, lightllm.server.api_start, redis, rpyc, torch; assert metadata.version('hypercorn') == '0.18.0'; assert redis.__version__ == '7.3.0'; assert rpyc.__version__ == '6.0.2'; print(torch.__version__)" + +EXPOSE 8002 +STOPSIGNAL SIGTERM +ENTRYPOINT ["/opt/nvidia/nvidia_entrypoint.sh"] + +HEALTHCHECK --interval=30s --timeout=5s --start-period=15m --retries=3 \ + CMD /opt/sglang/bin/python -c "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8002/v1/models', timeout=3).read()" || exit 1 + +# Exact 16K-input/256-output profile validated on an eight-H100 80 GB node. +# The larger KV capacity keeps a 64-request prompt wave resident, and decode +# CUDA graphs cover batches through 128. +CMD ["/opt/sglang/bin/python", "-m", "lightllm.server.api_server", \ + "--model_dir", "/model", \ + "--model_name", "glm-5.3-flash", \ + "--tp", "8", \ + "--host", "0.0.0.0", \ + "--port", "8002", \ + "--httpserver_workers", "16", \ + "--mem_fraction", ".90", \ + "--max_total_token_num", "1120000", \ + "--running_max_req_size", "256", \ + "--max_req_total_len", "532480", \ + "--batch_max_tokens", "32768", \ + "--chunked_prefill_size", "32768", \ + "--linear_att_ssm_data_type", "bfloat16", \ + "--linear_att_cache_size", "256", \ + "--graph_max_batch_size", "128", \ + "--graph_split_batch_size", "8", \ + "--graph_grow_step_size", "8", \ + "--graph_max_len_in_batch", "32768", \ + "--disable_flashinfer_allreduce", \ + "--enable_fused_shared_experts", \ + "--disable_dynamic_prompt_cache", \ + "--router_max_wait_tokens", "64", \ + "--disable_vision", \ + "--cache_capacity", "64", \ + "--schedule_time_interval", "0.001", \ + "--reasoning_parser", "glm45", \ + "--tool_call_parser", "glm47"] diff --git a/docker/Dockerfile.glm53-h100-overlay b/docker/Dockerfile.glm53-h100-overlay new file mode 100644 index 0000000000..72b46a925c --- /dev/null +++ b/docker/Dockerfile.glm53-h100-overlay @@ -0,0 +1,66 @@ +# Fast release path for environments where the original Docker Hub base is +# unavailable. The pinned base is the already-flattened GLM-5.3 runtime; this +# layer replaces only the public LightLLM source and the validated default CMD. +ARG BASE_IMAGE=registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm@sha256:3801951c8697edcab68ac5d38efc1807fcda0cba3d7946125a04103c302a3845 +FROM ${BASE_IMAGE} + +ARG OCI_CREATED +ARG OCI_REVISION +ARG OCI_SOURCE=https://github.com/sufubao/LightLLM +ARG OCI_VERSION +ARG BASE_NAME=registry.ms-sc-01.maoshanwangtech.com/ms-ccr/lightllm +ARG BASE_DIGEST=sha256:3801951c8697edcab68ac5d38efc1807fcda0cba3d7946125a04103c302a3845 + +WORKDIR /opt/lightllm + +COPY lightllm ./lightllm + +RUN test -f 'lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/chunk_gated_delta_rule_fwd_h/{BT=64,H=8,K=128,V=128}_NVIDIA_H100_80GB_HBM3.json' && \ + test -f 'lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/silu_and_mul_fwd:v1/{N=1536,out_dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json' && \ + test -f 'lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/silu_and_mul_fwd:v1/{N=2048,out_dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json' && \ + python -c "import lightllm.server.api_start; print('Updated LightLLM GLM-5.3 runtime import OK')" + +LABEL org.opencontainers.image.created="${OCI_CREATED}" \ + org.opencontainers.image.revision="${OCI_REVISION}" \ + org.opencontainers.image.source="${OCI_SOURCE}" \ + org.opencontainers.image.version="${OCI_VERSION}" \ + org.opencontainers.image.title="LightLLM GLM-5.3-Flash H100/H200 TP8 optimized text" \ + org.opencontainers.image.base.name="${BASE_NAME}" \ + org.opencontainers.image.base.digest="${BASE_DIGEST}" \ + ai.lightllm.model="GLM-5.3-Flash" \ + ai.lightllm.accelerator="NVIDIA H100/H200" \ + ai.lightllm.tensor-parallel-size="8" \ + ai.lightllm.profile="text-16k256-tp8-kpool-c256-no-prompt-cache" \ + ai.lightllm.security-profile="flattened-base-plus-source-overlay" + +# Exact 16K-input/256-output profile validated on an eight-H100 80 GB node. +# The larger KV capacity keeps a 64-request prompt wave resident, and decode +# CUDA graphs cover batches through 128. +CMD ["/opt/sglang/bin/python", "-m", "lightllm.server.api_server", \ + "--model_dir", "/model", \ + "--model_name", "glm-5.3-flash", \ + "--tp", "8", \ + "--host", "0.0.0.0", \ + "--port", "8002", \ + "--httpserver_workers", "16", \ + "--mem_fraction", ".90", \ + "--max_total_token_num", "1120000", \ + "--running_max_req_size", "256", \ + "--max_req_total_len", "532480", \ + "--batch_max_tokens", "32768", \ + "--chunked_prefill_size", "32768", \ + "--linear_att_ssm_data_type", "bfloat16", \ + "--linear_att_cache_size", "256", \ + "--graph_max_batch_size", "128", \ + "--graph_split_batch_size", "8", \ + "--graph_grow_step_size", "8", \ + "--graph_max_len_in_batch", "32768", \ + "--disable_flashinfer_allreduce", \ + "--enable_fused_shared_experts", \ + "--disable_dynamic_prompt_cache", \ + "--router_max_wait_tokens", "64", \ + "--disable_vision", \ + "--cache_capacity", "64", \ + "--schedule_time_interval", "0.001", \ + "--reasoning_parser", "glm45", \ + "--tool_call_parser", "glm47"] diff --git a/docker/glm53-h100.openvex.json b/docker/glm53-h100.openvex.json new file mode 100644 index 0000000000..90a83eebd1 --- /dev/null +++ b/docker/glm53-h100.openvex.json @@ -0,0 +1,48 @@ +{ + "@context": "https://openvex.dev/ns/v0.2.0", + "@id": "https://github.com/sufubao/LightLLM/security/vex/glm53-h100-sglang-unused-components", + "author": "https://github.com/sufubao/LightLLM", + "timestamp": "2026-08-28T11:45:00Z", + "version": 1, + "statements": [ + { + "vulnerability": { + "name": "CVE-2026-3059" + }, + "products": [ + { + "@id": "pkg:pypi/sglang@0.0.0.dev1%2Bgf609d677b" + } + ], + "status": "not_affected", + "justification": "vulnerable_code_not_present", + "impact_statement": "The image removes sglang/multimodal_gen; LightLLM only imports SGLang's Triton MoE kernels." + }, + { + "vulnerability": { + "name": "CVE-2026-3060" + }, + "products": [ + { + "@id": "pkg:pypi/sglang@0.0.0.dev1%2Bgf609d677b" + } + ], + "status": "not_affected", + "justification": "vulnerable_code_not_present", + "impact_statement": "The image removes sglang/srt/disaggregation; LightLLM does not expose SGLang's disaggregation service." + }, + { + "vulnerability": { + "name": "CVE-2026-3989" + }, + "products": [ + { + "@id": "pkg:pypi/sglang@0.0.0.dev1%2Bgf609d677b" + } + ], + "status": "not_affected", + "justification": "vulnerable_code_not_present", + "impact_statement": "The image removes SGLang's replay_request_dump.py developer utility." + } + ] +} diff --git a/docker/requirements-glm53-runtime.txt b/docker/requirements-glm53-runtime.txt new file mode 100644 index 0000000000..08740b3440 --- /dev/null +++ b/docker/requirements-glm53-runtime.txt @@ -0,0 +1,30 @@ +atomics==1.0.3 \ + --hash=sha256:19ed27f1ae3fe3353e9103e6d6f54af83e56a929f78ba061ae2e8e9435101daa +diffusers==0.38.0 \ + --hash=sha256:18e53f9e539096320470f62c6360a6fd5727ff28cffda566265316e13fcdb612 +frozendict==2.4.7 \ + --hash=sha256:972af65924ea25cf5b4d9326d549e69a9a4918d8a76a9d3a7cd174d98b237550 +h2==4.4.1 \ + --hash=sha256:0e25f1462b23c9cb82d9eb02e28bc706dac2a68cb457c6a0d74d63c8a2a5d0e6 +hpack==4.2.0 \ + --hash=sha256:858ac0b02280fa582b5080d68db0899c62a80375e0e5413a74970c5e518b6986 +hypercorn==0.18.0 \ + --hash=sha256:225e268f2c1c2f28f6d8f6db8f40cb8c992963610c5725e13ccfcddccb24b1cd +hyperframe==6.1.0 \ + --hash=sha256:b03380493a519fce58ea5af42e4a42317bf9bd425596f7a0835ffce80f1a42e5 +msgpack==1.2.2 \ + --hash=sha256:77c2e018417dc1d66f235e383877ee885b60ade9d29e494dd581e08af2cb1923 +plumbum==2.0.2 \ + --hash=sha256:a865771826c1d2cce0ca8b2a2c48aa0447ed7a518859e389a95f158e9c369547 +priority==2.0.0 \ + --hash=sha256:6f8eefce5f3ad59baf2c080a664037bb4725cd0a790d53d59ab4059288faf6aa +rpyc==6.0.2 \ + --hash=sha256:8072308ad30725bc281c42c011fc8c922be15f3eeda6eafb2917cafe1b6f00ec +redis==7.3.0 \ + --hash=sha256:9d4fcb002a12a5e3c3fbe005d59c48a2cc231f87fbb2f6b70c2d89bb64fec364 +setuptools==84.0.0 \ + --hash=sha256:51a52592b3b99e102b609654876bd65f19f999935166d1352678931132b0c670 +ujson==5.13.0 \ + --hash=sha256:ea7204e9fa7538bfbb1396e1ee8c2bbcd3818b3633ef5bb14d4fdea52994d14d +wsproto==1.3.2 \ + --hash=sha256:61eea322cdf56e8cc904bd3ad7573359a242ba65688716b0710a5eb12beab584 diff --git a/docker/split_glm53_rootfs.py b/docker/split_glm53_rootfs.py new file mode 100644 index 0000000000..e5af15989b --- /dev/null +++ b/docker/split_glm53_rootfs.py @@ -0,0 +1,195 @@ +#!/usr/bin/python3 +"""Move large runtime trees into deterministic, size-bounded image layers.""" + +from __future__ import annotations + +import errno +import os +import shutil +import stat +from pathlib import Path +from typing import NamedTuple + + +ATOM_LIMIT = 512 * 1024 * 1024 +BUCKET_COUNT = 32 +BUCKET_LIMIT = 1_250_000_000 +LAYER_ROOT = Path("/__image_layers") +SOURCES = ( + Path("/opt"), + Path("/root"), + Path("/sgl-workspace"), + Path("/usr/include"), + Path("/usr/lib/x86_64-linux-gnu"), + Path("/usr/libexec"), + Path("/usr/local"), + Path("/usr/share"), +) + + +class Node(NamedTuple): + path: Path + size: int + children: tuple["Node", ...] + + +def scan_tree(path: Path, directory_stats: dict[Path, os.stat_result]) -> Node: + metadata = path.lstat() + if not stat.S_ISDIR(metadata.st_mode): + return Node(path, metadata.st_size, ()) + + directory_stats[path] = metadata + children = tuple( + scan_tree(Path(entry.path), directory_stats) for entry in sorted(os.scandir(path), key=lambda item: item.name) + ) + return Node(path, metadata.st_size + sum(child.size for child in children), children) + + +def iter_atoms(node: Node): + if not node.children or node.size <= ATOM_LIMIT: + yield node.path, node.size + return + for child in node.children: + yield from iter_atoms(child) + + +def make_parent_directories( + source: Path, + destination: Path, + directory_stats: dict[Path, os.stat_result], + created_directories: dict[Path, Path], +) -> None: + source_parent = Path("/") + destination_parent = destination.parent + relative_parents = source.relative_to("/").parts[:-1] + destination_cursor = LAYER_ROOT / destination.relative_to(LAYER_ROOT).parts[0] + + for component in relative_parents: + source_parent /= component + destination_cursor /= component + if not destination_cursor.exists(): + destination_cursor.mkdir() + created_directories[destination_cursor] = source_parent + + assert destination_cursor == destination_parent + + +def restore_directory_metadata(destination: Path, source: Path, directory_stats: dict[Path, os.stat_result]) -> None: + metadata = directory_stats[source] + os.chown(destination, metadata.st_uid, metadata.st_gid, follow_symlinks=False) + os.chmod(destination, stat.S_IMODE(metadata.st_mode), follow_symlinks=False) + os.utime( + destination, + ns=(metadata.st_atime_ns, metadata.st_mtime_ns), + follow_symlinks=False, + ) + + +def copy_path( + source: Path, + destination: Path, + copied_inodes: dict[tuple[int, int], Path], +) -> None: + metadata = source.lstat() + inode = (metadata.st_dev, metadata.st_ino) + + if stat.S_ISLNK(metadata.st_mode): + destination.symlink_to(os.readlink(source)) + os.lchown(destination, metadata.st_uid, metadata.st_gid) + shutil.copystat(source, destination, follow_symlinks=False) + return + + if stat.S_ISDIR(metadata.st_mode): + destination.mkdir(mode=stat.S_IMODE(metadata.st_mode)) + for entry in sorted(os.scandir(source), key=lambda item: item.name): + copy_path(Path(entry.path), destination / entry.name, copied_inodes) + os.chown(destination, metadata.st_uid, metadata.st_gid) + shutil.copystat(source, destination, follow_symlinks=False) + return + + if stat.S_ISREG(metadata.st_mode): + previous = copied_inodes.get(inode) + if metadata.st_nlink > 1 and previous is not None: + destination.hardlink_to(previous) + else: + shutil.copyfile(source, destination, follow_symlinks=False) + copied_inodes[inode] = destination + os.chown(destination, metadata.st_uid, metadata.st_gid) + shutil.copystat(source, destination, follow_symlinks=False) + return + + if stat.S_ISFIFO(metadata.st_mode): + os.mkfifo(destination, stat.S_IMODE(metadata.st_mode)) + elif stat.S_ISCHR(metadata.st_mode) or stat.S_ISBLK(metadata.st_mode): + os.mknod(destination, metadata.st_mode, metadata.st_rdev) + else: + raise RuntimeError(f"unsupported file type while copying {source}") + os.chown(destination, metadata.st_uid, metadata.st_gid) + shutil.copystat(source, destination, follow_symlinks=False) + + +def move_path( + source: Path, + destination: Path, + copied_inodes: dict[tuple[int, int], Path], +) -> None: + try: + source.rename(destination) + return + except OSError as error: + if error.errno != errno.EXDEV: + raise + + copy_path(source, destination, copied_inodes) + if source.is_dir() and not source.is_symlink(): + shutil.rmtree(source) + else: + source.unlink() + + +def main() -> None: + Path(__file__).unlink() + LAYER_ROOT.mkdir(mode=0o755) + + directory_stats: dict[Path, os.stat_result] = {} + for source in SOURCES: + parent = source.parent + while parent != Path("/"): + directory_stats.setdefault(parent, parent.lstat()) + parent = parent.parent + nodes = tuple(scan_tree(source, directory_stats) for source in SOURCES) + atoms = sorted( + (atom for node in nodes for atom in iter_atoms(node)), + key=lambda item: (-item[1], str(item[0])), + ) + + bucket_sizes = [0] * BUCKET_COUNT + assignments: list[tuple[Path, int]] = [] + for source, size in atoms: + bucket = min(range(BUCKET_COUNT), key=lambda index: (bucket_sizes[index], index)) + bucket_sizes[bucket] += size + assignments.append((source, bucket)) + + largest_bucket = max(bucket_sizes) + if largest_bucket > BUCKET_LIMIT: + raise RuntimeError(f"largest rootfs bucket is {largest_bucket} bytes, over {BUCKET_LIMIT}") + + created_directories: dict[Path, Path] = {} + for index in range(BUCKET_COUNT): + (LAYER_ROOT / f"{index:02d}").mkdir(mode=0o755) + + copied_inodes: dict[tuple[int, int], Path] = {} + for source, bucket in sorted(assignments, key=lambda item: str(item[0])): + destination = LAYER_ROOT / f"{bucket:02d}" / source.relative_to("/") + make_parent_directories(source, destination, directory_stats, created_directories) + move_path(source, destination, copied_inodes) + + for destination, source in sorted(created_directories.items(), key=lambda item: len(item[0].parts), reverse=True): + restore_directory_metadata(destination, source, directory_stats) + + for index, size in enumerate(bucket_sizes): + print(f"rootfs layer {index:02d}: {size} uncompressed bytes") + + +if __name__ == "__main__": + main() diff --git a/lightllm/common/basemodel/attention/__init__.py b/lightllm/common/basemodel/attention/__init__.py index 10cd3b0864..1cd6d0aff7 100644 --- a/lightllm/common/basemodel/attention/__init__.py +++ b/lightllm/common/basemodel/attention/__init__.py @@ -13,6 +13,7 @@ # NSA backend from .nsa.flashmla_sparse import NsaFlashMlaSparseAttBackend from .nsa.fp8_flashmla_sparse import NsaFlashMlaFp8SparseAttBackend +from .nsa.tilelang_sparse import NsaTilelangSparseAttBackend from .create_utils import ( get_prefill_att_backend_class, diff --git a/lightllm/common/basemodel/attention/base_att.py b/lightllm/common/basemodel/attention/base_att.py index a7e2d8122a..6cd87e9f2d 100644 --- a/lightllm/common/basemodel/attention/base_att.py +++ b/lightllm/common/basemodel/attention/base_att.py @@ -120,6 +120,18 @@ class BasePrefillAttState(ABC): backend: BaseAttBackend = None infer_state: "InferStateInfo" = None + def copy_for_prefill_cuda_graph(self, new_state: "BasePrefillAttState"): + """Refresh fixed-address attention metadata before graph replay.""" + for attr_name, attr_value in vars(new_state).items(): + if isinstance(attr_value, torch.Tensor): + graph_attr = getattr(self, attr_name, None) + if ( + graph_attr is not None + and graph_attr.data_ptr() != attr_value.data_ptr() + and graph_attr.shape == attr_value.shape + ): + graph_attr.copy_(attr_value, non_blocking=True) + @abstractmethod def init_state(self): pass diff --git a/lightllm/common/basemodel/attention/create_utils.py b/lightllm/common/basemodel/attention/create_utils.py index f708357765..f7ac37943e 100644 --- a/lightllm/common/basemodel/attention/create_utils.py +++ b/lightllm/common/basemodel/attention/create_utils.py @@ -17,6 +17,7 @@ from .flashinfer.mla import MlaFlashInferAttBackend from .nsa.flashmla_sparse import NsaFlashMlaSparseAttBackend from .nsa.fp8_flashmla_sparse import NsaFlashMlaFp8SparseAttBackend +from .nsa.tilelang_sparse import NsaTilelangSparseAttBackend logger = init_logger(__name__) @@ -56,7 +57,8 @@ nsa_data_type_to_backend = { "None": { "flashmla_sparse": NsaFlashMlaSparseAttBackend, - # Future backends: "fa3", "tilelang", "aiter" + "tilelang": NsaTilelangSparseAttBackend, + # Future backends: "fa3", "aiter" }, "fp8kv_dsa": { "flashmla_sparse": NsaFlashMlaFp8SparseAttBackend, diff --git a/lightllm/common/basemodel/attention/linear/__init__.py b/lightllm/common/basemodel/attention/linear/__init__.py index bfc8445390..b1a3d1ca34 100644 --- a/lightllm/common/basemodel/attention/linear/__init__.py +++ b/lightllm/common/basemodel/attention/linear/__init__.py @@ -5,6 +5,7 @@ ) from .flashqla import FlashQlaLinearAttBackend from .triton import TritonLinearAttBackend +from .kda import KDALinearAttBackend __all__ = [ "LinearAttBackend", @@ -12,4 +13,5 @@ "LinearAttDecodeAttState", "FlashQlaLinearAttBackend", "TritonLinearAttBackend", + "KDALinearAttBackend", ] diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index ca6ceaec43..85a19107f8 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -3,7 +3,8 @@ from abc import ABC, abstractmethod from typing import TYPE_CHECKING from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl -from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.envs_utils import get_env_start_args, get_triton_autotune_level +from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneLevel from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_fn from lightllm.common.basemodel.triton_kernel.linear_att.fused_gdn_gating import fused_gdn_gating from lightllm.common.basemodel.triton_kernel.linear_att.gdn_decode_pack import conv_pack_gdn_decode_inputs @@ -247,7 +248,7 @@ def _init_fixed_mtp_decode_state(self, draft_step: int): dtype=torch.int32, device=self.infer_state.b_req_idx.device, ) - self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous() + self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0] self.b_num_accepted_tokens = self.infer_state.req_manager.req_to_mtp_state_index[self.b_conv_buffer_idx] + 1 self._init_mtp_ssm_buffer_idx(mtp_size) @@ -388,18 +389,38 @@ def _gdn_mtp_kernel( # #8b: b_num_accepted_tokens >= 1 is guaranteed upstream: init/cache restore set 1, # and MTP decode only writes values in [1, mtp_step+1]. The old per-layer per-step # .all() D2H sync stalled the GPU on the eager decode hot path; it is redundant here. - core_attn_out, _ = mtp_fused_recurrent_gated_delta_rule( - q=query, - k=key, - v=value, - initial_state=ssm_states, - cu_seqlens=cu_seqlens_q.to(torch.long), - ssm_state_indices=self.b_ssm_buffer_idx, - ssm_state_write_indices=self.b_ssm_buffer_idx, - num_accepted_tokens=self.b_num_accepted_tokens, - A_log=layer_weight.linear_A_log.weight, - dt_bias=layer_weight.linear_dt_bias.weight, - a_raw=a, - b_raw=b, + fixed_seq_len = ( + 0 + if backend.uses_dynamic_spec_verify_layout() + else backend.model.mtp_manager.get_decode_draft_step(backend.model.is_mtp_draft_model) + 1 + ) + tune_now = ( + get_triton_autotune_level() in [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE] + and getattr(infer_state, "is_cuda_graph", False) + and infer_state.microbatch_index == 0 + and layer_weight.layer_num_ == 0 + and not torch.cuda.is_current_stream_capturing() + and not Autotuner.is_autotune_warmup() ) + if tune_now: + Autotuner.start_autotune_warmup() + try: + core_attn_out, _ = mtp_fused_recurrent_gated_delta_rule( + q=query, + k=key, + v=value, + initial_state=ssm_states, + cu_seqlens=cu_seqlens_q, + ssm_state_indices=self.b_ssm_buffer_idx, + ssm_state_write_indices=self.b_ssm_buffer_idx, + num_accepted_tokens=self.b_num_accepted_tokens, + A_log=layer_weight.linear_A_log.weight, + dt_bias=layer_weight.linear_dt_bias.weight, + a_raw=a, + b_raw=b, + fixed_seq_len=fixed_seq_len, + ) + finally: + if tune_now: + Autotuner.end_autotune_warmup() return core_attn_out diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py new file mode 100644 index 0000000000..ac31c02c73 --- /dev/null +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -0,0 +1,302 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""KDA attention backend for GLM-5-Next.""" + +from __future__ import annotations + +import dataclasses +from typing import TYPE_CHECKING + +import torch + +from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( + causal_conv1d_fn, + causal_conv1d_update, +) +from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( + build_dynamic_mtp_linear_att_state_params, +) +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import ( + chunk_kda_with_fused_gate, + fused_recurrent_kda, +) +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.index import prepare_chunk_indices +from lightllm.utils.envs_utils import get_env_start_args + +if TYPE_CHECKING: + from lightllm.common.basemodel.basemodel import TpPartBaseModel + from lightllm.common.basemodel.infer_struct import InferStateInfo + + +class KDALinearAttBackend(BaseAttBackend): + def __init__(self, model: "TpPartBaseModel"): + super().__init__(model=model) + config = model.config["linear_attn_config"] + self.num_heads = config["num_heads"] + self.head_dim = config["head_dim"] + assert self.num_heads % model.tp_world_size_ == 0 + self.tp_num_heads = self.num_heads // model.tp_world_size_ + self.tp_projection_size = self.tp_num_heads * self.head_dim + self.conv_kernel_size = config["short_conv_kernel_size"] + self.lower_bound = config.get("gate_lower_bound", -5.0) + self.mtp_step = get_env_start_args().mtp_step + + def create_att_prefill_state(self, infer_state: "InferStateInfo"): + return KDAPrefillAttState(backend=self, infer_state=infer_state) + + def create_att_decode_state(self, infer_state: "InferStateInfo"): + return KDADecodeAttState(backend=self, infer_state=infer_state) + + def split_qkv(self, mixed_qkv: torch.Tensor): + return mixed_qkv.split(self.tp_projection_size, dim=-1) + + def reshape_qkv(self, value: torch.Tensor, *, decode: bool): + if decode: + return value.view(-1, 1, self.tp_num_heads, self.head_dim) + return value.view(1, -1, self.tp_num_heads, self.head_dim) + + +@dataclasses.dataclass +class KDAPrefillAttState(BasePrefillAttState): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + chunk_indices: torch.Tensor = None + seq_lens_cpu: list[int] = None + + def init_state(self): + self.b_conv_buffer_idx = self.infer_state.b_req_idx + self.b_ssm_buffer_idx = self.infer_state.b_req_idx * (self.backend.mtp_step + 1) + self.seq_lens_cpu = (self.infer_state.b1_cu_q_seq_len[1:] - self.infer_state.b1_cu_q_seq_len[:-1]).tolist() + # prepare_chunk_indices performs a GPU-to-CPU shape sync. Build it + # before entering CUDA Graph capture and copy its fixed-size contents + # through BasePrefillAttState on replay. + self.chunk_indices = prepare_chunk_indices(self.infer_state.b1_cu_q_seq_len, 64) + + def prefill_att( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + att_control: AttControl = AttControl(), + alloc_func=torch.empty, + ): + assert att_control.linear_att_prefill + params = att_control.linear_att_prefill_dict + layer_weight = params["layer_weight"] + layer_num = params["layer_num"] + mixed_qkv = params["mixed_qkv"] + raw_gate = params["raw_gate"] + raw_beta = params["raw_beta"] + backend: KDALinearAttBackend = self.backend + + conv_states, ssm_states = self.infer_state.req_manager.get_mamba_cache(layer_num) + # MTP widens each request's convolution cache so the verify path can + # retain all candidate states. Prefill only populates the canonical + # history, whose width remains kernel_size - 1. + if backend.mtp_step > 0: + conv_states = conv_states[:, :, : -backend.mtp_step] + mixed_qkv = causal_conv1d_fn( + mixed_qkv.transpose(0, 1), + layer_weight.get_merged_kda_conv_weight(), + bias=None, + query_start_loc=self.infer_state.b1_cu_q_seq_len, + cache_indices=self.b_conv_buffer_idx, + has_initial_state=self.infer_state.b_ready_cache_len > 0, + conv_states=conv_states, + activation="silu", + seq_lens_cpu=self.seq_lens_cpu, + ).transpose(0, 1) + + q, k, v = [backend.reshape_qkv(x, decode=False) for x in backend.split_qkv(mixed_qkv)] + raw_gate = raw_gate.view(1, -1, backend.tp_projection_size) + raw_beta = raw_beta.view(1, -1, backend.tp_num_heads) + + initial_state = ssm_states[self.b_ssm_buffer_idx].contiguous() + output, final_state = chunk_kda_with_fused_gate( + q=q, + k=k, + v=v, + raw_g=raw_gate.view(1, -1, backend.tp_num_heads, backend.head_dim), + beta=raw_beta.float().sigmoid(), + A_log=layer_weight.linear_A_log.weight, + g_bias=layer_weight.linear_dt_bias.weight, + initial_state=initial_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=self.infer_state.b1_cu_q_seq_len, + chunk_indices=self.chunk_indices, + safe_gate=True, + lower_bound=backend.lower_bound, + ) + ssm_states[self.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype, copy=False) + return output + + +@dataclasses.dataclass +class KDADecodeAttState(BaseDecodeAttState): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + b1_mtp_cu_q_seq_len: torch.Tensor = None + b_num_accepted_tokens: torch.Tensor = None + + def init_state(self): + draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) + if draft_step == 0: + self._init_normal_decode_state() + elif self.backend.uses_dynamic_spec_verify_layout(): + self._init_dynamic_mtp_decode_state(draft_step + 1) + else: + self._init_fixed_mtp_decode_state(draft_step) + + def _init_normal_decode_state(self): + self.b_conv_buffer_idx = self.infer_state.b_req_idx + self.b_ssm_buffer_idx = self.infer_state.b_req_idx + + def _init_dynamic_mtp_decode_state(self, mtp_size: int): + ( + self.b1_mtp_cu_q_seq_len, + self.b_conv_buffer_idx, + self.b_num_accepted_tokens, + ) = build_dynamic_mtp_linear_att_state_params( + b_req_idx=self.infer_state.b_req_idx, + b_mtp_index=self.infer_state.b_mtp_index, + req_to_mtp_state_index=self.infer_state.req_manager.req_to_mtp_state_index, + hold_req_id=self.infer_state.req_manager.HOLD_REQUEST_ID, + ) + self._init_mtp_ssm_buffer_idx(mtp_size) + + def _init_fixed_mtp_decode_state(self, draft_step: int): + mtp_size = draft_step + 1 + batch_size = self.infer_state.batch_size + assert batch_size % mtp_size == 0, ( + "KDA fixed-layout decode requires batch_size to be divisible by draft_step + 1, " + f"got batch_size={batch_size}, draft_step={draft_step}." + ) + + att_batch_size = batch_size // mtp_size + self.b1_mtp_cu_q_seq_len = torch.arange( + 0, + batch_size + 1, + mtp_size, + dtype=torch.int32, + device=self.infer_state.b_req_idx.device, + ) + self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous() + self.b_num_accepted_tokens = self.infer_state.req_manager.req_to_mtp_state_index[self.b_conv_buffer_idx] + 1 + self._init_mtp_ssm_buffer_idx(mtp_size) + + def _init_mtp_ssm_buffer_idx(self, mtp_size: int): + att_batch_size = self.b_conv_buffer_idx.shape[0] + b_ssm_buffer_start_idx = (self.b_conv_buffer_idx * mtp_size).view(att_batch_size, 1) + state_offsets = torch.arange( + mtp_size, + device=self.infer_state.b_req_idx.device, + dtype=self.infer_state.b_req_idx.dtype, + ).view(1, mtp_size) + self.b_ssm_buffer_idx = b_ssm_buffer_start_idx + state_offsets + + def decode_att( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + att_control: AttControl = AttControl(), + alloc_func=torch.empty, + ): + assert att_control.linear_att_decode + params = att_control.linear_att_decode_dict + layer_weight = params["layer_weight"] + layer_num = params["layer_num"] + mixed_qkv = params["mixed_qkv"] + raw_gate = params["raw_gate"] + raw_beta = params["raw_beta"] + backend: KDALinearAttBackend = self.backend + + conv_states, ssm_states = self.infer_state.req_manager.get_mamba_cache(layer_num) + draft_step = backend.model.mtp_manager.get_decode_draft_step(backend.model.is_mtp_draft_model) + if draft_step > 0: + return self._kda_mtp_kernel( + mixed_qkv=mixed_qkv, + raw_gate=raw_gate, + raw_beta=raw_beta, + conv_states=conv_states, + ssm_states=ssm_states, + layer_weight=layer_weight, + draft_step=draft_step, + ) + + mixed_qkv = causal_conv1d_update( + mixed_qkv, + conv_states, + layer_weight.get_merged_kda_conv_weight(), + bias=None, + activation="silu", + conv_state_indices=self.b_conv_buffer_idx, + ) + q, k, v = [backend.reshape_qkv(x, decode=True) for x in backend.split_qkv(mixed_qkv)] + raw_gate = raw_gate.view(-1, 1, backend.tp_projection_size) + raw_beta = raw_beta.view(-1, 1, backend.tp_num_heads) + output, _ = fused_recurrent_kda( + q=q, + k=k, + v=v, + raw_gate=raw_gate, + raw_beta=raw_beta, + a_log=layer_weight.linear_A_log.weight, + gate_bias=layer_weight.linear_dt_bias.weight, + initial_state=ssm_states, + lower_bound=backend.lower_bound, + inplace_final_state=True, + ssm_state_indices=self.b_ssm_buffer_idx, + ) + return output + + def _kda_mtp_kernel( + self, + mixed_qkv: torch.Tensor, + raw_gate: torch.Tensor, + raw_beta: torch.Tensor, + conv_states: torch.Tensor, + ssm_states: torch.Tensor, + layer_weight, + draft_step: int, + ): + from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, + ) + + backend: KDALinearAttBackend = self.backend + mixed_qkv = causal_conv1d_update_mtp( + mixed_qkv, + conv_states, + layer_weight.get_merged_kda_conv_weight(), + mtp_step=draft_step, + bias=None, + activation="silu", + conv_state_indices=self.b_conv_buffer_idx, + num_accepted_tokens=self.b_num_accepted_tokens, + query_start_loc=self.b1_mtp_cu_q_seq_len, + ) + + q, k, v = [backend.reshape_qkv(x, decode=False) for x in backend.split_qkv(mixed_qkv)] + raw_gate = raw_gate.view(1, -1, backend.tp_projection_size) + raw_beta = raw_beta.view(1, -1, backend.tp_num_heads) + assert self.b_ssm_buffer_idx.dim() == 2, "KDA MTP SSM buffer idx must be 2D [N, S+1]" + output, _ = fused_recurrent_kda( + q=q, + k=k, + v=v, + raw_gate=raw_gate, + raw_beta=raw_beta, + a_log=layer_weight.linear_A_log.weight, + gate_bias=layer_weight.linear_dt_bias.weight, + initial_state=ssm_states, + lower_bound=backend.lower_bound, + inplace_final_state=True, + cu_seqlens=self.b1_mtp_cu_q_seq_len.to(torch.long), + ssm_state_indices=self.b_ssm_buffer_idx, + ssm_state_write_indices=self.b_ssm_buffer_idx, + num_accepted_tokens=self.b_num_accepted_tokens, + ) + return output diff --git a/lightllm/common/basemodel/attention/nsa/__init__.py b/lightllm/common/basemodel/attention/nsa/__init__.py index f9db52dc2b..b183461504 100644 --- a/lightllm/common/basemodel/attention/nsa/__init__.py +++ b/lightllm/common/basemodel/attention/nsa/__init__.py @@ -10,6 +10,7 @@ NsaFlashMlaFp8SparsePrefillAttState, NsaFlashMlaFp8SparseDecodeAttState, ) +from .tilelang_sparse import NsaTilelangSparseAttBackend, NsaTilelangSparsePrefillAttState __all__ = [ "NsaFlashMlaSparseAttBackend", diff --git a/lightllm/common/basemodel/attention/nsa/flashmla_sparse.py b/lightllm/common/basemodel/attention/nsa/flashmla_sparse.py index c25f1d31e7..6bf4ca4334 100644 --- a/lightllm/common/basemodel/attention/nsa/flashmla_sparse.py +++ b/lightllm/common/basemodel/attention/nsa/flashmla_sparse.py @@ -3,15 +3,70 @@ import dataclasses import torch +import torch.distributed as dist from typing import Tuple, TYPE_CHECKING from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl from lightllm.utils.dist_utils import get_current_device_id +from lightllm.utils.dist_utils import get_current_rank_in_dp +from lightllm.utils.envs_utils import get_env_start_args if TYPE_CHECKING: from lightllm.common.basemodel.infer_struct import InferStateInfo +_TP_HEAD_TOKEN_TRANSPOSE_MIN_TOKENS = 4096 + + +def _copy_received_head_shards(received: torch.Tensor, output: torch.Tensor, world_size: int) -> None: + """Transpose all-to-all receive order from rank-major to token-major heads.""" + + tokens, local_heads, head_dim = received.shape + tokens_per_rank = tokens // world_size + output.view(tokens_per_rank, world_size, local_heads, head_dim).copy_( + received.view(world_size, tokens_per_rank, local_heads, head_dim).permute(1, 0, 2, 3) + ) + + +def _copy_token_shard_for_head_scatter(output: torch.Tensor, send: torch.Tensor, world_size: int) -> None: + """Transpose token-major global heads into all-to-all destination order.""" + + tokens_per_rank, global_heads, head_dim = output.shape + local_heads = global_heads // world_size + send.view(world_size, tokens_per_rank, local_heads, head_dim).copy_( + output.view(tokens_per_rank, world_size, local_heads, head_dim).permute(1, 0, 2, 3) + ) + + +def _should_use_tp_head_token_transpose( + q: torch.Tensor, + infer_state: "InferStateInfo", + required_heads: int, +) -> bool: + """Select the exact TP transpose only for its validated serving layout.""" + + args = get_env_start_args() + world_size = infer_state.dist_group.dp_world_size + return ( + world_size > 1 + and q.is_contiguous() + and q.shape[0] >= _TP_HEAD_TOKEN_TRANSPOSE_MIN_TOKENS + and q.shape[0] % world_size == 0 + and q.shape[1] * world_size == required_heads + and infer_state.max_cache_len == 0 + and not infer_state.need_dp_prefill_balance + and not infer_state.use_replicated_attention_ep + and not args.enable_tpsp_mix_mode + and not args.enable_prefill_cudagraph + and not args.enable_prefill_microbatch_overlap + and not args.enable_prefill_decode_mixed + ) + + +def _alloc_like(input_: torch.Tensor, shape: Tuple[int, ...]) -> torch.Tensor: + return torch.empty(shape, dtype=input_.dtype, device=input_.device) + + class NsaFlashMlaSparseAttBackend(BaseAttBackend): def __init__(self, model): super().__init__(model=model) @@ -86,13 +141,70 @@ def _nsa_prefill_att( if topk_mem_indices.ndim == 2: topk_mem_indices = topk_mem_indices.unsqueeze(1) + # The FlashMLA sparse kernels require 64 query heads on Hopper and + # 128 on Blackwell. Tensor parallelism can leave fewer local heads + # (GLM-5.3 has 64 / TP8 = 8), so pad the inactive heads and trim the + # result just as SGLang's DSA backend does. + num_tokens, num_heads, head_dim = q.shape + device_sm_major = torch.cuda.get_device_capability(q.device)[0] + required_heads = 128 if device_sm_major >= 10 else 64 + need_padding = num_heads % required_heads != 0 + if need_padding and _should_use_tp_head_token_transpose(q, self.infer_state, required_heads): + world_size = self.infer_state.dist_group.dp_world_size + rank = get_current_rank_in_dp() + tokens_per_rank = num_tokens // world_size + + received_q = _alloc_like(q, q.shape) + dist.all_to_all_single( + received_q, + q, + group=self.infer_state.dist_group.device_group, + ) + transposed_q = _alloc_like(q, (tokens_per_rank, world_size * num_heads, head_dim)) + _copy_received_head_shards(received_q, transposed_q, world_size) + del received_q + + token_start = rank * tokens_per_rank + local_indices = topk_mem_indices[token_start : token_start + tokens_per_rank] + transposed_out, _, _ = flash_mla_sparse_fwd( + q=transposed_q, + kv=kv, + indices=local_indices, + sm_scale=softmax_scale, + d_v=kv_lora_rank, + ) + + send_out = _alloc_like(q, q.shape) + _copy_token_shard_for_head_scatter(transposed_out, send_out, world_size) + del transposed_q, transposed_out + output = _alloc_like(q, q.shape) + dist.all_to_all_single( + output, + send_out, + group=self.infer_state.dist_group.device_group, + ) + del send_out + return output + + if need_padding: + assert required_heads % num_heads == 0, ( + f"num_heads {num_heads} cannot be padded to {required_heads}; " + "the tensor-parallel size is unsupported" + ) + q_input = q.new_zeros((num_tokens, required_heads, head_dim)) + q_input[:, :num_heads, :] = q + else: + q_input = q + mla_out, _, _ = flash_mla_sparse_fwd( - q=q, + q=q_input, kv=kv, indices=topk_mem_indices, sm_scale=softmax_scale, d_v=kv_lora_rank, ) + if need_padding: + mla_out = mla_out[:, :num_heads, :] return mla_out @@ -173,14 +285,20 @@ def _nsa_decode_att( q_nope, q_rope = q # Extract k_rope and kv_nope from the KV buffer - k_rope = kv[:, :, -qk_rope_head_dim:].view(-1, 1, 1, qk_rope_head_dim) - kv_nope = kv[:, :, :-qk_rope_head_dim].view(-1, 1, 1, kv_lora_rank) + only_qv = qk_rope_head_dim == 0 + if only_qv: + k_rope = None + kv_nope = kv[:, :, :kv_lora_rank].view(-1, 1, 1, kv_lora_rank) + else: + k_rope = kv[:, :, -qk_rope_head_dim:].view(-1, 1, 1, qk_rope_head_dim) + kv_nope = kv[:, :, :-qk_rope_head_dim].view(-1, 1, 1, kv_lora_rank) o_tensor = flash_attn_with_kvcache( - q=q_rope, + q=None if only_qv else q_rope, k_cache=k_rope, v_cache=kv_nope, qv=q_nope, + only_qv=only_qv, page_table=topk_mem_indices, cache_seqlens=self.nsa_cache_seqlens, cu_seqlens_q=self.infer_state.b1_cu_q_seq_len, diff --git a/lightllm/common/basemodel/attention/nsa/tilelang_sparse.py b/lightllm/common/basemodel/attention/nsa/tilelang_sparse.py new file mode 100644 index 0000000000..33578d6335 --- /dev/null +++ b/lightllm/common/basemodel/attention/nsa/tilelang_sparse.py @@ -0,0 +1,84 @@ +"""TileLang sparse prefill attention for NSA/DSA models. + +The kernel is provided by SGLang's kernel package. Decode intentionally keeps +using the FlashMLA implementation: TileLang is selected only for the much +larger prefill workload where it is advantageous on Hopper. +""" + +import dataclasses + +import torch + +from ..base_att import AttControl +from .flashmla_sparse import ( + NsaFlashMlaSparseAttBackend, + NsaFlashMlaSparsePrefillAttState, +) + + +def pad_sparse_indices(indices: torch.Tensor, block_size: int = 64) -> torch.Tensor: + """Mask-pad the sparse index table to the TileLang block width.""" + + if indices.ndim == 2: + indices = indices.unsqueeze(1) + if indices.ndim != 3: + raise ValueError(f"Expected a 2D or 3D sparse index tensor, got shape {tuple(indices.shape)}") + if block_size <= 0: + raise ValueError(f"block_size must be positive, got {block_size}") + + padding = (-indices.shape[-1]) % block_size + if padding: + indices = torch.cat( + (indices, indices.new_full((*indices.shape[:-1], padding), -1)), + dim=-1, + ) + return indices + + +class NsaTilelangSparseAttBackend(NsaFlashMlaSparseAttBackend): + """Use TileLang for prefill and FlashMLA for decode.""" + + def create_att_prefill_state(self, infer_state): + return NsaTilelangSparsePrefillAttState(backend=self, infer_state=infer_state) + + +@dataclasses.dataclass +class NsaTilelangSparsePrefillAttState(NsaFlashMlaSparsePrefillAttState): + def _nsa_prefill_att( + self, + q: torch.Tensor, + kv: torch.Tensor, + att_control: AttControl, + ) -> torch.Tensor: + from sglang.kernels.ops.attention.dsa.tilelang_kernel import tilelang_sparse_fwd + + nsa_dict = att_control.nsa_prefill_dict + # GLM packs its persistent MLA KV and FP8 indexer key in one allocation, + # so the MLA slice has a padded token stride. TileLang requires packed + # KV. For an uncached prefill, the freshly projected batch KV is the + # same ragged ordering addressed by topk_indices; materialize that + # compact view instead of copying the whole persistent cache. Prefix + # cache requests retain the FlashMLA path, which supports the padded + # persistent layout and global memory indices. + if self.infer_state.max_cache_len != 0: + return super()._nsa_prefill_att(q=q, kv=kv, att_control=att_control) + + compact_kv = nsa_dict["prefill_cache_kv"].contiguous() + topk_indices = pad_sparse_indices(nsa_dict["topk_indices"]) + if topk_indices.dtype != torch.int32: + topk_indices = topk_indices.to(torch.int32) + + output = tilelang_sparse_fwd( + q=q, + kv=compact_kv, + indices=topk_indices, + sm_scale=nsa_dict["softmax_scale"], + d_v=nsa_dict["kv_lora_rank"], + ) + # TileLang's generated kernel retains the synthetic batch dimension + # inserted by its Python wrapper. LightLLM uses token-major tensors. + if output.ndim == 4: + if output.shape[0] != 1: + raise RuntimeError(f"Unexpected TileLang sparse output shape {tuple(output.shape)}") + output = output.squeeze(0) + return output diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index e80f2b552f..46f49ebab5 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -1,4 +1,5 @@ import os +import math # os.environ["CUDA_LAUNCH_BLOCKING"] = "1" import gc @@ -23,6 +24,9 @@ from lightllm.common.basemodel.prefill_cuda_graph import PrefillCudaGraph from lightllm.common.quantization import Quantcfg from lightllm.common.basemodel.triton_kernel.gather_token_id import gather_token, gather_token_prefill_decode_mixed +from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_greedy import ( + is_vocab_parallel_greedy_enabled, +) from lightllm.utils.log_utils import init_logger from lightllm.utils.dist_utils import get_dp_world_size from lightllm.utils.profile_max_tokens import profile_mtp_weight_memory @@ -58,6 +62,7 @@ class TpPartBaseModel: is_mtp_draft_model = False + replicated_attention_ep = False # weight class pre_and_post_weight_class = None @@ -90,10 +95,11 @@ def __init__(self, kvargs): if get_env_start_args().enable_decode_microbatch_overlap else self.graph_max_batch_size ) + self.logical_graph_max_batch_size = self.graph_max_batch_size self.mtp_manager = MtpManager.get_instance() - self.graph_max_batch_size = self.graph_max_batch_size * self.mtp_manager.get_decode_batch_multiplier( - self.is_mtp_draft_model - ) + self.decode_batch_multiplier = self.mtp_manager.get_decode_batch_multiplier(self.is_mtp_draft_model) + cuda_graph_batch_multiplier = self.mtp_manager.get_decode_cuda_graph_batch_multiplier(self.is_mtp_draft_model) + self.graph_max_batch_size = self.graph_max_batch_size * cuda_graph_batch_multiplier self.graph_max_len_in_batch = kvargs.get("graph_max_len_in_batch", 8192) self.disable_cudagraph = kvargs.get("disable_cudagraph", False) @@ -103,6 +109,12 @@ def __init__(self, kvargs): self.mem_fraction = kvargs.get("mem_fraction", 0.9) self.tp_world_size_ = get_dp_world_size() self.enable_tpsp_mix_mode = get_env_start_args().enable_tpsp_mix_mode + self.use_replicated_attention_ep = ( + self.replicated_attention_ep + and self.enable_tpsp_mix_mode + and self.args.enable_ep_moe + and not self.is_mtp_draft_model + ) self.torch_memory_saver = TorchMemorySaverWrapper(self.args.enable_torch_memory_saver) self.prefill_graph: PrefillCudaGraph = None @@ -275,19 +287,33 @@ def _init_att_backend1(self): return def _init_cudagraph(self): - decode_batch_multiplier = self.mtp_manager.get_decode_batch_multiplier(self.is_mtp_draft_model) + cuda_graph_batch_multiplier = self.mtp_manager.get_decode_cuda_graph_batch_multiplier(self.is_mtp_draft_model) cuda_graph_grow_step_size = self.mtp_manager.get_decode_cuda_graph_grow_step_size(self.is_mtp_draft_model) + extra_batch_sizes = None + if self.mtp_manager.draft_model_needs_logical_batch_graphs(self.is_mtp_draft_model): + # Recurrent EAGLE alternates between a full-width verify-layout + # extend and one-row-per-request recursive draft forwards. Keep + # the ordinary logical schedule in addition to the widened one so + # neither phase falls back to eager execution or excessive padding. + extra_batch_sizes = CudaGraph.gen_cuda_graph_batch_sizes( + batch_step_size_before_split=1, + split_batch_size=self.args.graph_split_batch_size, + batch_step_size_after_split=self.args.graph_grow_step_size, + max_batch_size=self.logical_graph_max_batch_size, + tp_world_size=self.tp_world_size_, + ) self.graph = ( None if self.disable_cudagraph else CudaGraph( batch_step_size_before_split=cuda_graph_grow_step_size, - split_batch_size=self.args.graph_split_batch_size * decode_batch_multiplier, + split_batch_size=self.args.graph_split_batch_size * cuda_graph_batch_multiplier, batch_step_size_after_split=self.args.graph_grow_step_size * cuda_graph_grow_step_size, max_batch_size=self.graph_max_batch_size, max_len_in_batch=self.graph_max_len_in_batch, tp_world_size=self.tp_world_size_, capture_infer_cost=self.args.mtp_dynamic_verify, + extra_batch_sizes=extra_batch_sizes, ) ) if self.graph is not None: @@ -370,6 +396,31 @@ def _init_hidden_collector(self): @torch.no_grad() def forward(self, model_input: ModelInput): + if model_input.is_prefill: + pool_size = int(self.config.get("index_kpool", 1) or 1) + if pool_size > 1: + # These tensors are normally still on CPU here. Computing the + # batch-wide predicate once avoids synchronizing every NSA + # layer merely to decide whether its batched K-pool path is + # safe. Requiring both boundaries to align also makes every + # query chunk an integer number of pools. + model_input.kpool_prefill_aligned = bool( + torch.all( + (model_input.b_ready_cache_len.remainder(pool_size) == 0) + & (model_input.b_seq_len.remainder(pool_size) == 0) + & (model_input.b_input_len.remainder(pool_size) == 0) + ).item() + ) + else: + pool_size = int(self.config.get("index_kpool", 1) or 1) + model_input.kpool_decode_aligned = bool( + pool_size > 1 + and os.getenv("LIGHTLLM_ENABLE_KPOOL_DECODE_FASTPATH", "0").upper() in {"1", "ON", "TRUE"} + and self.args.mtp_mode is None + and self.args.disable_dynamic_prompt_cache + and self.args.chunked_prefill_size % pool_size == 0 + and torch.all(model_input.b_input_len.remainder(pool_size) == 0).item() + ) model_input.to_cuda() assert model_input.mem_indexes.is_cuda @@ -378,17 +429,30 @@ def forward(self, model_input: ModelInput): else: return self._decode(model_input) + def _is_cuda_graph_output_compatible(self, *model_inputs: ModelInput) -> bool: + """Whether inputs match the dense/sparse output captured at startup.""" + + return ( + self.is_mtp_draft_model + or not is_vocab_parallel_greedy_enabled() + or all(model_input.use_vocab_parallel_greedy for model_input in model_inputs) + ) + def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0): infer_state = self.infer_state_class() + infer_state.use_replicated_attention_ep = self.use_replicated_attention_ep infer_state.hidden_collector = self.hidden_collector_prototype.new_instance() infer_state.input_ids = model_input.input_ids infer_state.is_prefill = model_input.is_prefill infer_state.return_all_prompt_logics = self.return_all_prompt_logics + infer_state.use_vocab_parallel_greedy = self.is_mtp_draft_model or model_input.use_vocab_parallel_greedy infer_state.batch_size = model_input.batch_size infer_state.total_token_num = model_input.total_token_num infer_state.max_q_seq_len = model_input.max_q_seq_len infer_state.max_kv_seq_len = model_input.max_kv_seq_len infer_state.max_cache_len = model_input.max_cache_len + infer_state.kpool_prefill_aligned = model_input.kpool_prefill_aligned + infer_state.kpool_decode_aligned = model_input.kpool_decode_aligned assert model_input.b_req_idx.shape[0] == model_input.b_seq_len.shape[0] infer_state.b_req_idx = model_input.b_req_idx infer_state.b_seq_len = model_input.b_seq_len @@ -449,6 +513,9 @@ def _create_padded_decode_model_input(self, model_input: ModelInput, new_batch_s new_model_input.b_mtp_index, (0, padded_batch_size), mode="constant", value=0 ) new_model_input.b_seq_len = F.pad(new_model_input.b_seq_len, (0, padded_batch_size), mode="constant", value=2) + new_model_input.b_input_len = F.pad( + new_model_input.b_input_len, (0, padded_batch_size), mode="constant", value=2 + ) if new_model_input.b_position_delta is not None: new_model_input.b_position_delta = F.pad( new_model_input.b_position_delta, (0, padded_batch_size), mode="constant", value=0 @@ -534,6 +601,9 @@ def _create_unpad_decode_model_output(self, model_output: ModelOutput, origin_ba return model_output new_model_output = copy.copy(model_output) new_model_output.logits = new_model_output.logits[0:origin_batch_size] + if new_model_output.logits_token_ids is not None: + new_model_output.logits_token_ids = new_model_output.logits_token_ids[0:origin_batch_size] + new_model_output.logits_logsumexp = new_model_output.logits_logsumexp[0:origin_batch_size] new_model_output.mtp_collector = model_output.mtp_collector.unpad_decode( padded_batch_size=padded_batch_size, origin_batch_size=origin_batch_size, @@ -546,6 +616,9 @@ def _create_unpad_prefill_model_output( new_model_output = copy.copy(padded_model_output) # logits 始终只对应每个请求最后一个位置，移除 padding 的 req 对应的行。 new_model_output.logits = new_model_output.logits[0:origin_batch_size] + if new_model_output.logits_token_ids is not None: + new_model_output.logits_token_ids = new_model_output.logits_token_ids[0:origin_batch_size] + new_model_output.logits_logsumexp = new_model_output.logits_logsumexp[0:origin_batch_size] new_model_output.mtp_collector = padded_model_output.mtp_collector.unpad_prefill( origin_handle_token_num=origin_handle_token_num ) @@ -579,7 +652,13 @@ def _prefill( if self.args.enable_tpsp_mix_mode: infer_handle_token_num = triton.cdiv(infer_handle_token_num, self.tp_world_size_) * self.tp_world_size_ - if self.prefill_graph is not None and self.prefill_graph.can_run(handle_token_num=infer_handle_token_num): + if self.prefill_graph is not None and self.prefill_graph.can_run( + handle_token_num=infer_handle_token_num, + batch_size=model_input.batch_size, + max_q_seq_len=model_input.max_q_seq_len, + max_kv_seq_len=model_input.max_kv_seq_len, + max_cache_len=model_input.max_cache_len, + ): infer_handle_token_num = self.prefill_graph.find_closest_graph_handle_token_num( handle_token_num=infer_handle_token_num ) @@ -638,14 +717,19 @@ def _decode( # 向上对齐到 TP world size 的整数倍，保证后续切分得到合法 shape。 infer_batch_size = max(1, origin_batch_size) if self.args.enable_tpsp_mix_mode: - infer_batch_size = triton.cdiv(infer_batch_size, self.tp_world_size_) * self.tp_world_size_ + decode_alignment = math.lcm(self.tp_world_size_, self.decode_batch_multiplier) + infer_batch_size = triton.cdiv(infer_batch_size, decode_alignment) * decode_alignment # CUDA Graph 可能继续向上对齐 batch size，并因此加入 seq_len=2 的 # dummy request。先用最终可能出现的 KV 长度判断 graph，再统一 padding 一次。 infer_max_kv_seq_len = max(2, model_input.max_kv_seq_len) - use_cuda_graph = self.graph is not None and self.graph.can_run( - batch_size=infer_batch_size, - max_len_in_batch=infer_max_kv_seq_len, + use_cuda_graph = ( + self._is_cuda_graph_output_compatible(model_input) + and self.graph is not None + and self.graph.can_run( + batch_size=infer_batch_size, + max_len_in_batch=infer_max_kv_seq_len, + ) ) need_capture = False if use_cuda_graph: @@ -685,7 +769,8 @@ def _context_forward(self, infer_state: InferStateInfo): infer_state.prepare_prefill_dp_balance() input_embs = infer_state._all_to_all_balance_get(data=input_embs) - input_embs = self.pre_infer._tpsp_sp_split(input=input_embs, infer_state=infer_state) + if not self.use_replicated_attention_ep: + input_embs = self.pre_infer._tpsp_sp_split(input=input_embs, infer_state=infer_state) input_tensors = [input_embs] if Autotuner.is_autotune_warmup(): infer_state.hidden_collector = NoopHiddenCollector() @@ -705,7 +790,13 @@ def prefill_func(input_tensors, _infer_state): handle_token_num = infer_state.input_ids.shape[0] - if self.prefill_graph is not None and self.prefill_graph.can_run(handle_token_num=handle_token_num): + if self.prefill_graph is not None and self.prefill_graph.can_run( + handle_token_num=handle_token_num, + batch_size=infer_state.batch_size, + max_q_seq_len=infer_state.max_q_seq_len, + max_kv_seq_len=infer_state.max_kv_seq_len, + max_cache_len=infer_state.max_cache_len, + ): finded_handle_token_num = self.prefill_graph.find_closest_graph_handle_token_num( handle_token_num=handle_token_num ) @@ -728,7 +819,9 @@ def prefill_func(input_tensors, _infer_state): input_embs = output_tensors[0] - last_input_embs = self.post_infer._tpsp_allgather(input=input_embs, infer_state=infer_state) + last_input_embs = input_embs + if not self.use_replicated_attention_ep: + last_input_embs = self.post_infer._tpsp_allgather(input=input_embs, infer_state=infer_state) if infer_state.need_dp_prefill_balance: last_input_embs = infer_state._all_to_all_unbalance_get(data=last_input_embs) @@ -737,6 +830,8 @@ def prefill_func(input_tensors, _infer_state): hidden_collector.add_final_hidden(last_input_embs) model_output = ModelOutput( logits=predict_logits.contiguous(), + logits_token_ids=infer_state.logits_token_ids, + logits_logsumexp=infer_state.logits_logsumexp, mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), prompt_logics=infer_state.prompt_logics, ) @@ -751,14 +846,17 @@ def _token_forward(self, infer_state: InferStateInfo): input_ids = infer_state.input_ids cuda_input_ids = input_ids input_embs = self.pre_infer.token_forward(cuda_input_ids, infer_state, self.pre_post_weight) - input_embs = self.pre_infer._tpsp_sp_split(input=input_embs, infer_state=infer_state) + if not self.use_replicated_attention_ep: + input_embs = self.pre_infer._tpsp_sp_split(input=input_embs, infer_state=infer_state) for i in range(self.layers_num): layer = self.layers_infer[i] input_embs: torch.Tensor = layer.token_forward(input_embs, infer_state, self.trans_layers_weight[i]) hidden_collector.add(layer_index=i, hidden=input_embs) - last_input_embs = self.post_infer._tpsp_allgather(input=input_embs, infer_state=infer_state) + last_input_embs = input_embs + if not self.use_replicated_attention_ep: + last_input_embs = self.post_infer._tpsp_allgather(input=input_embs, infer_state=infer_state) predict_logits: torch.Tensor = self.post_infer.token_forward( last_input_embs, infer_state=infer_state, layer_weight=self.pre_post_weight ) @@ -766,6 +864,8 @@ def _token_forward(self, infer_state: InferStateInfo): hidden_collector.add_final_hidden(last_input_embs) model_output = ModelOutput( logits=predict_logits.contiguous(), + logits_token_ids=infer_state.logits_token_ids, + logits_logsumexp=infer_state.logits_logsumexp, mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), ) @@ -895,9 +995,14 @@ def _microbatch_overlap_decode_cuda(self, model_input0: ModelInput, model_input1 origin_batch_size1 = model_input1.batch_size max_len_in_batch = max(2, model_input0.max_kv_seq_len, model_input1.max_kv_seq_len) infer_batch_size = max(1, origin_batch_size0, origin_batch_size1) - infer_batch_size = triton.cdiv(infer_batch_size, self.tp_world_size_) * self.tp_world_size_ + decode_alignment = math.lcm(self.tp_world_size_, self.decode_batch_multiplier) + infer_batch_size = triton.cdiv(infer_batch_size, decode_alignment) * decode_alignment - if self.graph is not None and self.graph.can_run(infer_batch_size, max_len_in_batch): + if ( + self._is_cuda_graph_output_compatible(model_input0, model_input1) + and self.graph is not None + and self.graph.can_run(infer_batch_size, max_len_in_batch) + ): infer_batch_size = self.graph.find_closest_graph_batch_size(infer_batch_size) need_capture = self.graph.need_capture(infer_batch_size) padded_model_input0 = self._create_padded_decode_model_input(model_input0, infer_batch_size) @@ -1020,11 +1125,15 @@ def _overlap_tpsp_context_forward(self, infer_state: InferStateInfo, infer_state hidden_collector1.add_final_hidden(last_input_embs1) model_output = ModelOutput( logits=predict_logits.contiguous(), + logits_token_ids=infer_state.logits_token_ids, + logits_logsumexp=infer_state.logits_logsumexp, mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), prompt_logics=infer_state.prompt_logics, ) model_output1 = ModelOutput( logits=predict_logits1.contiguous(), + logits_token_ids=infer_state1.logits_token_ids, + logits_logsumexp=infer_state1.logits_logsumexp, mtp_collector=infer_state1.hidden_collector.finish_output(infer_state=infer_state1), prompt_logics=infer_state1.prompt_logics, ) @@ -1069,10 +1178,14 @@ def _overlap_tpsp_token_forward(self, infer_state: InferStateInfo, infer_state1: hidden_collector1.add_final_hidden(last_input_embs1) model_output = ModelOutput( logits=predict_logits.contiguous(), + logits_token_ids=infer_state.logits_token_ids, + logits_logsumexp=infer_state.logits_logsumexp, mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), ) model_output1 = ModelOutput( logits=predict_logits1.contiguous(), + logits_token_ids=infer_state1.logits_token_ids, + logits_logsumexp=infer_state1.logits_logsumexp, mtp_collector=infer_state1.hidden_collector.finish_output(infer_state=infer_state1), ) diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py index ae645d4b7b..cd0b4e75dd 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -20,6 +20,9 @@ class ModelInput: b_req_idx: torch.Tensor = None b_mtp_index: torch.Tensor = None b_seq_len: torch.Tensor = None + # Original prompt length per request/row. This remains on CPU and is used + # only for one-time model-specific fast-path eligibility decisions. + b_input_len: torch.Tensor = None # 在 prefill 阶段，用于在 enable_prefill_decode_mixed 开启下， # 用于标识请求是否为 decode 请求混合在 prefill 请求中。 # 其对应的 input_ids 需要特殊处理, 从 req_to_next_token_ids 中获取。 @@ -55,6 +58,17 @@ class ModelInput: # 的 draft 模型的输入 mtp_draft_input_hiddens: Optional[torch.Tensor] = None + # Whether every request in this prefill batch starts and ends on an + # index-kpool boundary. BaseModel computes this once while the length + # tensors are still on CPU; NSA layers reuse it without a per-layer GPU + # synchronization. + kpool_prefill_aligned: bool = False + kpool_decode_aligned: bool = False + + # The router enables sparse vocabulary output only when target sampling is + # exact, unmodified greedy. Draft models always consume greedy proposals. + use_vocab_parallel_greedy: bool = False + def to_cuda(self): self.check_input() @@ -85,6 +99,11 @@ def __post_init__(self): self.check_input() def check_input(self): + if self.b_input_len is None: + # Internal warmup/padding inputs predate this metadata. Their + # synthetic prompt occupies the whole sequence, so b_seq_len is + # the correct conservative default. + self.b_input_len = self.b_seq_len if self.input_ids is not None: assert ( self.input_ids.dtype == torch.int64 @@ -93,12 +112,14 @@ def check_input(self): assert self.b_req_idx is not None assert self.b_mtp_index is not None assert self.b_seq_len is not None + assert self.b_input_len is not None assert self.multimodal_params is not None assert self.mem_indexes is not None or self.mem_indexes_cpu is not None assert self.b_req_idx.shape == (self.batch_size,) assert self.b_mtp_index.shape == self.b_req_idx.shape assert self.b_seq_len.shape == self.b_req_idx.shape + assert self.b_input_len.shape == self.b_req_idx.shape assert len(self.multimodal_params) == self.batch_size if self.is_prefill: @@ -140,13 +161,15 @@ class ModelMtpOutputCollector: # - 未启用 MTP 时不收集投机特征，该字段同样为 None。 spec_hidden: Optional[torch.Tensor] = None - # DSpark block draft 模型直接生成的 token id，形状通常为 - # [request_count * block_size]。 - # - 仅 DSpark 启用 Markov head（markov_rank > 0）时由 head 直接生成并返回。 - # - DSpark 未启用 Markov head 时为 None，调用方从普通 logits 执行 argmax。 - # - Vanilla MTP、EAGLE、EAGLE3、DFlash 以及未启用 MTP 的模型均不使用该字段。 + # Draft head 直接生成的 token id。Block drafter 通常返回 + # [request_count * block_size]，autoregressive drafter 通常返回 + # [request_count]。未提供时，调用方从普通 logits 执行 argmax。 draft_token_ids: Optional[torch.Tensor] = None + # 与 draft_token_ids 一一对应的精确最大 token 概率。Vocab-parallel + # draft head 可直接返回该值，避免为了动态 verify 聚合完整词表。 + draft_token_probs: Optional[torch.Tensor] = None + # DSpark confidence head 输出的原始置信度 logits，形状通常为 # [request_count, block_size]，供动态 MTP verify 计算各 draft 位置的调度分数。 # - 仅 DSpark checkpoint 启用 confidence head 时返回；动态 verify 模式要求该字段存在。 @@ -159,6 +182,8 @@ def to_no_ref_tensor(self) -> None: self.spec_hidden = tensor_to_no_ref_tensor(self.spec_hidden) if self.draft_token_ids is not None: self.draft_token_ids = tensor_to_no_ref_tensor(self.draft_token_ids) + if self.draft_token_probs is not None: + self.draft_token_probs = tensor_to_no_ref_tensor(self.draft_token_probs) if self.confidence_logits is not None: self.confidence_logits = tensor_to_no_ref_tensor(self.confidence_logits) @@ -168,6 +193,8 @@ def unpad_decode(self, padded_batch_size: int, origin_batch_size: int) -> "Model collector.spec_hidden = collector.spec_hidden[:origin_batch_size] if collector.draft_token_ids is not None: collector.draft_token_ids = collector.draft_token_ids[:origin_batch_size] + if collector.draft_token_probs is not None: + collector.draft_token_probs = collector.draft_token_probs[:origin_batch_size] if collector.confidence_logits is not None: confidence_row_count = collector.confidence_logits.shape[0] assert confidence_row_count > 0 and padded_batch_size % confidence_row_count == 0 @@ -200,10 +227,59 @@ class ModelOutput: # 需要返回 prompt logprobs 信息时才会非空。 prompt_logics: Optional[torch.Tensor] = None + # Vocab-parallel outputs keep logits as logits while mapping each sparse + # column back to its global token id. logits_logsumexp is computed over the + # complete vocabulary, so sparse argmax probabilities remain exact. + # Both fields are None for historical dense logits. + logits_token_ids: Optional[torch.Tensor] = None + logits_logsumexp: Optional[torch.Tensor] = None + def __post_init__(self) -> None: if self.mtp_collector is None: self.mtp_collector = ModelMtpOutputCollector() + assert (self.logits_token_ids is None) == (self.logits_logsumexp is None) + if self.logits_token_ids is not None: + assert self.logits.ndim == 2 + assert self.logits_token_ids.shape == self.logits.shape + assert self.logits_token_ids.dtype in (torch.int32, torch.int64) + assert self.logits_token_ids.device == self.logits.device + assert self.logits_logsumexp.shape == (self.logits.shape[0],) + assert self.logits_logsumexp.dtype == torch.float32 + assert self.logits_logsumexp.device == self.logits.device def to_no_ref_tensor(self): self.logits = tensor_to_no_ref_tensor(self.logits) + if self.logits_token_ids is not None: + self.logits_token_ids = tensor_to_no_ref_tensor(self.logits_token_ids) + self.logits_logsumexp = tensor_to_no_ref_tensor(self.logits_logsumexp) self.mtp_collector.to_no_ref_tensor() + + @property + def has_vocab_parallel_logits(self) -> bool: + return self.logits_token_ids is not None + + def index_select_logits_rows(self, index: torch.Tensor) -> "ModelOutput": + """Select logit rows without dropping their vocabulary metadata.""" + + return ModelOutput( + logits=self.logits.index_select(0, index), + logits_token_ids=( + self.logits_token_ids.index_select(0, index) if self.logits_token_ids is not None else None + ), + logits_logsumexp=( + self.logits_logsumexp.index_select(0, index) if self.logits_logsumexp is not None else None + ), + ) + + @classmethod + def concat_logits_rows(cls, outputs: List["ModelOutput"]) -> "ModelOutput": + """Concatenate compatible dense or sparse-vocabulary logit rows.""" + + assert outputs + sparse = outputs[0].has_vocab_parallel_logits + assert all(output.has_vocab_parallel_logits == sparse for output in outputs) + return cls( + logits=torch.cat([output.logits for output in outputs], dim=0), + logits_token_ids=(torch.cat([output.logits_token_ids for output in outputs], dim=0) if sparse else None), + logits_logsumexp=(torch.cat([output.logits_logsumexp for output in outputs], dim=0) if sparse else None), + ) diff --git a/lightllm/common/basemodel/cuda_graph.py b/lightllm/common/basemodel/cuda_graph.py index 5849cccf54..b2467f7dfc 100644 --- a/lightllm/common/basemodel/cuda_graph.py +++ b/lightllm/common/basemodel/cuda_graph.py @@ -1,14 +1,18 @@ import os +import math import torch import torch.distributed as dist import copy import bisect import triton -from typing import Optional +from typing import Iterable, Optional from lightllm.utils.log_utils import init_logger from lightllm.utils.envs_utils import get_env_start_args from lightllm.distributed import dist_group_manager from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput +from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_greedy import ( + is_vocab_parallel_greedy_enabled, +) from lightllm.utils.torch_memory_saver_utils import ( TorchMemorySaverWrapper, MemoryTag, @@ -46,7 +50,11 @@ def gen_cuda_graph_batch_sizes( batch_sizes = sorted({size for size in batch_sizes if size < max_batch_size} | {max_batch_size}) if args.enable_tpsp_mix_mode: - batch_sizes = sorted({triton.cdiv(size, tp_world_size) * tp_world_size for size in batch_sizes}) + # Static speculative verification stores one request in a fixed + # block of ``batch_step_size_before_split`` rows. TP/SP padding + # must preserve both that block and an even split across TP ranks. + alignment = math.lcm(tp_world_size, batch_step_size_before_split) + batch_sizes = sorted({triton.cdiv(size, alignment) * alignment for size in batch_sizes}) assert batch_sizes[-1] == max_batch_size return batch_sizes @@ -59,6 +67,7 @@ def __init__( max_len_in_batch=8192, tp_world_size: int = 1, capture_infer_cost: bool = False, + extra_batch_sizes: Optional[Iterable[int]] = None, ): self.graph = {} self.tp_world_size = tp_world_size @@ -78,6 +87,11 @@ def __init__( max_batch_size=self.max_batch_size, tp_world_size=self.tp_world_size, ) + if extra_batch_sizes is not None: + self.cuda_graph_batch_sizes = sorted( + set(self.cuda_graph_batch_sizes) + | {int(batch_size) for batch_size in extra_batch_sizes if 0 < int(batch_size) <= self.max_batch_size} + ) logger.info(f"cuda graph batch_sizes: {self.cuda_graph_batch_sizes}") def can_run(self, batch_size, max_len_in_batch): @@ -279,6 +293,7 @@ def warmup(self, model): b_position_delta=torch.zeros(batch_size, dtype=torch.int32, device="cuda"), is_prefill=False, multimodal_params=[{"images": [], "audios": []} for _ in range(batch_size)], + use_vocab_parallel_greedy=is_vocab_parallel_greedy_enabled(), **model._gen_special_model_input(batch_size), ) model_output: ModelOutput = model.forward(model_input) @@ -340,6 +355,7 @@ def warmup_overlap(self, model): b_shared_radix_node_id=b_shared_radix_node_id, b_position_delta=torch.zeros(batch_size, dtype=torch.int32, device="cuda"), multimodal_params=[{"images": [], "audios": []} for _ in range(batch_size)], + use_vocab_parallel_greedy=is_vocab_parallel_greedy_enabled(), **model._gen_special_model_input(batch_size), ) decode_batches.append(micro_batch) diff --git a/lightllm/common/basemodel/hidden_collector.py b/lightllm/common/basemodel/hidden_collector.py index 3eb946fe82..605ed68b59 100644 --- a/lightllm/common/basemodel/hidden_collector.py +++ b/lightllm/common/basemodel/hidden_collector.py @@ -91,6 +91,7 @@ def add_mtp_outputs( self, draft_token_ids: Optional[torch.Tensor], confidence_logits: Optional[torch.Tensor], + draft_token_probs: Optional[torch.Tensor] = None, ) -> None: """Collect optional token/confidence outputs produced by an MTP head. @@ -133,6 +134,7 @@ class MtpHeadOutputCollector(NoopHiddenCollector): def __init__(self) -> None: self.draft_token_ids: Optional[torch.Tensor] = None + self.draft_token_probs: Optional[torch.Tensor] = None self.confidence_logits: Optional[torch.Tensor] = None def new_instance(self) -> HiddenCollector: @@ -142,16 +144,20 @@ def add_mtp_outputs( self, draft_token_ids: Optional[torch.Tensor], confidence_logits: Optional[torch.Tensor], + draft_token_probs: Optional[torch.Tensor] = None, ) -> None: self.draft_token_ids = draft_token_ids + self.draft_token_probs = draft_token_probs self.confidence_logits = confidence_logits def finish_output(self, infer_state) -> ModelMtpOutputCollector: output = ModelMtpOutputCollector( draft_token_ids=self.draft_token_ids, + draft_token_probs=self.draft_token_probs, confidence_logits=self.confidence_logits, ) self.draft_token_ids = None + self.draft_token_probs = None self.confidence_logits = None return output diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index 29648aa78e..f81275279a 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -45,6 +45,8 @@ def __init__(self): # max_cache_len 用于 prefill 阶段标识请求中最大 cache的kv 的长度 self.max_cache_len: int = None self.is_prefill: bool = None + self.kpool_prefill_aligned: bool = False + self.kpool_decode_aligned: bool = False self.mem_manager: MemoryManager = None self.req_manager: ReqManager = None @@ -52,6 +54,9 @@ def __init__(self): self.mem_index: torch.Tensor = None self.return_all_prompt_logics: bool = False + self.use_vocab_parallel_greedy: bool = False + self.logits_token_ids: Optional[torch.Tensor] = None + self.logits_logsumexp: Optional[torch.Tensor] = None # 在开启 return_all_prompt_logics 模式时，保存整个 prefill 阶段每一个 # token 位置的 logits，供后续回传 prompt logprobs 信息使用。 # 仅在 prefill 阶段且需要返回 prompt logprobs 时才会被填充。 @@ -394,4 +399,8 @@ def copy_for_prefill_cuda_graph(self, new_infer_state: "InferStateInfo"): attr_ = getattr(self, attr_name, None) if attr_ is not None and attr_.data_ptr() != attr_value.data_ptr() and attr_.shape == attr_value.shape: attr_.copy_(attr_value, non_blocking=True) + + self.prefill_att_state.copy_for_prefill_cuda_graph(new_infer_state.prefill_att_state) + if self.prefill_att_state1 is not None: + self.prefill_att_state1.copy_for_prefill_cuda_graph(new_infer_state.prefill_att_state1) return diff --git a/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py b/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py index f0cc129c09..da57c8f81b 100755 --- a/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py +++ b/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py @@ -6,6 +6,7 @@ from lightllm.distributed import all_reduce from typing import Tuple from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor +from lightllm.utils.envs_utils import get_env_start_args class TransformerLayerInferTpl(TransformerLayerInfer): @@ -102,6 +103,13 @@ def _context_attention_wrapper_run( self, q: torch.Tensor, cache_kv: torch.Tensor, infer_state: InferStateInfo, layer_weight ) -> torch.Tensor: if torch.cuda.is_current_stream_capturing(): + # Exact-layout graphs can opt into capturing attention as part of + # the main graph. This avoids one CPU callback and one graph split + # per attention layer. The fixed layout is validated by + # PrefillCudaGraph before capture starts. + if bool(getattr(get_env_start_args(), "prefill_cudagraph_capture_attention", False)): + return self._context_attention_kernel(q, cache_kv, infer_state, layer_weight) + q = q.contiguous() # cache_kv is None for layers that own no K/V slot (e.g. gemma4 # KV-shared layers, which read K/V from a prior layer's cache and @@ -110,6 +118,24 @@ def _context_attention_wrapper_run( cache_kv = cache_kv.contiguous() if cache_kv is not None else None _q = tensor_to_no_ref_tensor(q) _cache_kv = tensor_to_no_ref_tensor(cache_kv) if cache_kv is not None else None + + # Some attention implementations stash graph-produced tensors in + # a short-lived infer-state dict between QKV projection and the + # CPU attention callback (for example NSA indexer inputs). Python + # attribute assignment is not replayed by CUDA Graph, and the + # shape-probing call below may consume/delete the dict. Preserve + # fixed-address, non-owning views and restore them for every CPU + # callback invocation. + callback_tensor_dicts = {} + for attr_name, attr_value in vars(infer_state).items(): + if ( + isinstance(attr_value, dict) + and attr_value + and all(isinstance(value, torch.Tensor) for value in attr_value.values()) + ): + callback_tensor_dicts[attr_name] = { + key: tensor_to_no_ref_tensor(value.contiguous()) for key, value in attr_value.items() + } pre_capture_graph = infer_state.prefill_cuda_graph_get_current_capture_graph() pre_capture_graph.__exit__(None, None, None) @@ -135,6 +161,8 @@ def get_o_shape_dtype_device(): _o = tensor_to_no_ref_tensor(o) def att_func(new_infer_state: InferStateInfo): + for attr_name, attr_value in callback_tensor_dicts.items(): + setattr(new_infer_state, attr_name, attr_value) tmp_o = self._context_attention_kernel(_q, _cache_kv, new_infer_state, layer_weight) assert tmp_o.shape == _o.shape _o.copy_(tmp_o) diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py index 7f369c4fd8..367e424a22 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py @@ -68,6 +68,9 @@ def __init__( routed_expert_counter_tensor=self.routed_expert_counter_tensor, auto_update_redundancy_expert=self.auto_update_redundancy_expert, ) + self.fuse_moe_impl.swiglu_limit = self.swiglu_limit + self.fuse_moe_impl.swiglu_alpha = self.swiglu_alpha + self.fuse_moe_impl.swiglu_clamp_up_add_one = self.swiglu_clamp_up_add_one self.lock = threading.Lock() self._create_weight() @@ -79,6 +82,9 @@ def _init_config(self, network_config: Dict[str, Any]): self.num_experts_per_tok = network_config["num_experts_per_tok"] self.routed_scaling_factor = network_config.get("routed_scaling_factor", 1.0) self.scoring_func = network_config.get("scoring_func", "softmax") + self.swiglu_limit = network_config.get("swiglu_limit") + self.swiglu_alpha = network_config.get("swiglu_alpha", 1.0) + self.swiglu_clamp_up_add_one = network_config.get("swiglu_clamp_up_add_one", True) def _init_redundancy_expert_params(self): self.redundancy_expert_num = get_redundancy_expert_num() @@ -285,14 +291,36 @@ def load_hf_weights(self, weights): self._load_weight(self.redundancy_expert_idx_to_local_idx, weights) def verify_load(self): - weight_load_ok = all(all(_weight_pack.load_ok) for _weight_pack in self.w1_list + self.w2_list + self.w3_list) + if getattr(self, "_sm90_mega_moe_weights_prepared", False): + weight_load_ok = all(all(_weight_pack.load_ok) for _weight_pack in self.w2_list) + else: + weight_load_ok = all( + all(_weight_pack.load_ok) for _weight_pack in self.w1_list + self.w2_list + self.w3_list + ) per_expert_scale_load_ok = ( True if self.per_expert_scale is None else getattr(self.per_expert_scale, "load_ok", False) ) e_score_correction_bias_load_ok = ( True if self.e_score_correction_bias is None else getattr(self.e_score_correction_bias, "load_ok", False) ) - return weight_load_ok and per_expert_scale_load_ok and e_score_correction_bias_load_ok + load_ok = weight_load_ok and per_expert_scale_load_ok and e_score_correction_bias_load_ok + if load_ok and self.enable_ep_moe and not getattr(self, "_sm90_mega_moe_weights_prepared", False): + from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import ( + prepare_sm90_mega_moe_weights, + use_sm90_mega_moe, + ) + + if use_sm90_mega_moe(self.quant_method): + prepare_sm90_mega_moe_weights(self.w13) + # The loader-only gate/up views retain the original storage. + # Drop them after replacement so all layers do not keep a + # second full copy of their routed-expert L1 weights. + self.w1 = None + self.w3 = None + self.w1_list = [] + self.w3_list = [] + self._sm90_mega_moe_weights_prepared = True + return load_ok def _create_weight(self): intermediate_size = self.split_inter_size diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py index a5ba656c9c..1878b4071e 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py @@ -12,7 +12,9 @@ get_ep_num_sms, masked_group_gemm, chunked_expanded_moe_forward, + legacy_normal_moe_forward, quantize_fused_experts_input, + use_sm90_mega_moe, ) from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd from lightllm.common.triton_utils.autotuner import Autotuner @@ -77,16 +79,20 @@ def _fused_experts( router_logits: Optional[torch.Tensor] = None, is_prefill: Optional[bool] = None, ): + fused_topk_ids = topk_ids if use_sm90_mega_moe(self.quant_method) else topk_ids.to(torch.long) output = fused_experts( hidden_states=input_tensor, w13=w13, w2=w2, topk_weights=topk_weights, - topk_idx=topk_ids.to(torch.long), + topk_idx=fused_topk_ids, num_experts=self.total_expert_num_contain_redundancy, # number of all experts contain redundancy quant_method=self.quant_method, is_prefill=is_prefill, previous_event=None, # for overlap + swiglu_limit=self.swiglu_limit, + swiglu_alpha=self.swiglu_alpha, + swiglu_clamp_up_add_one=self.swiglu_clamp_up_add_one, ) return output @@ -163,6 +169,42 @@ def dispatch( overlap_event: Optional[Any] = None, ): buffer = dist_group_manager.ep_buffer + if dist_group_manager.ep_prefill_uses_legacy_buffer: + ( + num_tokens_per_rank, + num_tokens_per_rdma_rank, + num_tokens_per_expert, + is_token_in_rank, + layout_event, + ) = buffer.get_dispatch_layout( + topk_idx, + self.total_expert_num_contain_redundancy, + previous_event=overlap_event, + async_finish=False, + allocate_on_comm_stream=False, + ) + (recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, _,) = buffer.dispatch( + qinput_tensor, + topk_idx=topk_idx, + topk_weights=topk_weights, + num_tokens_per_rank=num_tokens_per_rank, + num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, + is_token_in_rank=is_token_in_rank, + num_tokens_per_expert=num_tokens_per_expert, + previous_event=layout_event, + async_finish=False, + allocate_on_comm_stream=False, + expert_alignment=128, + ) + return ( + recv_x, + recv_topk_idx, + recv_topk_weights, + num_recv_tokens_per_expert_list, + handle, + lambda: None, + ) + num_max_tokens_per_rank = get_deepep_num_max_dispatch_tokens_per_rank_prefill() recv_x, recv_topk_idx, recv_topk_weights, handle, event = buffer.dispatch( qinput_tensor, @@ -206,6 +248,9 @@ def masked_group_gemm( w2_weight, w2_scale, expected_m=expected_m, + swiglu_limit=self.swiglu_limit, + swiglu_alpha=self.swiglu_alpha, + swiglu_clamp_up_add_one=self.swiglu_clamp_up_add_one, ) def prefilled_group_gemm( @@ -223,6 +268,22 @@ def prefilled_group_gemm( ): w13_weight, w13_scale = w13.weight, w13.weight_scale w2_weight, w2_scale = w2.weight, w2.weight_scale + if dist_group_manager.ep_prefill_uses_legacy_buffer: + return legacy_normal_moe_forward( + num_recv_tokens_per_expert_list, + recv_x, + recv_topk_idx, + recv_topk_weights, + w13_weight, + w13_scale, + w2_weight, + w2_scale, + hidden_dtype, + self.swiglu_limit, + self.swiglu_alpha, + self.swiglu_clamp_up_add_one, + ) + assert recv_topk_idx is None all_tokens = sum(num_recv_tokens_per_expert_list) if all_tokens > 0: @@ -239,6 +300,9 @@ def prefilled_group_gemm( block_size_k=self.quant_method.block_size, workspace=dist_group_manager.get_deep_ep_prefill_moe_workspace(microbatch_index), hidden_dtype=hidden_dtype, + swiglu_limit=self.swiglu_limit, + swiglu_alpha=self.swiglu_alpha, + swiglu_clamp_up_add_one=self.swiglu_clamp_up_add_one, ) else: gather_out = torch.empty( @@ -254,7 +318,13 @@ def prefilled_group_gemm( N = w13_weight.shape[1] _gemm_out_a = torch.zeros((1, N), device=recv_x[0].device, dtype=hidden_dtype) _silu_out = torch.zeros((1, N // 2), device=recv_x[0].device, dtype=hidden_dtype) - silu_and_mul_fwd(_gemm_out_a.view(-1, N), _silu_out) + silu_and_mul_fwd( + _gemm_out_a.view(-1, N), + _silu_out, + limit=self.swiglu_limit, + alpha=self.swiglu_alpha, + clamp_up_add_one=self.swiglu_clamp_up_add_one, + ) _gemm_out_a, _silu_out = None, None del recv_x return gather_out @@ -278,6 +348,17 @@ def combine( overlap_event: Optional[Any] = None, ): # normal combine + if dist_group_manager.ep_prefill_uses_legacy_buffer: + combined_x, _, _ = dist_group_manager.ep_buffer.combine( + gemm_out_b, + handle, + topk_weights=None, + previous_event=overlap_event, + async_finish=False, + allocate_on_comm_stream=False, + ) + return combined_x, lambda: None + combined_x, _, event = dist_group_manager.ep_buffer.combine( gemm_out_b, handle, diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py index 1d6a38c069..cbeea19baa 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py @@ -1,3 +1,8 @@ +import os +from contextlib import contextmanager, nullcontext +from functools import lru_cache +from types import SimpleNamespace + import torch from typing import Callable, Optional from lightllm.common.quantization.no_quant import WeightPack @@ -5,6 +10,161 @@ from .base_impl import FuseMoeBaseImpl +def _use_sglang_triton_moe() -> bool: + return os.getenv("LIGHTLLM_USE_SGLANG_TRITON_MOE", "0").upper() in { + "ON", + "TRUE", + "1", + } + + +@lru_cache(maxsize=1) +def _get_sglang_fused_experts_impl(): + """Load SGLang's tuned Triton MoE without its server RuntimeContext. + + LightLLM and SGLang share the same standard block-FP8 expert layout, but + SGLang's standalone kernel helpers consult two unrelated process-global + server flags. Supply their disabled defaults locally so the kernel can be + used as an optional backend without initializing an SGLang server. + """ + + from sglang.srt.layers.moe.moe_runner.triton_utils import ( + fused_moe as sglang_fused_moe, + ) + from sglang.srt.layers.moe.moe_runner.triton_utils import ( + fused_moe_triton_config as sglang_fused_moe_config, + ) + + # Despite the upstream flag name, for this in-place standalone invocation + # the optimization only folds the top-k expert sum into the down GEMM. It + # does not perform the TP all-reduce, which remains owned by LightLLM. + enable_fused_sum = os.getenv("LIGHTLLM_SGLANG_FUSED_MOE_SUM", "1").upper() in { + "ON", + "TRUE", + "1", + } + standalone_exec = SimpleNamespace( + moe=SimpleNamespace(enable_fused_moe_sum_all_reduce=enable_fused_sum), + deterministic=SimpleNamespace(enable_deterministic_inference=False), + ) + sglang_fused_moe.get_exec = lambda: standalone_exec + sglang_fused_moe_config.get_exec = lambda: standalone_exec + from sglang.srt.layers.moe.moe_runner.triton_utils import override_config + + return sglang_fused_moe, override_config + + +@contextmanager +def _override_sglang_moe_configs( + sglang_fused_moe, + override_config, + up_config: dict, + down_config: Optional[dict], +): + """Temporarily select independently measured up/down MoE configs.""" + + if down_config is None: + with override_config(up_config): + yield + return + + original = sglang_fused_moe.try_get_optimal_moe_config + + def resolve_config(*args, return_down_config=False, **kwargs): + if return_down_config: + return up_config, (down_config, None) + return up_config + + sglang_fused_moe.try_get_optimal_moe_config = resolve_config + try: + yield + finally: + sglang_fused_moe.try_get_optimal_moe_config = original + + +@lru_cache(maxsize=None) +def _get_sglang_triton_moe_configs( + w13_shape: tuple[int, ...], + w2_shape: tuple[int, ...], + topk: int, + is_prefill: bool, + token_count: int, +): + """Return measured H100 (up, down) configs for GLM-5's TP8 shape.""" + + if "H100" not in torch.cuda.get_device_name(torch.cuda.current_device()): + return None + if is_prefill: + if token_count < 8192: + up_config = { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3, + } + return up_config, None + + # The two GLM expert projections have different output widths and + # diverge at large prefill batches. Measurements on H100 TP8 show + # that N=128 remains best for the 4096->512 up projection, while the + # 256->4096 down projection crosses over to N=64 at roughly 24K + # tokens. Selecting them independently avoids making the faster up + # projection pay for the down projection's narrower tile. + up_config = { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3, + } + down_config = { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64 if token_count >= 24576 else 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32 if token_count >= 24576 else 8, + "num_warps": 4, + "num_stages": 3, + } + return up_config, down_config + if w13_shape != (289, 512, 4096): + return None + if w2_shape != (289, 4096, 256) or topk != 9: + return None + up_config = { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3, + } + # The draft graph has eight physical rows; the MTP2 verify graph has 24. + # Separate down-projection sweeps found different winners for those two + # hot shapes while retaining BLOCK_SIZE_M=16 for shared route alignment. + if token_count <= 8: + down_config = { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 3, + } + else: + down_config = { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + } + return up_config, down_config + + class FuseMoeTriton(FuseMoeBaseImpl): def __init__( self, @@ -92,6 +252,52 @@ def _fused_experts( w2_weight, w2_scale = w2.weight, w2.weight_scale use_fp8_w8a8 = w13_weight.dtype == torch.float8_e4m3fn + if _use_sglang_triton_moe(): + block_size = getattr(self.quant_method, "block_size", None) + if not use_fp8_w8a8 or block_size != 128: + raise RuntimeError( + "LIGHTLLM_USE_SGLANG_TRITON_MOE currently requires " + "block-wise FP8 expert weights with block size 128" + ) + if getattr(self, "swiglu_limit", None) is not None and getattr(self, "swiglu_clamp_up_add_one", True): + raise RuntimeError("SGLang Triton MoE does not support clamp_up_add_one=True") + + sglang_fused_moe, override_config = _get_sglang_fused_experts_impl() + tuned_configs = _get_sglang_triton_moe_configs( + tuple(w13_weight.shape), + tuple(w2_weight.shape), + topk_ids.shape[1], + bool(is_prefill), + input_tensor.shape[0], + ) + config_context = ( + _override_sglang_moe_configs( + sglang_fused_moe, + override_config, + *tuned_configs, + ) + if tuned_configs is not None + else nullcontext() + ) + with config_context: + sglang_fused_moe.fused_experts_impl( + hidden_states=input_tensor, + w1=w13_weight, + w2=w2_weight, + topk_weights=topk_weights, + topk_ids=topk_ids, + inplace=True, + use_fp8_w8a8=True, + w1_scale=w13_scale, + w2_scale=w2_scale, + block_shape=[block_size, block_size], + routed_scaling_factor=1.0, + filter_expert=False, + swiglu_limit=getattr(self, "swiglu_limit", None), + gate_up_interleaved=False, + ) + return input_tensor + from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe import fused_experts fused_experts( @@ -104,6 +310,9 @@ def _fused_experts( use_fp8_w8a8=use_fp8_w8a8, w1_scale=w13_scale, w2_scale=w2_scale, + limit=getattr(self, "swiglu_limit", None), + alpha=(getattr(self, "swiglu_alpha", 1.0) if getattr(self, "swiglu_limit", None) is not None else None), + clamp_up_add_one=getattr(self, "swiglu_clamp_up_add_one", True), ) return input_tensor diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/norm_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/norm_weight.py index ee9d1923c3..fe80974569 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/norm_weight.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/norm_weight.py @@ -73,6 +73,17 @@ def __call__( class GatedRMSNormWeight(RMSNormWeight): + def __init__( + self, + dim: int, + weight_name: str, + data_type: torch.dtype, + activation: str = "silu", + ): + super().__init__(dim=dim, weight_name=weight_name, data_type=data_type) + assert activation in ("silu", "sigmoid") + self.activation = activation + def _triton_forward( self, input: torch.Tensor, @@ -86,7 +97,15 @@ def _triton_forward( ), f"input.ndim: {input.ndim} != 2 or weight.ndim: {self.weight.ndim} != 1" if out is None: out = alloc_func(input.shape, dtype=input.dtype, device=input.device) - return gated_rmsnorm_forward(x=input, weight=self.weight, bias=None, eps=eps, z=gate_value, out=out) + return gated_rmsnorm_forward( + x=input, + weight=self.weight, + bias=None, + eps=eps, + z=gate_value, + out=out, + activation=self.activation, + ) def _cuda_forward( self, diff --git a/lightllm/common/basemodel/mtp_manager.py b/lightllm/common/basemodel/mtp_manager.py index be6c477b99..16bcdbebb5 100644 --- a/lightllm/common/basemodel/mtp_manager.py +++ b/lightllm/common/basemodel/mtp_manager.py @@ -16,6 +16,7 @@ class MtpManager: _instance: ClassVar[Optional["MtpManager"]] = None _CHAINED_DRAFT_MODES = ("vanilla_with_att", "vanilla_no_att") _RECURRENT_DRAFT_MODES = ("eagle_with_att", "eagle_no_att", "eagle3") + _RECURRENT_ATTN_DRAFT_MODES = ("eagle_with_att", "eagle3") _BLOCK_DRAFT_MODES = ("dspark", "dflash") @classmethod @@ -41,7 +42,10 @@ def get_decode_batch_multiplier(self, is_draft_model: bool) -> int: if not is_draft_model: return verify_width - # Chained MTP runs every draft module over the expanded verify layout. + # Chained MTP runs every draft module over the expanded verify layout, + # but each physical row is still an independent one-token draft decode + # from the attention backend's point of view. CUDA Graph sizing must + # account for the wider physical batch separately (see below). if spec_mode in self._CHAINED_DRAFT_MODES: return 1 @@ -55,12 +59,39 @@ def get_decode_batch_multiplier(self, is_draft_model: bool) -> int: return 1 + def get_decode_cuda_graph_batch_multiplier(self, is_draft_model: bool) -> int: + """Return physical decode rows per logical request for graph sizing. + + Chained draft models keep normal one-token attention semantics, while + their proposer forwards the complete target verification layout through + every draft depth. Consequently a logical batch of ``N`` requests has + ``N * (mtp_step + 1)`` physical rows and needs graphs captured at that + width even though :meth:`get_decode_batch_multiplier` returns one. + """ + + # Attention-backed recurrent EAGLE normally decodes one row per + # request, but its first forward after every target verification runs + # over the complete expanded verify layout to commit draft KV. That + # extend forward must be graph-covered as well. + if is_draft_model and self.args.mtp_mode in ( + *self._CHAINED_DRAFT_MODES, + *self._RECURRENT_ATTN_DRAFT_MODES, + ): + return self.args.mtp_step + 1 + return self.get_decode_batch_multiplier(is_draft_model=is_draft_model) + + def draft_model_needs_logical_batch_graphs(self, is_draft_model: bool) -> bool: + """Whether a widened draft graph also needs one-row-per-request sizes.""" + + return is_draft_model and self.args.mtp_mode in self._RECURRENT_ATTN_DRAFT_MODES + def get_decode_cuda_graph_grow_step_size(self, is_draft_model: bool) -> int: """Return the batch-size stride used to capture decode CUDA Graphs.""" - # Draft model CUDA Graphs follow the drafter's physical decode layout. + # Draft model CUDA Graphs follow the drafter's physical forward layout, + # which is wider than its attention semantics for chained MTP. if is_draft_model: - return self.get_decode_batch_multiplier(is_draft_model=True) + return self.get_decode_cuda_graph_batch_multiplier(is_draft_model=True) # Main model CUDA Graphs use unit growth for dynamically compacted verify rows. else: if self.args.mtp_dynamic_verify: diff --git a/lightllm/common/basemodel/prefill_cuda_graph.py b/lightllm/common/basemodel/prefill_cuda_graph.py index bf6039a48f..244bfd0114 100644 --- a/lightllm/common/basemodel/prefill_cuda_graph.py +++ b/lightllm/common/basemodel/prefill_cuda_graph.py @@ -10,6 +10,9 @@ from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor from lightllm.distributed import dist_group_manager from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput +from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_greedy import ( + is_vocab_parallel_greedy_enabled, +) from .infer_struct import InferStateInfo from .cuda_graph import CudaGraph @@ -33,16 +36,64 @@ def __init__(self, decode_cuda_graph: CudaGraph, tp_world_size: int): if self.args.batch_max_tokens is not None: self.max_handle_token_num = min(self.max_handle_token_num, self.args.batch_max_tokens) - graph_handle_token_nums = ( - list(range(4, 33, 4)) - + list(range(48, 257, 16)) - + list(range(288, 513, 32)) - + list(range(576, 1024 + 1, 64)) - + list(range(1280, 4096 + 1, 256)) - + list(range(4608, self.max_handle_token_num + 1, 512)) - ) - graph_handle_token_nums = [e for e in graph_handle_token_nums if e <= self.max_handle_token_num] - graph_handle_token_nums.append(self.max_handle_token_num) + configured_token_nums = getattr(self.args, "prefill_cudagraph_token_nums", None) + configured_batch_sizes = getattr(self.args, "prefill_cudagraph_batch_sizes", None) + self.use_exact_token_nums = configured_token_nums is not None + self.capture_attention = bool(getattr(self.args, "prefill_cudagraph_capture_attention", False)) + if self.capture_attention and not self.use_exact_token_nums: + raise ValueError( + "--prefill_cudagraph_capture_attention requires " + "--prefill_cudagraph_token_nums and --prefill_cudagraph_batch_sizes" + ) + self.exact_batch_size_by_token_num = {} + if self.use_exact_token_nums: + if configured_batch_sizes is None: + raise ValueError("--prefill_cudagraph_batch_sizes is required with " "--prefill_cudagraph_token_nums") + if len(configured_token_nums) != len(configured_batch_sizes): + raise ValueError( + "--prefill_cudagraph_token_nums and --prefill_cudagraph_batch_sizes " + "must contain the same number of entries" + ) + + graph_handle_token_nums = [] + for raw_token_num, raw_batch_size in zip(configured_token_nums, configured_batch_sizes): + token_num = int(raw_token_num) + batch_size = int(raw_batch_size) + if not 0 < token_num <= self.max_handle_token_num: + continue + if batch_size <= 0: + raise ValueError("--prefill_cudagraph_batch_sizes entries must be positive") + if token_num % batch_size != 0: + raise ValueError( + "Each --prefill_cudagraph_token_nums entry must be divisible by its paired " + "--prefill_cudagraph_batch_sizes entry" + ) + old_batch_size = self.exact_batch_size_by_token_num.get(token_num) + if old_batch_size is not None and old_batch_size != batch_size: + raise ValueError( + f"Conflicting prefill CUDA Graph batch sizes for token count {token_num}: " + f"{old_batch_size} and {batch_size}" + ) + self.exact_batch_size_by_token_num[token_num] = batch_size + graph_handle_token_nums.append(token_num) + if not graph_handle_token_nums: + raise ValueError( + "--prefill_cudagraph_token_nums must contain a positive value no larger than " + "--prefill_cudagraph_max_handle_token and --batch_max_tokens" + ) + else: + if configured_batch_sizes is not None: + raise ValueError("--prefill_cudagraph_batch_sizes requires --prefill_cudagraph_token_nums") + graph_handle_token_nums = ( + list(range(4, 33, 4)) + + list(range(48, 257, 16)) + + list(range(288, 513, 32)) + + list(range(576, 1024 + 1, 64)) + + list(range(1280, 4096 + 1, 256)) + + list(range(4608, self.max_handle_token_num + 1, 512)) + ) + graph_handle_token_nums = [e for e in graph_handle_token_nums if e <= self.max_handle_token_num] + graph_handle_token_nums.append(self.max_handle_token_num) graph_handle_token_nums = list(set[int](graph_handle_token_nums)) graph_handle_token_nums.sort() @@ -55,8 +106,25 @@ def __init__(self, decode_cuda_graph: CudaGraph, tp_world_size: int): self.graph_handle_token_nums = graph_handle_token_nums logger.info(f"prefill cuda graph graph_handle_token_nums: {self.graph_handle_token_nums}") + if self.exact_batch_size_by_token_num: + logger.info( + "prefill cuda graph exact layouts (token_num -> batch_size): " f"{self.exact_batch_size_by_token_num}" + ) - def can_run(self, handle_token_num: int): + def can_run( + self, + handle_token_num: int, + batch_size: Optional[int] = None, + max_q_seq_len: Optional[int] = None, + max_kv_seq_len: Optional[int] = None, + max_cache_len: Optional[int] = None, + ): + if self.use_exact_token_nums: + configured_batch_size = self.exact_batch_size_by_token_num.get(handle_token_num) + if configured_batch_size is None or batch_size != configured_batch_size: + return False + uniform_seq_len = handle_token_num // configured_batch_size + return max_q_seq_len == uniform_seq_len and max_kv_seq_len == uniform_seq_len and max_cache_len == 0 return handle_token_num <= self.max_handle_token_num def need_capture(self, handle_token_num: int): @@ -191,23 +259,31 @@ def warmup(self, model): # prefill cuda graph init for handle_token_num in self.graph_handle_token_nums[::-1]: - logger.info(f"Capture prefill cudagraph, handle_token_num: {handle_token_num}") + batch_size = self.exact_batch_size_by_token_num.get(handle_token_num, 1) + if handle_token_num % batch_size != 0: + raise ValueError( + f"Prefill CUDA Graph token count {handle_token_num} is not divisible by batch size {batch_size}" + ) + seq_len = handle_token_num // batch_size + logger.info( + "Capture prefill cudagraph, " + f"handle_token_num: {handle_token_num}, batch_size: {batch_size}, seq_len: {seq_len}" + ) total_token_num = handle_token_num input_ids = torch.tensor([1 for _ in range(total_token_num)], dtype=torch.int64, device="cuda") mem_indexes = model.mem_manager.alloc(len(input_ids)).cuda() - b_req_idx = torch.tensor([model.req_manager.HOLD_REQUEST_ID], dtype=torch.int32, device="cuda") - b_seq_len = torch.empty(1, dtype=torch.int32, device="cuda") - b_seq_len.fill_(total_token_num) - b_mtp_index = torch.zeros(1, dtype=torch.int32, device="cuda") - b_is_decode_req = torch.zeros(1, dtype=torch.bool, device="cuda") - b_ready_cache_len = torch.zeros(1, dtype=torch.int32, device="cuda") - b_prefill_start_loc = torch.zeros(1, dtype=torch.int32, device="cuda") + b_req_idx = torch.full((batch_size,), model.req_manager.HOLD_REQUEST_ID, dtype=torch.int32, device="cuda") + b_seq_len = torch.full((batch_size,), seq_len, dtype=torch.int32, device="cuda") + b_mtp_index = torch.zeros(batch_size, dtype=torch.int32, device="cuda") + b_is_decode_req = torch.zeros(batch_size, dtype=torch.bool, device="cuda") + b_ready_cache_len = torch.zeros(batch_size, dtype=torch.int32, device="cuda") + b_prefill_start_loc = torch.arange(0, total_token_num, seq_len, dtype=torch.int32, device="cuda") model_input = ModelInput( - batch_size=1, + batch_size=batch_size, total_token_num=total_token_num, - max_q_seq_len=total_token_num, - max_kv_seq_len=total_token_num, + max_q_seq_len=seq_len, + max_kv_seq_len=seq_len, max_cache_len=0, input_ids=input_ids, mem_indexes=mem_indexes, @@ -218,8 +294,9 @@ def warmup(self, model): b_ready_cache_len=b_ready_cache_len, b_prefill_start_loc=b_prefill_start_loc, is_prefill=True, - b_prefill_has_output_cpu=[False], - multimodal_params=[{"images": [], "audios": []}], + b_prefill_has_output_cpu=[False for _ in range(batch_size)], + multimodal_params=[{"images": [], "audios": []} for _ in range(batch_size)], + use_vocab_parallel_greedy=is_vocab_parallel_greedy_enabled(), **model._gen_special_model_input(token_num=total_token_num), ) model_output: ModelOutput = model.forward(model_input) @@ -281,6 +358,7 @@ def warmup_overlap(self, model): is_prefill=True, b_prefill_has_output_cpu=[False], multimodal_params=[{"images": [], "audios": []}], + use_vocab_parallel_greedy=is_vocab_parallel_greedy_enabled(), **model._gen_special_model_input(token_num=total_token_num), ) diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_legacy_layout.py b/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_legacy_layout.py new file mode 100644 index 0000000000..0a40488422 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_legacy_layout.py @@ -0,0 +1,253 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _ep_scatter_offsets( + padded_tokens_per_expert, + valid_tokens_per_expert, + expert_start_loc, + m_indices, + num_experts: tl.constexpr, + block_experts: tl.constexpr, + block_tokens: tl.constexpr, +): + expert_id = tl.program_id(0) + expert_offsets = tl.arange(0, block_experts) + counts = tl.load( + padded_tokens_per_expert + expert_offsets, + mask=expert_offsets < num_experts, + other=0, + ) + starts = tl.cumsum(counts) - counts + tl.store(expert_start_loc + expert_offsets, starts, mask=expert_offsets < num_experts) + + expert_start = tl.load(expert_start_loc + expert_id) + padded_count = tl.load(padded_tokens_per_expert + expert_id) + valid_count = tl.load(valid_tokens_per_expert + expert_id) + offsets = tl.arange(0, block_tokens) + for token_start in tl.range(0, padded_count, block_tokens, num_stages=4): + token_offsets = token_start + offsets + tl.store( + m_indices + expert_start + token_offsets, + tl.where(token_offsets < valid_count, expert_id, -1), + ) + + +@triton.jit +def _ep_scatter_tokens( + total_token_num, + expert_start_loc, + recv_x, + recv_x_stride0, + recv_x_stride1, + recv_x_scale, + recv_x_scale_stride0, + recv_x_scale_stride1, + recv_topk, + recv_topk_stride0, + recv_topk_stride1, + output_tensor, + output_tensor_stride0, + output_tensor_stride1, + output_tensor_scale, + output_tensor_scale_stride0, + output_tensor_scale_stride1, + output_index, + output_index_stride0, + output_index_stride1, + topk_num: tl.constexpr, + hidden_size: tl.constexpr, + hidden_size_pad: tl.constexpr, + scale_hidden_size: tl.constexpr, + scale_hidden_size_pad: tl.constexpr, +): + start_token_id = tl.program_id(0) + grid_size = tl.num_programs(0) + hidden_offsets = tl.arange(0, hidden_size_pad) + hidden_mask = hidden_offsets < hidden_size + scale_offsets = tl.arange(0, scale_hidden_size_pad) + scale_mask = scale_offsets < scale_hidden_size + + for token_id_int32 in range(start_token_id, total_token_num, grid_size): + token_id = token_id_int32.to(tl.int64) + token = tl.load( + recv_x + token_id * recv_x_stride0 + hidden_offsets * recv_x_stride1, + mask=hidden_mask, + ) + token_scale = tl.load( + recv_x_scale + token_id * recv_x_scale_stride0 + scale_offsets * recv_x_scale_stride1, + mask=scale_mask, + ) + for topk_offset_int32 in tl.range(0, topk_num, 1, num_stages=4): + topk_offset = topk_offset_int32.to(tl.int64) + expert_id = tl.load(recv_topk + token_id * recv_topk_stride0 + topk_offset * recv_topk_stride1) + if expert_id >= 0: + destination_int32 = tl.atomic_add(expert_start_loc + expert_id, 1) + destination = destination_int32.to(tl.int64) + tl.store( + output_index + token_id * output_index_stride0 + topk_offset * output_index_stride1, + destination_int32, + ) + tl.store( + output_tensor + destination * output_tensor_stride0 + hidden_offsets * output_tensor_stride1, + token, + mask=hidden_mask, + ) + tl.store( + output_tensor_scale + + destination * output_tensor_scale_stride0 + + scale_offsets * output_tensor_scale_stride1, + token_scale, + mask=scale_mask, + ) + + +@torch.no_grad() +def ep_scatter( + recv_x: torch.Tensor, + recv_x_scale: torch.Tensor, + recv_topk: torch.Tensor, + padded_tokens_per_expert: torch.Tensor, + valid_tokens_per_expert: torch.Tensor, + expert_start_loc: torch.Tensor, + output_tensor: torch.Tensor, + output_tensor_scale: torch.Tensor, + m_indices: torch.Tensor, + output_index: torch.Tensor, +): + block_tokens = 128 + num_experts = padded_tokens_per_expert.shape[0] + hidden_size = recv_x.shape[1] + scale_hidden_size = recv_x_scale.shape[1] + assert m_indices.shape[0] % block_tokens == 0 + + _ep_scatter_offsets[(num_experts,)]( + padded_tokens_per_expert, + valid_tokens_per_expert, + expert_start_loc, + m_indices, + num_experts=num_experts, + block_experts=triton.next_power_of_2(num_experts), + block_tokens=block_tokens, + num_warps=8, + ) + _ep_scatter_tokens[(min(recv_topk.shape[0], 8192),)]( + recv_topk.shape[0], + expert_start_loc, + recv_x, + recv_x.stride(0), + recv_x.stride(1), + recv_x_scale, + recv_x_scale.stride(0), + recv_x_scale.stride(1), + recv_topk, + recv_topk.stride(0), + recv_topk.stride(1), + output_tensor, + output_tensor.stride(0), + output_tensor.stride(1), + output_tensor_scale, + output_tensor_scale.stride(0), + output_tensor_scale.stride(1), + output_index, + output_index.stride(0), + output_index.stride(1), + topk_num=recv_topk.shape[1], + hidden_size=hidden_size, + hidden_size_pad=triton.next_power_of_2(hidden_size), + scale_hidden_size=scale_hidden_size, + scale_hidden_size_pad=triton.next_power_of_2(scale_hidden_size), + num_warps=8, + ) + + +@triton.jit +def _ep_gather_kernel( + total_token_num, + input_tensor, + input_tensor_stride0, + input_tensor_stride1, + recv_topk_ids, + recv_topk_ids_stride0, + recv_topk_ids_stride1, + recv_topk_weights, + recv_topk_weights_stride0, + recv_topk_weights_stride1, + input_index, + input_index_stride0, + input_index_stride1, + output_tensor, + output_tensor_stride0, + output_tensor_stride1, + topk_num: tl.constexpr, + block_hidden: tl.constexpr, +): + hidden_block_int32 = tl.program_id(0) + hidden_block = hidden_block_int32.to(tl.int64) + start_token_int32 = tl.program_id(1) + grid_size = tl.num_programs(1) + hidden_offsets = tl.arange(0, block_hidden) + + for token_int32 in range(start_token_int32, total_token_num, grid_size): + token = token_int32.to(tl.int64) + accumulator = tl.zeros([block_hidden], dtype=tl.float32) + for topk_offset_int32 in range(0, topk_num): + topk_offset = topk_offset_int32.to(tl.int64) + expert_id = tl.load(recv_topk_ids + token * recv_topk_ids_stride0 + topk_offset * recv_topk_ids_stride1) + if expert_id >= 0: + source_int32 = tl.load(input_index + token * input_index_stride0 + topk_offset * input_index_stride1) + source = source_int32.to(tl.int64) + weight = tl.load( + recv_topk_weights + token * recv_topk_weights_stride0 + topk_offset * recv_topk_weights_stride1 + ) + value = tl.load( + input_tensor + + source * input_tensor_stride0 + + hidden_block * block_hidden + + hidden_offsets * input_tensor_stride1 + ) + accumulator += value.to(tl.float32) * weight + tl.store( + output_tensor + + token * output_tensor_stride0 + + hidden_block * block_hidden + + hidden_offsets * output_tensor_stride1, + accumulator.to(output_tensor.dtype.element_ty), + ) + + +@torch.no_grad() +def ep_gather( + input_tensor: torch.Tensor, + recv_topk_ids: torch.Tensor, + recv_topk_weights: torch.Tensor, + input_index: torch.Tensor, + output_tensor: torch.Tensor, +): + hidden_size = input_tensor.shape[1] + block_hidden = 1024 if hidden_size % 1024 == 0 else 128 + assert hidden_size % block_hidden == 0 + grid = (triton.cdiv(hidden_size, block_hidden), min(output_tensor.shape[0], 1024)) + _ep_gather_kernel[grid]( + output_tensor.shape[0], + input_tensor, + input_tensor.stride(0), + input_tensor.stride(1), + recv_topk_ids, + recv_topk_ids.stride(0), + recv_topk_ids.stride(1), + recv_topk_weights, + recv_topk_weights.stride(0), + recv_topk_weights.stride(1), + input_index, + input_index.stride(0), + input_index.stride(1), + output_tensor, + output_tensor.stride(0), + output_tensor.stride(1), + topk_num=recv_topk_ids.shape[1], + block_hidden=block_hidden, + num_warps=2, + ) diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py index e10adf7758..ed3bf33cfe 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py @@ -1009,6 +1009,7 @@ def fused_experts_impl( layout="blocked", limit=None, alpha=None, + clamp_up_add_one=True, ): # Check constraints. assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch" @@ -1087,6 +1088,7 @@ def fused_experts_impl( intermediate_cache2.view(-1, N // 2), limit=limit, alpha=alpha, + clamp_up_add_one=clamp_up_add_one, layout=layout, ) @@ -1133,6 +1135,7 @@ def inplace_fused_experts_impl( layout: str = "blocked", alpha: Optional[float] = None, limit: Optional[float] = None, + clamp_up_add_one: bool = True, ) -> None: fused_experts_impl( hidden_states, @@ -1152,6 +1155,7 @@ def inplace_fused_experts_impl( layout=layout, alpha=alpha, limit=limit, + clamp_up_add_one=clamp_up_add_one, ) @@ -1173,6 +1177,7 @@ def inplace_fused_experts_impl_fake( layout: str = "blocked", alpha: Optional[float] = None, limit: Optional[float] = None, + clamp_up_add_one: bool = True, ) -> None: pass @@ -1203,6 +1208,7 @@ def outplace_fused_experts_impl( layout: str = "blocked", alpha: Optional[float] = None, limit: Optional[float] = None, + clamp_up_add_one: bool = True, ) -> None: return fused_experts_impl( hidden_states, @@ -1222,6 +1228,7 @@ def outplace_fused_experts_impl( layout=layout, alpha=alpha, limit=limit, + clamp_up_add_one=clamp_up_add_one, ) @@ -1243,6 +1250,7 @@ def outplace_fused_experts_impl_fake( layout: str = "blocked", alpha: Optional[float] = None, limit: Optional[float] = None, + clamp_up_add_one: bool = True, ) -> None: return torch.empty_like(hidden_states) @@ -1274,6 +1282,7 @@ def fused_experts( layout: str = "blocked", alpha: Optional[float] = None, limit: Optional[float] = None, + clamp_up_add_one: bool = True, ): if inplace: torch.ops.lightllm.inplace_fused_experts_impl( @@ -1293,6 +1302,7 @@ def fused_experts( layout=layout, alpha=alpha, limit=limit, + clamp_up_add_one=clamp_up_add_one, ) return hidden_states else: @@ -1313,4 +1323,5 @@ def fused_experts( layout=layout, alpha=alpha, limit=limit, + clamp_up_add_one=clamp_up_add_one, ) diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py index 58d4d45514..1eef8114f0 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py @@ -1,5 +1,7 @@ """Fused MoE kernel.""" +import os + import torch import triton import triton.language as tl @@ -12,7 +14,9 @@ ) from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import ( per_token_group_quant_fp8, + tma_align_input_scale, ) +from lightllm.common.basemodel.triton_kernel.fused_moe.deepep_legacy_layout import ep_gather, ep_scatter from lightllm.common.basemodel.triton_kernel.fused_moe.deepep_expanded_layout_kernels import ( ep_build_m_indices, ep_compact_metadata, @@ -24,7 +28,7 @@ get_deepep_num_max_dispatch_tokens_per_rank_decode, ) from lightllm.common.triton_utils.autotuner import Autotuner -from lightllm.utils.device_utils import is_sm100_gpu +from lightllm.utils.device_utils import is_sm90_gpu, is_sm100_gpu from lightllm.utils.sgl_utils import HAS_SGL_KERNEL from lightllm.utils.tensor_buffer_manager import TensorBufferManager @@ -51,6 +55,21 @@ def use_sm100_mega_moe(quant_method: Any) -> bool: return is_sm100_gpu() and quant_method.method_name == "fp4fp8-b32-deepgemm" +def use_sm90_mega_moe(quant_method: Any) -> bool: + return ( + is_sm90_gpu() + and os.getenv("LIGHTLLM_ENABLE_SM90_MEGA_MOE", "0").upper() in {"1", "ON", "TRUE"} + and quant_method.method_name == "fp8w8a8-b128-deepgemm" + and HAS_DEEPGEMM + and hasattr(deep_gemm, "fp8_mega_moe") + and hasattr(deep_gemm, "mega_moe_pre_dispatch_sm90") + ) + + +def use_mega_moe(quant_method: Any) -> bool: + return use_sm90_mega_moe(quant_method) or use_sm100_mega_moe(quant_method) + + def check_ep_expert_dtype(quant_method: Any): expert_dtype = getattr(quant_method, "method_name", None) if expert_dtype not in SUPPORTED_EP_EXPERT_DTYPES: @@ -75,6 +94,9 @@ def masked_group_gemm( w2: torch.Tensor, w2_scale: torch.Tensor, expected_m: int, + swiglu_limit=None, + swiglu_alpha=1.0, + swiglu_clamp_up_add_one=True, ): padded_m = recv_x[0].shape[1] E, N, _ = w1.shape @@ -86,7 +108,16 @@ def masked_group_gemm( qsilu_out = torch.empty((E, padded_m, N // 2), dtype=w1.dtype, device=recv_x[0].device) _deepgemm_grouped_fp8_nt_masked(recv_x, (w1, w1_scale), gemm_out_a, masked_m, expected_m) - silu_and_mul_masked_post_quant_fwd(gemm_out_a, qsilu_out, qsilu_out_scale, block_size, masked_m) + silu_and_mul_masked_post_quant_fwd( + gemm_out_a, + qsilu_out, + qsilu_out_scale, + block_size, + masked_m, + limit=swiglu_limit, + alpha=swiglu_alpha, + clamp_up_add_one=swiglu_clamp_up_add_one, + ) del gemm_out_a gemm_out_b = torch.empty_like(recv_x[0], device=recv_x[0].device, dtype=dtype) _deepgemm_grouped_fp8_nt_masked((qsilu_out, qsilu_out_scale), (w2, w2_scale), gemm_out_b, masked_m, expected_m) @@ -120,6 +151,76 @@ def _get_mega_moe_cumulative_stats(num_local_experts: int, device: torch.device, return stats +def prepare_sm90_mega_moe_weights(w13: Any) -> None: + """Replace standard gate/up FP8 weights with SM90 Mega-MoE layout.""" + + if getattr(w13, "sm90_mega_moe_prepared", False): + return + weight = w13.weight + num_groups, n, *rest = weight.shape + granularity = 8 + half = n // 2 + assert half % granularity == 0 + gate = weight[:, :half].reshape(num_groups, half // granularity, granularity, *rest) + up = weight[:, half:].reshape(num_groups, half // granularity, granularity, *rest) + w13.weight = torch.stack((gate, up), dim=2).reshape(num_groups, n, *rest) + w13.sm90_mega_moe_prepared = True + + +def _sm90_mega_moe_impl( + hidden_states: torch.Tensor, + w13: Any, + w2: Any, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation_clamp: Optional[float], +): + buffer = getattr(dist_group_manager, "ep_mega_moe_buffer", None) + if buffer is None: + raise RuntimeError("SM90 Mega MoE buffer is not initialized") + if hidden_states.shape[0] > buffer.num_max_tokens_per_rank: + raise RuntimeError( + f"SM90 Mega MoE got {hidden_states.shape[0]} tokens, exceeding " + f"num_max_tokens_per_rank={buffer.num_max_tokens_per_rank}" + ) + if not getattr(w13, "sm90_mega_moe_prepared", False): + raise RuntimeError("SM90 Mega MoE weights were not prepared after loading") + + num_tokens = hidden_states.shape[0] + deep_gemm.mega_moe_pre_dispatch_sm90( + hidden_states, + topk_ids.to(torch.int32), + topk_weights.to(torch.float32), + buffer.x, + buffer.x_sf, + buffer.topk_idx, + buffer.topk_weights, + num_tokens=num_tokens, + group_size=128, + routed_scaling_factor=1.0, + ) + # Match DeepGEMM's SM90 integration exactly. In particular, do not pass + # cumulative expert statistics here: the production SGLang path leaves the + # optional pointer null, and the SM90 persistent collective must not gain an + # extra per-layer side effect across repeated decode launches. + output = torch.empty( + (max(num_tokens, 1), hidden_states.shape[1]), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + deep_gemm.fp8_mega_moe( + output, + (w13.weight, w13.weight_scale), + (w2.weight, w2.weight_scale), + buffer, + recipe=(128, 128, 128), + activation="swiglu", + activation_clamp=activation_clamp, + fast_math=True, + ) + return output[:num_tokens] + + def mega_moe_impl( hidden_states: torch.Tensor, w13: Any, @@ -127,7 +228,17 @@ def mega_moe_impl( topk_weights: torch.Tensor, topk_ids: torch.Tensor, quant_method: Any, + swiglu_limit: Optional[float] = None, ): + if use_sm90_mega_moe(quant_method): + return _sm90_mega_moe_impl( + hidden_states, + w13, + w2, + topk_weights, + topk_ids, + activation_clamp=swiglu_limit, + ) if not (HAS_DEEPGEMM and hasattr(deep_gemm, "fp8_fp4_mega_moe")): raise RuntimeError("deep_gemm does not provide fp8-fp4 Mega MoE kernel") @@ -174,7 +285,7 @@ def quantize_fused_experts_input( quant_method: Any, ): check_ep_expert_dtype(quant_method) - if use_sm100_mega_moe(quant_method): + if use_mega_moe(quant_method): from deep_gemm.utils import per_token_cast_to_fp8 return per_token_cast_to_fp8( @@ -191,6 +302,114 @@ def quantize_fused_experts_input( return per_token_group_quant_fp8(hidden_states, block_size_k, dtype=w13.weight.dtype) +def legacy_normal_moe_forward( + num_recv_tokens_per_expert_list: List[int], + recv_x: Tuple[torch.Tensor, torch.Tensor], + recv_topk_idx: torch.Tensor, + recv_topk_weights: torch.Tensor, + w1: torch.Tensor, + w1_scale: torch.Tensor, + w2: torch.Tensor, + w2_scale: torch.Tensor, + hidden_dtype: torch.dtype, + swiglu_limit=None, + swiglu_alpha=1.0, + swiglu_clamp_up_add_one=True, +) -> torch.Tensor: + """Run MoE on the token layout returned by DeepEP's legacy normal buffer.""" + hidden_size = recv_x[0].shape[1] + intermediate_twice = w1.shape[1] + intermediate_size = intermediate_twice // 2 + block_size_k = w1.shape[2] // w1_scale.shape[2] + all_tokens = sum(num_recv_tokens_per_expert_list) + gather_out = torch.empty( + (recv_x[0].shape[0], hidden_size), + device=recv_x[0].device, + dtype=hidden_dtype, + ) + if all_tokens == 0: + if Autotuner.is_autotune_warmup(): + gemm_out_a = torch.zeros((1, intermediate_twice), device=recv_x[0].device, dtype=hidden_dtype) + silu_out = torch.zeros((1, intermediate_size), device=recv_x[0].device, dtype=hidden_dtype) + silu_and_mul_fwd( + gemm_out_a, + silu_out, + limit=swiglu_limit, + alpha=swiglu_alpha, + clamp_up_add_one=swiglu_clamp_up_add_one, + ) + return gather_out + + input_tensor = torch.empty( + (all_tokens, hidden_size), + device=recv_x[0].device, + dtype=recv_x[0].dtype, + ) + input_scale = torch.empty( + (all_tokens, hidden_size // block_size_k), + device=recv_x[0].device, + dtype=recv_x[1].dtype, + ) + m_indices = torch.empty(all_tokens, device=recv_x[0].device, dtype=torch.int32) + output_index = torch.empty_like(recv_topk_idx, dtype=torch.int32) + padded_counts = torch.tensor( + num_recv_tokens_per_expert_list, + dtype=torch.int32, + pin_memory=True, + device="cpu", + ).cuda(non_blocking=True) + valid_expert_ids = recv_topk_idx[recv_topk_idx >= 0].to(torch.int64) + valid_counts = torch.bincount(valid_expert_ids, minlength=w1.shape[0]).to(torch.int32) + expert_start_loc = torch.empty_like(padded_counts) + ep_scatter( + recv_x[0], + recv_x[1], + recv_topk_idx, + padded_counts, + valid_counts, + expert_start_loc, + input_tensor, + input_scale, + m_indices, + output_index, + ) + + gemm_out_a = torch.empty( + (all_tokens, intermediate_twice), + device=recv_x[0].device, + dtype=hidden_dtype, + ) + input_scale = tma_align_input_scale(input_scale) + deepgemm_grouped_fp8_nt_contiguous((input_tensor, input_scale), (w1, w1_scale), gemm_out_a, m_indices) + silu_out = torch.empty( + (all_tokens, intermediate_size), + device=recv_x[0].device, + dtype=hidden_dtype, + ) + silu_and_mul_fwd( + gemm_out_a, + silu_out, + limit=swiglu_limit, + alpha=swiglu_alpha, + clamp_up_add_one=swiglu_clamp_up_add_one, + ) + quant_silu = per_token_group_quant_fp8( + silu_out, + block_size_k, + dtype=w2.dtype, + column_major_scales=True, + scale_tma_aligned=True, + ) + gemm_out_b = torch.empty( + (all_tokens, hidden_size), + device=recv_x[0].device, + dtype=hidden_dtype, + ) + deepgemm_grouped_fp8_nt_contiguous(quant_silu, (w2, w2_scale), gemm_out_b, m_indices) + ep_gather(gemm_out_b, recv_topk_idx, recv_topk_weights, output_index, gather_out) + return gather_out + + def fused_experts( hidden_states: torch.Tensor, w13: Any, @@ -201,10 +420,21 @@ def fused_experts( quant_method: Any, is_prefill: Optional[bool], previous_event: Optional[Any] = None, + swiglu_limit=None, + swiglu_alpha=1.0, + swiglu_clamp_up_add_one=True, ): check_ep_expert_dtype(quant_method) - if use_sm100_mega_moe(quant_method): - return mega_moe_impl(hidden_states, w13, w2, topk_weights, topk_idx, quant_method) + if use_mega_moe(quant_method): + return mega_moe_impl( + hidden_states, + w13, + w2, + topk_weights, + topk_idx, + quant_method, + swiglu_limit=swiglu_limit, + ) buffer = dist_group_manager.ep_buffer if is_prefill else dist_group_manager.ep_low_latency_buffer return fused_experts_impl( @@ -222,6 +452,9 @@ def fused_experts( w1_scale=w13.weight_scale, w2_scale=w2.weight_scale, previous_event=previous_event, + swiglu_limit=swiglu_limit, + swiglu_alpha=swiglu_alpha, + swiglu_clamp_up_add_one=swiglu_clamp_up_add_one, ) @@ -240,6 +473,9 @@ def fused_experts_impl( w1_scale: Optional[torch.Tensor] = None, w2_scale: Optional[torch.Tensor] = None, previous_event: Optional[Any] = None, + swiglu_limit=None, + swiglu_alpha=1.0, + swiglu_clamp_up_add_one=True, ): # Check constraints. assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch" @@ -263,6 +499,57 @@ def fused_experts_impl( if is_prefill: qinput_tensor, input_scale = per_token_group_quant_fp8(hidden_states, block_size_k, dtype=w1.dtype) allocate_on_comm_stream = previous_event is not None + if dist_group_manager.ep_prefill_uses_legacy_buffer: + ( + num_tokens_per_rank, + num_tokens_per_rdma_rank, + num_tokens_per_expert, + is_token_in_rank, + layout_event, + ) = buffer.get_dispatch_layout( + topk_idx, + num_experts, + previous_event=previous_event, + async_finish=False, + allocate_on_comm_stream=False, + ) + (recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, _,) = buffer.dispatch( + (qinput_tensor, input_scale), + topk_idx=topk_idx, + topk_weights=topk_weights, + num_tokens_per_rank=num_tokens_per_rank, + num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, + is_token_in_rank=is_token_in_rank, + num_tokens_per_expert=num_tokens_per_expert, + previous_event=layout_event, + async_finish=False, + allocate_on_comm_stream=False, + expert_alignment=128, + ) + gather_out = legacy_normal_moe_forward( + num_recv_tokens_per_expert_list, + recv_x, + recv_topk_idx, + recv_topk_weights, + w1, + w1_scale, + w2, + w2_scale, + hidden_states.dtype, + swiglu_limit, + swiglu_alpha, + swiglu_clamp_up_add_one, + ) + combined_x, _, _ = buffer.combine( + gather_out, + handle, + topk_weights=None, + previous_event=previous_event, + async_finish=False, + allocate_on_comm_stream=False, + ) + return combined_x + # Expanded dispatch directly produces expert-contiguous, alignment-padded inputs: # recv_x[0]: [num_expanded_tokens, hidden] # recv_x[1]: [num_expanded_tokens, hidden // block_size_k], with a @@ -309,6 +596,9 @@ def fused_experts_impl( block_size_k=block_size_k, workspace=dist_group_manager.get_deep_ep_prefill_moe_workspace(), hidden_dtype=hidden_states.dtype, + swiglu_limit=swiglu_limit, + swiglu_alpha=swiglu_alpha, + swiglu_clamp_up_add_one=swiglu_clamp_up_add_one, ) else: gather_out = torch.empty( @@ -324,7 +614,13 @@ def fused_experts_impl( N = w1.shape[1] _gemm_out_a = torch.zeros((1, N), device=hidden_states.device, dtype=hidden_states.dtype) _silu_out = torch.zeros((1, N // 2), device=hidden_states.device, dtype=hidden_states.dtype) - silu_and_mul_fwd(_gemm_out_a.view(-1, N), _silu_out) + silu_and_mul_fwd( + _gemm_out_a.view(-1, N), + _silu_out, + limit=swiglu_limit, + alpha=swiglu_alpha, + clamp_up_add_one=swiglu_clamp_up_add_one, + ) _gemm_out_a, _silu_out = None, None del recv_x @@ -350,7 +646,19 @@ def fused_experts_impl( return_recv_hook=False, ) # deepgemm - gemm_out_b = masked_group_gemm(recv_x, masked_m, hidden_states.dtype, w1, w1_scale, w2, w2_scale, expected_m) + gemm_out_b = masked_group_gemm( + recv_x, + masked_m, + hidden_states.dtype, + w1, + w1_scale, + w2, + w2_scale, + expected_m, + swiglu_limit, + swiglu_alpha, + swiglu_clamp_up_add_one, + ) # low latency combine combined_x, event_overlap, hook = buffer.low_latency_combine( gemm_out_b, topk_idx, topk_weights, handle, async_finish=False, return_recv_hook=False @@ -468,6 +776,9 @@ def chunked_expanded_moe_forward( block_size_k: int, workspace: torch.Tensor, # [workspace_bytes], uint8 hidden_dtype: torch.dtype, # scalar dtype descriptor + swiglu_limit=None, + swiglu_alpha=1.0, + swiglu_clamp_up_add_one=True, ): """Run bounded expanded MoE and rewrite metadata for dense DeepEP combine.""" alignment = 128 @@ -535,7 +846,13 @@ def chunked_expanded_moe_forward( gemm_out_a, m_indices[chunk_start:chunk_end], ) - silu_and_mul_fwd(gemm_out_a, silu_out) + silu_and_mul_fwd( + gemm_out_a, + silu_out, + limit=swiglu_limit, + alpha=swiglu_alpha, + clamp_up_add_one=swiglu_clamp_up_add_one, + ) workspace_manager.free(gemm_out_a) del gemm_out_a @@ -588,6 +905,8 @@ def _deepgemm_grouped_fp8_nt_masked( expected_m: int, ): if HAS_DEEPGEMM: + if hasattr(deep_gemm, "fp8_m_grouped_gemm_nt_masked"): + return deep_gemm.fp8_m_grouped_gemm_nt_masked(input_tuple, w_tuple, out, masked_m, expected_m) if hasattr(deep_gemm, "m_grouped_fp8_gemm_nt_masked"): return deep_gemm.m_grouped_fp8_gemm_nt_masked(input_tuple, w_tuple, out, masked_m, expected_m) if hasattr(deep_gemm, "m_grouped_gemm_fp8_fp8_bf16_nt_masked"): diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_topk.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_topk.py index fb0323cd4b..a8bc8b9e58 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_topk.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_topk.py @@ -89,6 +89,125 @@ def argsort(x, x_1, ids, dim: tl.core.constexpr = None, descending: tl.core.cons return x, x_1, ids +@triton.jit +def single_group_sigmoid_topk_kernel( + gating_output_ptr, + gating_output_stride_m, + correction_bias_ptr, + out_topk_weights, + out_topk_weights_stride_m, + out_topk_ids, + out_topk_ids_stride_m, + total_expert_num, + HAS_CORRECTION_BIAS: tl.constexpr, + EXPERT_BLOCK_SIZE: tl.constexpr, + TOPK_BLOCK_SIZE: tl.constexpr, + TOPK_NUM: tl.constexpr, + RENORMALIZE: tl.constexpr, +): + """Select biased sigmoid top-k when every expert belongs to one group. + + GLM-5.3-Flash declares one expert group. The generic grouped kernel still + materializes a scratch buffer, synchronizes twice, and bitonic-sorts the + full power-of-two expert block. Repeated reductions are substantially + cheaper for the small fixed top-k used by MoE decode. + """ + + token_index = tl.program_id(axis=0) + offs_n = tl.arange(0, EXPERT_BLOCK_SIZE) + valid = offs_n < total_expert_num + hidden_states = tl.load( + gating_output_ptr + token_index * gating_output_stride_m + offs_n, + mask=valid, + other=-float("inf"), + ).to(tl.float32) + old_scores = tl.sigmoid(hidden_states) + if HAS_CORRECTION_BIAS: + scores = old_scores + tl.load(correction_bias_ptr + offs_n, mask=valid, other=0.0) + else: + scores = old_scores + scores = tl.where(valid, scores, -float("inf")) + + for topk_index in tl.static_range(0, TOPK_NUM): + selected_index = tl.argmax(scores, axis=0) + selected_weight = tl.sum(tl.where(offs_n == selected_index, old_scores, 0.0), axis=0) + tl.store( + out_topk_weights + token_index * out_topk_weights_stride_m + topk_index, + selected_weight, + ) + tl.store( + out_topk_ids + token_index * out_topk_ids_stride_m + topk_index, + selected_index, + ) + scores = tl.where(offs_n == selected_index, -float("inf"), scores) + + if RENORMALIZE: + topk_offs = tl.arange(0, TOPK_BLOCK_SIZE) + topk_mask = topk_offs < TOPK_NUM + weights = tl.load( + out_topk_weights + token_index * out_topk_weights_stride_m + topk_offs, + mask=topk_mask, + other=0.0, + ) + weight_sum = tl.sum(weights, axis=0) + tl.store( + out_topk_weights + token_index * out_topk_weights_stride_m + topk_offs, + weights / weight_sum, + mask=topk_mask, + ) + + +@triton.jit +def single_group_sigmoid_topk_bitonic_kernel( + gating_output_ptr, + gating_output_stride_m, + correction_bias_ptr, + out_topk_weights, + out_topk_weights_stride_m, + out_topk_ids, + out_topk_ids_stride_m, + total_expert_num, + HAS_CORRECTION_BIAS: tl.constexpr, + EXPERT_BLOCK_SIZE: tl.constexpr, + TOPK_NUM: tl.constexpr, + RENORMALIZE: tl.constexpr, +): + """Scratch-free single-group path retaining the legacy bitonic order.""" + + token_index = tl.program_id(axis=0) + offs_n = tl.arange(0, EXPERT_BLOCK_SIZE) + hidden_states = tl.load( + gating_output_ptr + token_index * gating_output_stride_m + offs_n, + mask=offs_n < total_expert_num, + other=-10000000.0, + ).to(tl.float32) + old_scores = tl.sigmoid(hidden_states) + if HAS_CORRECTION_BIAS: + scores = old_scores + tl.load( + correction_bias_ptr + offs_n, + mask=offs_n < total_expert_num, + other=-10000000.0, + ) + else: + scores = old_scores + + _, sorted_scores, sorted_indexes = argsort(scores, old_scores, offs_n, descending=True) + if RENORMALIZE: + sum_scores = tl.sum(tl.where(offs_n < TOPK_NUM, sorted_scores, 0.0)) + sorted_scores = sorted_scores / sum_scores + + tl.store( + out_topk_weights + token_index * out_topk_weights_stride_m + offs_n, + sorted_scores, + mask=offs_n < TOPK_NUM, + ) + tl.store( + out_topk_ids + token_index * out_topk_ids_stride_m + offs_n, + sorted_indexes, + mask=offs_n < TOPK_NUM, + ) + + @triton.jit def grouped_topk_kernel( gating_output_ptr, @@ -212,6 +331,7 @@ def triton_grouped_topk( topk_group: int = 0, scoring_func: str = "softmax", group_score_used_topk_num=2, + use_single_group_fast_path: bool = True, ): if correction_bias is not None: @@ -220,6 +340,29 @@ def triton_grouped_topk( has_correction_bias = False token_num, total_expert_num = gating_output.shape + + if use_single_group_fast_path and num_expert_group == 1 and topk_group == 1 and scoring_func == "sigmoid": + out_topk_weights = torch.empty((token_num, topk), dtype=torch.float32, device="cuda") + out_topk_ids = torch.empty((token_num, topk), dtype=torch.long, device="cuda") + single_group_sigmoid_topk_kernel[(token_num,)]( + gating_output, + gating_output.stride(0), + correction_bias, + out_topk_weights, + out_topk_weights.stride(0), + out_topk_ids, + out_topk_ids.stride(0), + total_expert_num, + HAS_CORRECTION_BIAS=has_correction_bias, + EXPERT_BLOCK_SIZE=triton.next_power_of_2(total_expert_num), + TOPK_BLOCK_SIZE=triton.next_power_of_2(topk), + TOPK_NUM=topk, + RENORMALIZE=renormalize, + num_warps=1, + num_stages=1, + ) + return out_topk_weights, out_topk_ids + if gating_output.dtype == torch.float64: dtype = torch.float64 else: diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py index a63d92692e..13d7540cc3 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py @@ -24,6 +24,7 @@ def _silu_and_mul_kernel_fast( NEED_MASK: tl.constexpr, layout: tl.constexpr = "blocked", # "blocked" or "interleaved" USE_LIMIT_AND_ALPHA: tl.constexpr = False, + CLAMP_UP_ADD_ONE: tl.constexpr = True, USE_TANH_APPROXIMATE_GELU: tl.constexpr = False, ): stride_input_m = tl.cast(stride_input_m, dtype=tl.int64) @@ -70,11 +71,9 @@ def _silu_and_mul_kernel_fast( up = tl.minimum(tl.maximum(up, -limit), limit) gate = 1 / (1 + tl.exp(-gate * alpha)) * gate gate = gate.to(input_ptr.dtype.element_ty) - tl.store( - output_ptr + out_offsets, - (up + 1) * gate, - mask=mask, - ) + if CLAMP_UP_ADD_ONE: + up += 1 + tl.store(output_ptr + out_offsets, up * gate, mask=mask) else: if USE_TANH_APPROXIMATE_GELU: # tanh-approx GELU, matching Gemma's gelu_pytorch_tanh MLP. @@ -120,6 +119,7 @@ def silu_and_mul_fwd( layout="blocked", limit=None, alpha=None, + clamp_up_add_one=True, run_config=None, ): assert input.stride(-1) == 1 @@ -171,6 +171,7 @@ def silu_and_mul_fwd( num_warps=num_warps, layout=layout, USE_LIMIT_AND_ALPHA=USE_LIMIT_AND_ALPHA, + CLAMP_UP_ADD_ONE=clamp_up_add_one, USE_TANH_APPROXIMATE_GELU=ffn_use_tanh_approximate_gelu(), ) return diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py index aa91f15ed9..f3b3426acf 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py @@ -24,9 +24,13 @@ def _silu_and_mul_post_quant_kernel( size_n, fp8_max, fp8_min, + limit: tl.constexpr, + alpha: tl.constexpr, BLOCK_N: tl.constexpr, NUM_STAGE: tl.constexpr, USE_TANH_APPROXIMATE_GELU: tl.constexpr = False, + USE_LIMIT_AND_ALPHA: tl.constexpr = False, + CLAMP_UP_ADD_ONE: tl.constexpr = True, ): expert_id = tl.program_id(2) token_id = tl.program_id(1) @@ -51,7 +55,13 @@ def _silu_and_mul_post_quant_kernel( for token_index in tl.range(token_id, token_num_cur_expert, block_num_per_expert, num_stages=NUM_STAGE): gate = tl.load(input_ptr_offs + token_index * stride_input_1, mask=offs_in_d < size_n, other=0.0).to(tl.float32) up = tl.load(input_ptr_offs + token_index * stride_input_1 + size_n, mask=offs_in_d < size_n, other=0.0) - if USE_TANH_APPROXIMATE_GELU: + if USE_LIMIT_AND_ALPHA: + gate = tl.minimum(gate, limit) + up = tl.minimum(tl.maximum(up, -limit), limit) + gate = gate / (1 + tl.exp(-gate * alpha)) + if CLAMP_UP_ADD_ONE: + up += 1 + elif USE_TANH_APPROXIMATE_GELU: gate_cubed = gate * gate * gate tanh_arg = 0.7978845608028654 * (gate + 0.044715 * gate_cubed) tanh_val = 2.0 / (1.0 + tl.exp(-2.0 * tanh_arg)) - 1.0 @@ -80,6 +90,9 @@ def silu_and_mul_masked_post_quant_fwd( output_scale: torch.Tensor, quant_group_size: int, masked_m: torch.Tensor, + limit=None, + alpha=None, + clamp_up_add_one=True, ): """ input shape [expert_num, token_num_padded, hidden_dim] @@ -121,6 +134,7 @@ def silu_and_mul_masked_post_quant_fwd( finfo = torch.finfo(torch.float8_e4m3fn) fp8_max = finfo.max fp8_min = -fp8_max + assert (limit is None and alpha is None) or (limit is not None and alpha is not None) _silu_and_mul_post_quant_kernel[grid]( input, @@ -133,9 +147,13 @@ def silu_and_mul_masked_post_quant_fwd( size_n, fp8_max, fp8_min, + limit=limit, + alpha=alpha, BLOCK_N=BLOCK_N, NUM_STAGE=NUM_STAGES, USE_TANH_APPROXIMATE_GELU=ffn_use_tanh_approximate_gelu(), + USE_LIMIT_AND_ALPHA=limit is not None, + CLAMP_UP_ADD_ONE=clamp_up_add_one, num_warps=num_warps, ) return diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py b/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py index 2bf325340f..2afeec8f60 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py @@ -1,6 +1,6 @@ # Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/attention/mamba/causal_conv1d.py -from typing import Optional +from typing import List, Optional import torch @@ -15,6 +15,7 @@ def causal_conv1d_fn( conv_states: Optional[torch.Tensor] = None, activation: Optional[str] = "silu", pad_slot_id: int = -1, + seq_lens_cpu: Optional[List[int]] = None, **kwargs, ): """ @@ -48,6 +49,32 @@ def causal_conv1d_fn( """ if activation not in [None, "silu", "swish"]: raise NotImplementedError("activation must be None, silu, or swish") + # The compiled CUDA kernel requires sequence-contiguous input. KDA's + # projection is token-major, so its dim-major transpose would otherwise + # materialize a large copy at every linear-attention layer. SGLang's + # Triton fallback accepts arbitrary input/output strides and is already a + # runtime dependency in the optimized GLM-5 deployment image. + if x.stride(-1) != 1 and seq_lens_cpu is not None: + try: + from sglang.kernels.ops.mamba.causal_conv1d_triton import ( + causal_conv1d_fn as causal_conv1d_strided, + ) + except ImportError: + pass + else: + return causal_conv1d_strided( + x, + weight, + bias, + conv_states=conv_states, + query_start_loc=query_start_loc, + seq_lens_cpu=seq_lens_cpu, + cache_indices=cache_indices, + has_initial_state=has_initial_state, + activation=activation, + pad_slot_id=pad_slot_id, + ) + from sgl_kernel import causal_conv1d_fwd if x.stride(-1) != 1: diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py index cd3b0962a3..aa91129088 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py @@ -8,8 +8,12 @@ # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang from .chunk import chunk_gated_delta_rule from .fused_recurrent import fused_recurrent_gated_delta_rule +from .kda import chunk_kda_with_fused_gate, fused_recurrent_kda, kda_safe_gate __all__ = [ "chunk_gated_delta_rule", "fused_recurrent_gated_delta_rule", + "fused_recurrent_kda", + "chunk_kda_with_fused_gate", + "kda_safe_gate", ] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_delta_h.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_delta_h.py index 97933b2ac2..2ab864eaf4 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_delta_h.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_delta_h.py @@ -24,6 +24,7 @@ { "USE_G": lambda args: args["g"] is not None, "USE_GK": lambda args: args["gk"] is not None, + "USE_EXP2": lambda args: args["use_exp2"], "USE_INITIAL_STATE": lambda args: args["h0"] is not None, "STORE_FINAL_STATE": lambda args: args["ht"] is not None, "SAVE_NEW_VALUE": lambda args: args["v_new"] is not None, @@ -43,6 +44,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( ht, cu_seqlens, chunk_offsets, + use_exp2, T, H: tl.constexpr, Hg: tl.constexpr, @@ -52,6 +54,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( BV: tl.constexpr, USE_G: tl.constexpr, USE_GK: tl.constexpr, + USE_EXP2: tl.constexpr, USE_INITIAL_STATE: tl.constexpr, STORE_FINAL_STATE: tl.constexpr, SAVE_NEW_VALUE: tl.constexpr, @@ -169,7 +172,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( mask=(o_k1 < K), other=0.0, ) - b_h1 *= exp(b_gk_last1)[:, None] + b_h1 *= (tl.exp2(b_gk_last1) if USE_EXP2 else exp(b_gk_last1))[:, None] if K > 64: o_k2 = 64 + o_k1 b_gk_last2 = tl.load( @@ -177,7 +180,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( mask=(o_k2 < K), other=0.0, ) - b_h2 *= exp(b_gk_last2)[:, None] + b_h2 *= (tl.exp2(b_gk_last2) if USE_EXP2 else exp(b_gk_last2))[:, None] if K > 128: o_k3 = 128 + o_k1 b_gk_last3 = tl.load( @@ -185,7 +188,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( mask=(o_k3 < K), other=0.0, ) - b_h3 *= exp(b_gk_last3)[:, None] + b_h3 *= (tl.exp2(b_gk_last3) if USE_EXP2 else exp(b_gk_last3))[:, None] if K > 192: o_k4 = 192 + o_k1 b_gk_last4 = tl.load( @@ -193,7 +196,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( mask=(o_k4 < K), other=0.0, ) - b_h4 *= exp(b_gk_last4)[:, None] + b_h4 *= (tl.exp2(b_gk_last4) if USE_EXP2 else exp(b_gk_last4))[:, None] b_v = b_v.to(k.dtype.element_ty) p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) @@ -264,6 +267,8 @@ def chunk_gated_delta_rule_fwd_h( chunk_size: int = 64, # SY: remove this argument and force chunk size 64? save_new_value: bool = True, cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.Tensor | None = None, + use_exp2: bool = False, run_config=None, ) -> tuple[torch.Tensor, torch.Tensor]: # This kernel is slightly different from fla to support Q/K with different head numbers. @@ -272,7 +277,8 @@ def chunk_gated_delta_rule_fwd_h( H = u.shape[-2] BT = chunk_size - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) # N: the actual number of sequences in the batch with either equal or variable lengths if cu_seqlens is None: N, NT, chunk_offsets = B, triton.cdiv(T, BT), None @@ -311,6 +317,7 @@ def chunk_gated_delta_rule_fwd_h( ht=final_state, cu_seqlens=cu_seqlens, chunk_offsets=chunk_offsets, + use_exp2=use_exp2, T=T, H=H, Hg=Hg, diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py index 5dfbd6e4ab..3e782b69b6 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py @@ -8,6 +8,8 @@ # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang # ruff: noqa: E501 +import os + import torch import triton @@ -16,6 +18,9 @@ from .op import exp +ENABLE_FAST_MTP_KDA = os.getenv("LIGHTLLM_ENABLE_FAST_MTP_KDA", "0") == "1" + + @triton.heuristics( { "USE_INITIAL_STATE": lambda args: args["h0"] is not None, @@ -45,6 +50,7 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( a_raw, # [B*T, HV] raw alpha values (before softplus) b_raw, # [B*T, HV] raw beta values (before sigmoid) scale, + kda_lower_bound, N: tl.int64, # num of sequences T: tl.int64, # num of tokens B: tl.constexpr, @@ -105,8 +111,12 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( if FUSE_GATING: # Fused gating: load per-head constants once, compute g/beta inline per token b_A_log = tl.load(A_log + i_hv).to(tl.float32) - b_dt_bias = tl.load(dt_bias + i_hv).to(tl.float32) - p_a_raw = a_raw + bos * stride_a_tok + i_hv + if IS_KDA: + p_dt_bias = dt_bias + i_hv * K + o_k + p_a_raw = a_raw + bos * stride_a_tok + i_hv * K + o_k + else: + b_dt_bias = tl.load(dt_bias + i_hv).to(tl.float32) + p_a_raw = a_raw + bos * stride_a_tok + i_hv p_b_raw = b_raw + bos * stride_b_tok + i_hv else: if IS_BETA_HEADWISE: @@ -151,16 +161,23 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( b_q = b_q * scale # [BK, BV] if FUSE_GATING: - # Compute g = -exp(A_log) * softplus(a_raw + dt_bias) inline - b_a = tl.load(p_a_raw).to(tl.float32) - x = b_a + b_dt_bias - softplus_x = tl.where( - SOFTPLUS_BETA * x <= SOFTPLUS_THRESHOLD, - (1.0 / SOFTPLUS_BETA) * tl.log(1.0 + tl.exp(SOFTPLUS_BETA * x)), - x, - ) - b_g = -tl.exp(b_A_log) * softplus_x - b_h *= exp(b_g) + if IS_KDA: + # GLM-5 safe gate is a K-wide vector per value head. + b_a = tl.load(p_a_raw, mask=mask_k, other=0).to(tl.float32) + b_dt_bias = tl.load(p_dt_bias, mask=mask_k, other=0).to(tl.float32) + b_gk = kda_lower_bound * tl.sigmoid(tl.exp(b_A_log) * (b_a + b_dt_bias)) + b_h *= exp(b_gk[:, None]) + else: + # GDN scalar gate: g = -exp(A_log) * softplus(a_raw + dt_bias). + b_a = tl.load(p_a_raw).to(tl.float32) + x = b_a + b_dt_bias + softplus_x = tl.where( + SOFTPLUS_BETA * x <= SOFTPLUS_THRESHOLD, + (1.0 / SOFTPLUS_BETA) * tl.log(1.0 + tl.exp(SOFTPLUS_BETA * x)), + x, + ) + b_g = -tl.exp(b_A_log) * softplus_x + b_h *= exp(b_g) # Compute beta = sigmoid(b_raw) inline b_b = tl.load(p_b_raw).to(tl.float32) b_beta = tl.sigmoid(b_b) @@ -267,6 +284,8 @@ def fused_recurrent_gated_delta_rule_fwd( a_raw: torch.Tensor | None = None, b_raw: torch.Tensor | None = None, out: torch.Tensor | None = None, + is_kda: bool = False, + kda_lower_bound: float = -5.0, ) -> tuple[torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] HV = v.shape[2] @@ -277,15 +296,28 @@ def fused_recurrent_gated_delta_rule_fwd( q, stride_q_tok = _ensure_qkv_token_strided(q, H * K) k, stride_k_tok = _ensure_qkv_token_strided(k, H * K) v, stride_v_tok = _ensure_qkv_token_strided(v, HV * V) - a_raw, stride_a_tok = _ensure_gate_token_strided(a_raw, HV) + a_raw, stride_a_tok = _ensure_gate_token_strided(a_raw, HV * K if is_kda else HV) b_raw, stride_b_tok = _ensure_gate_token_strided(b_raw, HV) BK = triton.next_power_of_2(K) + is_spec_verify = ENABLE_FAST_MTP_KDA and cu_seqlens is not None and num_accepted_tokens is not None if T == 1: # Decode path: use larger BV to reduce kernel instances (4 blocks instead of 16) # and more warps for better SM utilization at T=1 where there's no pipelining benefit BV = min(triton.next_power_of_2(V), 32) num_warps = 4 num_stages = 1 + elif is_spec_verify: + # MTP verification is represented as one packed tensor, so ``T`` is the + # total physical token count even though every sequence is only a few + # tokens long. Treating this as prefill creates 16 single-warp value + # tiles per head for GLM-5 (V=128). The short recurrent loop benefits + # from the same BV=32 shape used by SGLang's KDA target-verify kernel, + # cutting the launch grid to four value tiles. Keep this opt-in because + # the wider tile can alter bf16 rounding at the 1e-5 level even though + # the recurrence is mathematically identical. + BV = min(triton.next_power_of_2(V), 32) + num_warps = 1 + num_stages = 3 else: # Prefill path: small BV for better pipelining across sequence length BV = min(triton.next_power_of_2(V), 8) @@ -346,6 +378,7 @@ def fused_recurrent_gated_delta_rule_fwd( a_raw=a_raw, b_raw=b_raw, scale=scale, + kda_lower_bound=kda_lower_bound, N=N, T=T, B=B, @@ -371,7 +404,7 @@ def fused_recurrent_gated_delta_rule_fwd( IS_BETA_HEADWISE=False if fuse_gating else (beta.ndim == v.ndim), USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, INPLACE_FINAL_STATE=inplace_final_state, - IS_KDA=False, + IS_KDA=is_kda, FUSE_GATING=fuse_gating, num_warps=num_warps, num_stages=num_stages, @@ -402,6 +435,8 @@ def forward( a_raw: torch.Tensor | None = None, b_raw: torch.Tensor | None = None, out: torch.Tensor | None = None, + is_kda: bool = False, + kda_lower_bound: float = -5.0, ): # q/k/v/a_raw/b_raw may be non-contiguous column views of one projection # output; the kernel handles them via per-token strides (no copies). @@ -424,6 +459,8 @@ def forward( a_raw=a_raw, b_raw=b_raw, out=out, + is_kda=is_kda, + kda_lower_bound=kda_lower_bound, ) return o, final_state @@ -449,6 +486,8 @@ def fused_recurrent_gated_delta_rule( a_raw: torch.Tensor | None = None, b_raw: torch.Tensor | None = None, out: torch.Tensor | None = None, + is_kda: bool = False, + kda_lower_bound: float = -5.0, ) -> tuple[torch.Tensor, torch.Tensor]: r""" Args: @@ -531,5 +570,7 @@ def fused_recurrent_gated_delta_rule( a_raw, b_raw, out, + is_kda, + kda_lower_bound, ) return o, final_state diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/index.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/index.py index 8b1d59fc63..d09f3e82a2 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/index.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/index.py @@ -27,4 +27,7 @@ def prepare_chunk_indices(cu_seqlens: torch.LongTensor, chunk_size: int) -> torc @tensor_cache def prepare_chunk_offsets(cu_seqlens: torch.LongTensor, chunk_size: int) -> torch.LongTensor: - return torch.cat([cu_seqlens.new_tensor([0]), triton.cdiv(prepare_lens(cu_seqlens), chunk_size)]).cumsum(-1) + # new_tensor([0]) stages a regular CPU tensor before copying it to CUDA, + # which is illegal inside CUDA Graph capture. new_zeros allocates and + # initializes directly on the source tensor's device. + return torch.cat([cu_seqlens.new_zeros(1), triton.cdiv(prepare_lens(cu_seqlens), chunk_size)]).cumsum(-1) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py new file mode 100644 index 0000000000..0cf7ce8fd8 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py @@ -0,0 +1,1136 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang + +"""KDA helpers built on LightLLM's continuous-batching recurrent kernel.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from .chunk_delta_h import chunk_gated_delta_rule_fwd_h +from .cumsum import chunk_local_cumsum +from .fused_recurrent import ENABLE_FAST_MTP_KDA, fused_recurrent_gated_delta_rule +from .index import prepare_chunk_indices +from .l2norm import l2norm_fwd +from .op import exp2, log +from .solve_tril import solve_tril + + +FLA_CHUNK_SIZE = 64 +RCP_LN2 = 1.4426950216293335 + + +def cdiv(a: int, b: int) -> int: + return -(a // -b) + + +def next_power_of_2(n: int) -> int: + return 1 if n < 1 else 1 << (n - 1).bit_length() + + +def kda_safe_gate( + raw_gate: torch.Tensor, + a_log: torch.Tensor, + gate_bias: torch.Tensor, + lower_bound: float = -5.0, +) -> torch.Tensor: + """GLM-5 bounded KDA decay in fp32. + + ``raw_gate`` is ``[..., heads, key_dim]``; ``a_log`` is per-head and + ``gate_bias`` is per head/key coordinate. + """ + + head_count = a_log.numel() + key_dim = gate_bias.numel() // head_count + gate = raw_gate.float().view(*raw_gate.shape[:-1], head_count, key_dim) + amplitude = a_log.float().reshape(*((1,) * (gate.ndim - 2)), head_count, 1).exp() + bias = gate_bias.float().reshape(*((1,) * (gate.ndim - 2)), head_count, key_dim) + return lower_bound * torch.sigmoid(amplitude * (gate + bias)) + + +def fused_recurrent_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_gate: torch.Tensor, + raw_beta: torch.Tensor, + a_log: torch.Tensor, + gate_bias: torch.Tensor, + initial_state: torch.Tensor, + *, + lower_bound: float = -5.0, + inplace_final_state: bool = True, + cu_seqlens: torch.Tensor | None = None, + ssm_state_indices: torch.Tensor | None = None, + ssm_state_write_indices: torch.Tensor | None = None, + num_accepted_tokens: torch.Tensor | None = None, + out: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run GLM-5 KDA with bounded decay and sigmoid beta.""" + + if ENABLE_FAST_MTP_KDA: + value_heads = v.shape[2] + key_dim = q.shape[-1] + assert a_log.numel() == value_heads + assert gate_bias.numel() == value_heads * key_dim + # Keep the raw projections in their packed token-major views and form + # the GLM-5 safe gate plus beta inside the recurrent kernel. This + # removes the fp32 cast/add/exp/sigmoid/multiply launches from every KDA + # layer, which is especially important for the model's 34 KDA layers. + return fused_recurrent_gated_delta_rule( + q=q, + k=k, + v=v, + initial_state=initial_state, + inplace_final_state=inplace_final_state, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + ssm_state_write_indices=ssm_state_write_indices, + num_accepted_tokens=num_accepted_tokens, + use_qk_l2norm_in_kernel=True, + A_log=a_log, + dt_bias=gate_bias, + a_raw=raw_gate.reshape(-1, value_heads * key_dim), + b_raw=raw_beta.reshape(-1, value_heads), + out=out, + is_kda=True, + kda_lower_bound=lower_bound, + ) + + gate = kda_safe_gate(raw_gate, a_log, gate_bias, lower_bound) + beta = raw_beta.float().sigmoid() + return fused_recurrent_gated_delta_rule( + q=q, + k=k, + v=v, + g=gate, + beta=beta, + initial_state=initial_state, + inplace_final_state=inplace_final_state, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + ssm_state_write_indices=ssm_state_write_indices, + num_accepted_tokens=num_accepted_tokens, + use_qk_l2norm_in_kernel=True, + out=out, + is_kda=True, + ) + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[ + triton.Config({"BK": BK}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["BC"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( + q, + k, + g, + beta, + A, + Aqk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + NC: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + i_i, i_j = i_c // NC, i_c % NC + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT + i_i * BC >= T: + return + if i_i <= i_j: + return + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + g += (bos * H + i_h) * K + A += (bos * H + i_h) * BT + Aqk += (bos * H + i_h) * BT + + p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)) + + b_A = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk = tl.zeros([BC, BC], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_g = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + b_kt = tl.make_block_ptr(k, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) + p_gk = tl.make_block_ptr(g, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) + + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + # [BK,] + b_gn = tl.load(g + (i_t * BT + i_i * BC) * H * K + o_k, mask=m_k, other=0) + # [BC, BK] + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) * exp2(b_g - b_gn[None, :]) + # [BK, BC] + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kt = tl.load(b_kt, boundary_check=(0, 1)) + # [BC, BC] + b_ktg = b_kt * exp2(b_gn[:, None] - b_gk) + b_A += tl.dot(b_k, b_ktg) + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_qg = b_q * exp2(b_g - b_gn[None, :]) * scale + b_Aqk += tl.dot(b_qg, b_ktg) + + b_A *= b_b[:, None] + + p_A = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) + tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) + p_Aqk = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[triton.Config({}, num_warps=num_warps) for num_warps in [1, 2, 4, 8]], + key=["BK", "BT"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( + q, + k, + g, + beta, + A, + Aqk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT + i_i * BC >= T: + return + + o_i = tl.arange(0, BC) + o_k = tl.arange(0, BK) + m_k = o_k < K + m_A = (i_t * BT + i_i * BC + o_i) < T + o_A = (bos + i_t * BT + i_i * BC + o_i) * H * BT + i_h * BT + i_i * BC + + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT + i_i * BC, 0), + (BC, BK), + (1, 0), + ) + p_k = tl.make_block_ptr( + k + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT + i_i * BC, 0), + (BC, BK), + (1, 0), + ) + p_g = tl.make_block_ptr( + g + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT + i_i * BC, 0), + (BC, BK), + (1, 0), + ) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + + p_b = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h + b_k = b_k * tl.load(p_b, mask=m_A, other=0)[:, None] + + p_kt = k + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k + p_gk = g + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k + + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + b_kt = tl.load(p_kt, mask=m_k, other=0).to(tl.float32) + b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) + b_ktg = b_kt[None, :] * exp2(b_g - b_gk[None, :]) + b_A = tl.sum(b_k * b_ktg, 1) + b_A = tl.where(o_i > j, b_A, 0.0) + b_Aqk = tl.sum(b_q * b_ktg, 1) + b_Aqk = tl.where(o_i >= j, b_Aqk * scale, 0.0) + tl.store(A + o_A + j, b_A, mask=m_A) + tl.store(Aqk + o_A + j, b_Aqk, mask=m_A) + p_kt += H * K + p_gk += H * K + + +def chunk_kda_scaled_dot_kkt_fwd( + q: torch.Tensor, + k: torch.Tensor, + gk: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + scale: float | None = None, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, + output_dtype: torch.dtype = torch.float32, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Compute beta * K * K^T. + + Args: + k (torch.Tensor): + The key tensor of shape `[B, T, H, K]`. + beta (torch.Tensor): + The beta tensor of shape `[B, T, H]`. + gk (torch.Tensor): + The cumulative sum of the gate tensor of shape `[B, T, H, K]` applied to the key tensor. Default: `None`. + cu_seqlens (torch.Tensor): + The cumulative sequence lengths of the input tensor. + Default: None + chunk_size (int): + The chunk size. Default: 64. + output_dtype (torch.dtype): + The dtype of the output tensor. Default: `torch.float32` + + Returns: + beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size. + """ + B, T, H, K = k.shape + assert K <= 256 + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + BC = min(16, BT) + NC = cdiv(BT, BC) + BK = max(next_power_of_2(K), 16) + A = torch.zeros(B, T, H, BT, device=k.device, dtype=output_dtype) + Aqk = torch.zeros(B, T, H, BT, device=k.device, dtype=output_dtype) + grid = (NT, NC * NC, B * H) + chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter[grid]( + q=q, + k=k, + g=gk, + beta=beta, + A=A, + Aqk=Aqk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + NC=NC, + ) + + grid = (NT, NC, B * H) + chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra[grid]( + q=q, + k=k, + g=gk, + beta=beta, + A=A, + Aqk=Aqk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + BK=BK, + ) + return A, Aqk + + +@triton.heuristics( + { + "STORE_QG": lambda args: args["qg"] is not None, + "STORE_KG": lambda args: args["kg"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"], +) +@triton.jit(do_not_specialize=["T"]) +def recompute_w_u_fwd_kernel( + q, + k, + qg, + kg, + v, + beta, + w, + u, + A, + gk, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + STORE_QG: tl.constexpr, + STORE_KG: tl.constexpr, + IS_VARLEN: tl.constexpr, + DOT_PRECISION: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)) + + p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + b_A = tl.load(p_A, boundary_check=(0, 1)) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_u = tl.make_block_ptr( + u + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = (b_v * b_b[:, None]).to(b_v.dtype) + b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + + for i_k in range(tl.cdiv(K, BK)): + p_w = tl.make_block_ptr( + w + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_k = tl.make_block_ptr( + k + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = b_k * b_b[:, None] + + p_gk = tl.make_block_ptr( + gk + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kb *= exp2(b_gk) + if STORE_QG: + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_qg = tl.make_block_ptr( + qg + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_qg = b_q * exp2(b_gk) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + if STORE_KG: + last_idx = min(i_t * BT + BT, T) - 1 + + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + b_gn = tl.load(gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0) + b_kg = b_k * exp2(b_gn - b_gk) + + p_kg = tl.make_block_ptr( + kg + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + + b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + + +def recompute_w_u_fwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + q: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = A.shape[-1] + BK = 64 + BV = 64 + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + w = torch.empty_like(k) + u = torch.empty_like(v) + kg = torch.empty_like(k) if gk is not None else None + recompute_w_u_fwd_kernel[(NT, B * H)]( + q=q, + k=k, + qg=None, + kg=kg, + v=v, + beta=beta, + w=w, + u=u, + A=A, + gk=gk, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + DOT_PRECISION="ieee", + ) + return w, u, None, kg + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[ + triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for BV in [64, 128] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["BT"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_gla_fwd_kernel_o( + q, + v, + g, + h, + o, + A, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_g = tl.make_block_ptr( + g + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_h = tl.make_block_ptr( + h + (i_tg * H + i_h) * K * V, + (K, V), + (V, 1), + (i_k * BK, i_v * BV), + (BK, BV), + (1, 0), + ) + + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BT, BK] + b_g = tl.load(p_g, boundary_check=(0, 1)) + # [BT, BK] + b_qg = (b_q * exp2(b_g)).to(b_q.dtype) + # [BV, BK] + b_h = tl.load(p_h, boundary_check=(0, 1)) + # [BT, BV] + if i_k >= 0: + b_o += tl.dot(b_qg, b_h.to(b_qg.dtype)) + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, BT] + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) + b_o += tl.dot(b_A, b_v, allow_tf32=False) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_gla_fwd_o_gk( + q: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + A: torch.Tensor, + h: torch.Tensor, + o: torch.Tensor, + scale: float, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, +): + B, T, H, K, V = *q.shape, v.shape[-1] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + def grid(meta): + return (cdiv(V, meta["BV"]), NT, B * H) + + chunk_gla_fwd_kernel_o[grid]( + q=q, + v=v, + g=g, + h=h, + o=o, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + ) + return o + + +@triton.heuristics( + { + "HAS_BIAS": lambda args: args["g_bias"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[triton.Config({"BD": BD}, num_warps=num_warps) for BD in [32, 64] for num_warps in [2, 4, 8]], + key=["H", "D", "BT", "IS_VARLEN"], +) +@triton.jit(do_not_specialize=["T"]) +def kda_gate_cumsum_fwd_kernel( + g, + A, + y, + g_bias, + cu_seqlens, + chunk_indices, + cumsum_scale, + beta, + threshold, + SAFE_GATE: tl.constexpr, + LOWER_BOUND: tl.constexpr, + T, + H: tl.constexpr, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + HAS_BIAS: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos = i_b * T + + p_g = tl.make_block_ptr( + g + (bos * H + i_h) * D, + (T, D), + (H * D, 1), + (i_t * BT, i_d * BD), + (BT, BD), + (1, 0), + ) + p_y = tl.make_block_ptr( + y + (bos * H + i_h) * D, + (T, D), + (H * D, 1), + (i_t * BT, i_d * BD), + (BT, BD), + (1, 0), + ) + + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + if HAS_BIAS: + o_d = i_d * BD + tl.arange(0, BD) + b_bias = tl.load(g_bias + i_h * D + o_d, mask=o_d < D, other=0.0).to(tl.float32) + b_g = b_g + b_bias[None, :] + + b_a = tl.load(A + i_h).to(tl.float32) + b_a = tl.exp(b_a) if SAFE_GATE else -tl.exp(b_a) + if SAFE_GATE: + # y = lower_bound * sigmoid(exp(A) * (g + g_bias)); bounded to + # (lower_bound, 0). Mirrors the SGlang safe_gate branch used by GLM5-Next + # checkpoints whose linear_attn_config["safe_gate"] is True. + b_gate = LOWER_BOUND / (1.0 + tl.exp(-(b_a * b_g))) + else: + b_g_scaled = b_g * beta + b_softplus = tl.where( + b_g_scaled > threshold, + b_g, + (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)), + ) + b_gate = b_a * b_softplus + + # Out-of-bounds rows (load returns 0, but softplus/bias can still make + # b_gate non-zero) participate in the dot product. They only contribute to + # out-of-bounds output rows, which are masked away by `boundary_check` on + # the store, so visible output matches unfused gate + chunk-local cumsum. + o_t = tl.arange(0, BT) + m_cumsum = tl.where(o_t[:, None] >= o_t[None, :], 1.0, 0.0) + b_y = tl.dot(m_cumsum, b_gate, allow_tf32=False) * cumsum_scale + tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + + +def fused_kda_gate_chunk_cumsum( + raw_g: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None = None, + beta: float = 1.0, + threshold: float = 20.0, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, + output_dtype: torch.dtype | None = torch.float, + safe_gate: bool = False, + lower_bound: float = -5.0, +) -> torch.Tensor: + if cu_seqlens is not None: + assert raw_g.shape[0] == 1, "Only batch size 1 is supported when cu_seqlens are provided" + B, T, H, D = raw_g.shape + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = cdiv(T, chunk_size) if cu_seqlens is None else len(chunk_indices) + + A_log = A_log.reshape(-1) + if g_bias is not None: + g_bias = g_bias.reshape(-1) + y = torch.empty_like(raw_g, dtype=output_dtype or raw_g.dtype) + + def grid(meta): + return (cdiv(meta["D"], meta["BD"]), NT, B * H) + + kda_gate_cumsum_fwd_kernel[grid]( + g=raw_g, + A=A_log, + y=y, + g_bias=g_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + # RCP_LN2 folds in the natural-log -> log2 conversion so downstream + # exp2-based kernels reproduce exp(g). Keep this in sync with the + # `use_exp2=True` path in `_chunk_kda_fwd_with_cumulative_g`. + cumsum_scale=RCP_LN2, + beta=beta, + threshold=threshold, + SAFE_GATE=safe_gate, + LOWER_BOUND=lower_bound, + T=T, + H=H, + D=D, + BT=chunk_size, + ) + return y + + +def _chunk_kda_fwd_with_cumulative_g( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, +): + # `g` must already be chunk-local cumulatively-summed AND scaled by + # RCP_LN2 (so the downstream exp2-based kernels reproduce exp(g)). + # Use `chunk_kda_fwd` or `chunk_kda_with_fused_gate_fwd` instead of + # calling this helper directly unless that invariant is upheld. + # the intra Aqk is kept in fp32 + # the computation has very marginal effect on the entire throughput + A, Aqk = chunk_kda_scaled_dot_kkt_fwd( + q=q, + k=k, + gk=g, + beta=beta, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + output_dtype=torch.float32, + ) + A = solve_tril( + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + output_dtype=k.dtype, + ) + w, u, _, kg = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + gk=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + del A + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=kg, + w=w, + u=u, + gk=g, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + use_exp2=True, + ) + del w, u, kg + o = chunk_gla_fwd_o_gk( + q=q, + v=v_new, + g=g, + A=Aqk, + h=h, + o=v, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + del Aqk, v_new, h + return o, final_state + + +def chunk_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None = None, +): + chunk_size = FLA_CHUNK_SIZE + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None + g = chunk_local_cumsum( + g, + chunk_size=chunk_size, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + # KDA evaluates cumulative gate decays with exp2. Convert from natural-log + # space so exp(x) is preserved as exp2(x / ln(2)). + g = g * RCP_LN2 + return _chunk_kda_fwd_with_cumulative_g( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + + +def chunk_kda_with_fused_gate_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + safe_gate: bool = False, + lower_bound: float = -5.0, +): + chunk_size = FLA_CHUNK_SIZE + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + g = fused_kda_gate_chunk_cumsum( + raw_g, + A_log=A_log, + g_bias=g_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + safe_gate=safe_gate, + lower_bound=lower_bound, + ) + return _chunk_kda_fwd_with_cumulative_g( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + + +def chunk_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.Tensor | None = None, + **kwargs, +): + if scale is None: + scale = k.shape[-1] ** -0.5 + + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q.contiguous()) + k = l2norm_fwd(k.contiguous()) + + o, final_state = chunk_kda_fwd( + q=q, + k=k, + v=v.contiguous(), + g=g.contiguous(), + beta=beta.contiguous(), + scale=scale, + initial_state=initial_state.contiguous(), + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + return o, final_state + + +def chunk_kda_with_fused_gate( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + safe_gate: bool = False, + lower_bound: float = -5.0, + **kwargs, +): + """Run chunk KDA from raw gate projection using fused gate+cumsum.""" + if scale is None: + scale = k.shape[-1] ** -0.5 + + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q.contiguous()) + k = l2norm_fwd(k.contiguous()) + + o, final_state = chunk_kda_with_fused_gate_fwd( + q=q, + k=k, + v=v.contiguous(), + raw_g=raw_g.contiguous(), + beta=beta.contiguous(), + A_log=A_log, + g_bias=g_bias, + scale=scale, + initial_state=initial_state.contiguous() if initial_state is not None else None, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + safe_gate=safe_gate, + lower_bound=lower_bound, + ) + return o, final_state diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/op.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/op.py index 2f69aa981d..6bfb1db11c 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/op.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/op.py @@ -13,6 +13,7 @@ from .utils import is_gather_supported exp = tl.exp +exp2 = tl.exp2 log = tl.log log2 = tl.log2 diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py index b5b6cfc369..8b0293bfd5 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py @@ -411,6 +411,7 @@ def merge_16x16_to_64x64_inverse_kernel( def solve_tril( A: torch.Tensor, cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, output_dtype: torch.dtype = torch.float, ) -> torch.Tensor: """ @@ -433,7 +434,8 @@ def solve_tril( output_dtype = A.dtype if output_dtype is None else output_dtype B, T, H, BT = A.shape - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT) Ai = torch.zeros_like(A, dtype=output_dtype) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/mtp_fused_recurrent.py b/lightllm/common/basemodel/triton_kernel/linear_att/mtp_fused_recurrent.py index 2eb5ef5333..37642369fc 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/mtp_fused_recurrent.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/mtp_fused_recurrent.py @@ -15,6 +15,46 @@ import triton import triton.language as tl +from lightllm.common.triton_utils.autotuner import autotune + + +_MTP_RECURRENT_BV_SIZES = (4, 8, 16, 32, 64) + + +def _get_mtp_fused_recurrent_configs(): + return [ + {"BV": bv, "num_warps": num_warps, "num_stages": num_stages} + for bv in _MTP_RECURRENT_BV_SIZES + for num_warps in (1, 2, 4, 8) + for num_stages in (1, 2, 3) + ] + + +def _get_mtp_fused_recurrent_static_key(q, v, initial_state, fixed_seq_len): + return { + "H": q.shape[2], + "HV": v.shape[2], + "K": q.shape[3], + "V": v.shape[3], + "dtype": str(q.dtype), + "state_dtype": str(initial_state.dtype), + "fixed_seq_len": int(fixed_seq_len), + } + + +def _get_mtp_fused_recurrent_run_key(q, cu_seqlens): + sequence_count = int(cu_seqlens.shape[0] - 1) + total_tokens = int(q.shape[1]) + return sequence_count * 10 ** 9 + total_tokens + + +def _default_mtp_fused_recurrent_config(V): + return { + "BV": min(triton.next_power_of_2(V), 8), + "num_warps": 1, + "num_stages": 3, + } + # --------------------------------------------------------------------------- # Triton kernel @@ -59,14 +99,17 @@ def _fused_recurrent_gated_delta_rule_fwd_kernel( stride_write_indices_tok: tl.constexpr, SOFTPLUS_BETA: tl.constexpr, SOFTPLUS_THRESHOLD: tl.constexpr, + FIXED_SEQ_LEN: tl.constexpr, ): i_v, i_n, i_hv = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_h = i_hv // (HV // H) - bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int64), - tl.load(cu_seqlens + i_n + 1).to(tl.int64), - ) - T = eos - bos + if FIXED_SEQ_LEN > 0: + bos = i_n * FIXED_SEQ_LEN + T: tl.constexpr = FIXED_SEQ_LEN + else: + bos = tl.load(cu_seqlens + i_n).to(tl.int64) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos if T == 0: return @@ -133,11 +176,18 @@ def _fused_recurrent_gated_delta_rule_fwd_kernel( # --------------------------------------------------------------------------- -# Public API — directly launches the triton kernel (no autograd.Function) +# Autotuned implementation — directly launches the Triton kernel # --------------------------------------------------------------------------- -def mtp_fused_recurrent_gated_delta_rule( +@autotune( + kernel_name="mtp_fused_recurrent_gated_delta_rule:v1", + configs_gen_func=_get_mtp_fused_recurrent_configs, + static_key_func=_get_mtp_fused_recurrent_static_key, + run_key_func=_get_mtp_fused_recurrent_run_key, + mutates_args=["initial_state"], +) +def _mtp_fused_recurrent_gated_delta_rule_autotuned( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, @@ -150,6 +200,8 @@ def mtp_fused_recurrent_gated_delta_rule( dt_bias: torch.Tensor, a_raw: torch.Tensor, b_raw: torch.Tensor, + fixed_seq_len: int, + run_config: dict = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Fused recurrent gated delta rule with fused gating (GDN layer). @@ -170,6 +222,8 @@ def mtp_fused_recurrent_gated_delta_rule( dt_bias: ``[HV]`` per-head dt bias. a_raw: ``[T, HV]`` raw alpha. b_raw: ``[T, HV]`` raw beta. + fixed_seq_len: Compile-time sequence length for dense fixed-width MTP + verification. Zero keeps the variable-length ``cu_seqlens`` path. Returns: ``(o, final_state)`` where ``o`` is ``[1, T, HV, V]`` and @@ -184,6 +238,12 @@ def mtp_fused_recurrent_gated_delta_rule( V = v.shape[-1] HV = v.shape[2] N = len(cu_seqlens) - 1 + fixed_seq_len = int(fixed_seq_len) + assert fixed_seq_len >= 0 + if fixed_seq_len: + assert q.shape[1] == N * fixed_seq_len, ( + f"fixed_seq_len={fixed_seq_len} requires {N * fixed_seq_len} tokens, " f"got {q.shape[1]}" + ) q, stride_q_tok = _ensure_qkv_token_strided(q) k, stride_k_tok = _ensure_qkv_token_strided(k) v, stride_v_tok = _ensure_qkv_token_strided(v) @@ -191,17 +251,18 @@ def mtp_fused_recurrent_gated_delta_rule( b_raw, stride_b_tok = _ensure_gate_token_strided(b_raw) BK = triton.next_power_of_2(K) assert K == BK, f"K={K} must be a power of 2" - BV = min(triton.next_power_of_2(V), 8) - num_warps = 1 - num_stages = 3 + if run_config is None: + run_config = _default_mtp_fused_recurrent_config(V) + BV = run_config["BV"] + num_warps = run_config.get("num_warps", 1) + num_stages = run_config["num_stages"] NV = triton.cdiv(V, BV) - o = q.new_empty(v.shape) + output = q.new_empty(v.shape) final_state = initial_state - stride_init_state_token = initial_state.stride(0) stride_final_state_token = final_state.stride(0) - stride_o_tok = o.stride(1) + stride_o_tok = output.stride(1) assert stride_o_tok == HV * V, f"stride_o_tok={stride_o_tok} must be HV*V" stride_state_hv = K * V @@ -216,7 +277,7 @@ def mtp_fused_recurrent_gated_delta_rule( q=q, k=k, v=v, - o=o, + o=output, h0=initial_state, ht=final_state, cu_seqlens=cu_seqlens, @@ -249,10 +310,49 @@ def mtp_fused_recurrent_gated_delta_rule( stride_write_indices_tok=stride_write_indices_tok, SOFTPLUS_BETA=1.0, SOFTPLUS_THRESHOLD=20.0, + FIXED_SEQ_LEN=fixed_seq_len, num_warps=num_warps, num_stages=num_stages, ) - return o, final_state + return output, final_state + + +def mtp_fused_recurrent_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, + ssm_state_indices: torch.Tensor, + ssm_state_write_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + a_raw: torch.Tensor, + b_raw: torch.Tensor, + fixed_seq_len: int = 0, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run the autotuned fused recurrent GDN kernel. + + ``fixed_seq_len=0`` preserves the variable-length behavior of the original + public API. The private autotuned implementation always receives the value + explicitly so its cache key does not depend on generic default handling. + """ + return _mtp_fused_recurrent_gated_delta_rule_autotuned( + q=q, + k=k, + v=v, + initial_state=initial_state, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + ssm_state_write_indices=ssm_state_write_indices, + num_accepted_tokens=num_accepted_tokens, + A_log=A_log, + dt_bias=dt_bias, + a_raw=a_raw, + b_raw=b_raw, + fixed_seq_len=fixed_seq_len, + ) # --------------------------------------------------------------------------- diff --git a/lightllm/common/basemodel/triton_kernel/norm/gated_rmsnorm.py b/lightllm/common/basemodel/triton_kernel/norm/gated_rmsnorm.py index c62c5eb5d2..59f5ab2434 100644 --- a/lightllm/common/basemodel/triton_kernel/norm/gated_rmsnorm.py +++ b/lightllm/common/basemodel/triton_kernel/norm/gated_rmsnorm.py @@ -18,20 +18,27 @@ def gated_rmsnorm_forward_kernel( Z, # pointer to the other branch (required, not optional) stride_x_row, # how much to increase the pointer when moving by 1 row stride_y_row, - stride_z_row, + stride_z_token, + stride_z_head, M, # number of rows in X N, # number of columns in X eps, # epsilon to avoid division by zero BLOCK_N: tl.constexpr, HAS_BIAS: tl.constexpr, NORM_BEFORE_GATE: tl.constexpr, + SIGMOID_GATE: tl.constexpr, + Z_HEADS: tl.constexpr, ): # Map the program id to the row of X and Y it should compute. row = tl.program_id(0) group = tl.program_id(1) X += row * stride_x_row + group * N Y += row * stride_y_row + group * N - Z += row * stride_z_row + group * N + # X is flattened to [tokens * heads, N], while GDN can keep Z as a + # zero-copy [tokens, heads, N] slice of its packed projection output. + z_token = row // Z_HEADS + z_head = row % Z_HEADS + Z += z_token * stride_z_token + z_head * stride_z_head + group * N W += group * N if HAS_BIAS: B += group * N @@ -40,7 +47,7 @@ def gated_rmsnorm_forward_kernel( x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) if not NORM_BEFORE_GATE: z = tl.load(Z + cols, mask=cols < N).to(tl.float32) - x *= z * tl.sigmoid(z) + x *= tl.sigmoid(z) if SIGMOID_GATE else z * tl.sigmoid(z) # RMS norm: compute variance directly without mean subtraction xbar = tl.where(cols < N, x, 0.0) var = tl.sum(xbar * xbar, axis=0) / N @@ -55,7 +62,7 @@ def gated_rmsnorm_forward_kernel( y = x_hat * w + b if HAS_BIAS else x_hat * w if NORM_BEFORE_GATE: z = tl.load(Z + cols, mask=mask).to(tl.float32) - y *= z * tl.sigmoid(z) + y *= tl.sigmoid(z) if SIGMOID_GATE else z * tl.sigmoid(z) # Write output tl.store(Y + cols, y, mask=mask) @@ -102,6 +109,7 @@ def gated_rmsnorm_forward( out: torch.Tensor = None, group_size: int = None, norm_before_gate: bool = True, + activation: str = "silu", run_config: dict = None, ): M, N = x.shape @@ -112,8 +120,21 @@ def gated_rmsnorm_forward( assert x.stride(-1) == 1 # z is required for gated_rmsnorm assert z is not None, "z cannot be None for gated_rmsnorm_forward" + assert activation in ("silu", "sigmoid"), f"unsupported gate activation: {activation}" + # Accept GDN's strided 3D gate without materializing a flattened copy. + assert z.ndim in (2, 3), f"z must be [M, N] or [tokens, heads, N], got shape={z.shape}" assert z.stride(-1) == 1 - assert z.shape == (M, N) + if z.ndim == 2: + assert z.shape == (M, N) + z_heads = 1 + stride_z_token = z.stride(0) + stride_z_head = 0 + else: + assert z.shape[-1] == N, f"z.shape[-1]={z.shape[-1]} must match N={N}" + assert z.shape[0] * z.shape[1] == M, f"z token/head rows={z.shape[0] * z.shape[1]} must match M={M}" + z_heads = z.shape[1] + stride_z_token = z.stride(0) + stride_z_head = z.stride(1) assert weight.shape == (N,) assert weight.stride(-1) == 1 if bias is not None: @@ -156,12 +177,15 @@ def gated_rmsnorm_forward( z, x.stride(0), out.stride(0), - z.stride(0), + stride_z_token, + stride_z_head, M, group_size, eps, BLOCK_N=BLOCK_N, NORM_BEFORE_GATE=norm_before_gate, + SIGMOID_GATE=activation == "sigmoid", + Z_HEADS=z_heads, num_warps=num_warps, ) return out diff --git a/lightllm/common/basemodel/triton_kernel/post_process/greedy_sample.py b/lightllm/common/basemodel/triton_kernel/post_process/greedy_sample.py new file mode 100644 index 0000000000..8e7e99e1ab --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/post_process/greedy_sample.py @@ -0,0 +1,109 @@ +"""Local greedy statistics for distributed vocabulary shards.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _greedy_sample_stage1_kernel( + logits, + partial_max, + partial_sum, + partial_argmax, + stride_row, + stride_col, + vocab_size: tl.constexpr, + num_blocks: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0) + block = tl.program_id(1) + offsets = block * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + values = tl.load( + logits + row * stride_row + offsets * stride_col, + mask=offsets < vocab_size, + other=-float("inf"), + ) + values = values.to(tl.float32) + + block_max = tl.max(values, axis=0) + block_sum = tl.sum(tl.exp(values - block_max), axis=0) + block_argmax = tl.argmax(values, axis=0) + block * BLOCK_SIZE + output_offset = row * num_blocks + block + tl.store(partial_max + output_offset, block_max) + tl.store(partial_sum + output_offset, block_sum) + tl.store(partial_argmax + output_offset, block_argmax) + + +@triton.jit +def _greedy_sample_stage2_stats_kernel( + partial_max, + partial_sum, + partial_argmax, + output_stats, + output_argmax, + num_blocks: tl.constexpr, + batch_size: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0) + offsets = tl.arange(0, BLOCK_SIZE) + mask = offsets < num_blocks + input_offset = row * num_blocks + offsets + block_max = tl.load(partial_max + input_offset, mask=mask, other=-float("inf")) + block_sum = tl.load(partial_sum + input_offset, mask=mask, other=0.0) + block_argmax = tl.load(partial_argmax + input_offset, mask=mask, other=0x7FFFFFFF) + + global_max = tl.max(block_max, axis=0) + global_sum = tl.sum(block_sum * tl.exp(block_max - global_max), axis=0) + candidate_ids = tl.where(block_max == global_max, block_argmax, 0x7FFFFFFF) + global_argmax = tl.min(candidate_ids, axis=0) + tl.store(output_stats + row, global_max) + tl.store(output_stats + batch_size + row, global_max + tl.log(global_sum)) + tl.store(output_argmax + row, global_argmax) + + +def _launch_stage1(logits: torch.Tensor, scratch: torch.Tensor, block_size: int, num_blocks: int) -> None: + batch_size, vocab_size = logits.shape + _greedy_sample_stage1_kernel[(batch_size, num_blocks)]( + logits, + scratch[0], + scratch[1], + scratch[2], + logits.stride(0), + logits.stride(1), + vocab_size=vocab_size, + num_blocks=num_blocks, + BLOCK_SIZE=block_size, + num_warps=8, + ) + + +@torch.no_grad() +def greedy_sample_local_stats(logits: torch.Tensor, alloc_func=torch.empty) -> torch.Tensor: + """Return local max, logsumexp and argmax rows for distributed greedy sampling.""" + + assert logits.ndim == 2 and logits.is_cuda and logits.is_contiguous() + batch_size, vocab_size = logits.shape + block_size = 4096 + num_blocks = triton.cdiv(vocab_size, block_size) + scratch = alloc_func((3, batch_size, num_blocks), dtype=torch.float32, device=logits.device) + # The third FP32 row carries INT32 argmax bits. Keeping one fixed-size + # payload gives the distributed reducer a single collective without losing + # token-id precision through a numeric int-to-float conversion. + output_stats = alloc_func((3, batch_size), dtype=torch.float32, device=logits.device) + + _launch_stage1(logits, scratch, block_size, num_blocks) + _greedy_sample_stage2_stats_kernel[(batch_size,)]( + scratch[0], + scratch[1], + scratch[2], + output_stats, + output_stats[2].view(torch.int32), + num_blocks=num_blocks, + batch_size=batch_size, + BLOCK_SIZE=triton.next_power_of_2(num_blocks), + num_warps=4, + ) + return output_stats diff --git a/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_greedy.py b/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_greedy.py new file mode 100644 index 0000000000..d7dea3e861 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_greedy.py @@ -0,0 +1,126 @@ +"""Greedy sampling directly from tensor-parallel vocabulary shards.""" + +import torch +import triton +import triton.language as tl + +from lightllm.common.basemodel.triton_kernel.post_process.greedy_sample import ( + greedy_sample_local_stats, +) +from lightllm.common.basemodel.triton_kernel.transpose_convert import ( + transpose_convert_2d, +) +from lightllm.distributed.communication_op import all_gather_into_tensor +from lightllm.utils.envs_utils import enable_env_vars + + +VOCAB_PARALLEL_GREEDY_ENV = "LIGHTLLM_VOCAB_PARALLEL_GREEDY" + + +def is_vocab_parallel_greedy_enabled() -> bool: + return enable_env_vars(VOCAB_PARALLEL_GREEDY_ENV) + + +@triton.jit +def _combine_vocab_parallel_stats_kernel( + gathered_stats, + gathered_argmax, + output_logits, + output_token_ids, + output_logsumexp, + token_num, + vocab_size: tl.constexpr, + tp_world_size: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + token_offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + token_mask = token_offsets < token_num + rank_stride = 3 * token_num + + global_max = tl.full((BLOCK_SIZE,), -float("inf"), tl.float32) + global_id = tl.full((BLOCK_SIZE,), 0x7FFFFFFF, tl.int32) + for rank in tl.static_range(tp_world_size): + rank_base = rank * rank_stride + local_max = tl.load( + gathered_stats + rank_base + token_offsets, + mask=token_mask, + other=-float("inf"), + ) + local_id = tl.load( + gathered_argmax + rank_base + 2 * token_num + token_offsets, + mask=token_mask, + other=0x7FFFFFFF, + ) + local_id += (rank * vocab_size) // tp_world_size + wins = (local_max > global_max) | ((local_max == global_max) & (local_id < global_id)) + global_max = tl.where(wins, local_max, global_max) + global_id = tl.where(wins, local_id, global_id) + + global_sum = tl.zeros((BLOCK_SIZE,), tl.float32) + for rank in tl.static_range(tp_world_size): + rank_base = rank * rank_stride + local_lse = tl.load( + gathered_stats + rank_base + token_num + token_offsets, + mask=token_mask, + other=-float("inf"), + ) + global_sum += tl.exp(local_lse - global_max) + + tl.store(output_logits + token_offsets, global_max, mask=token_mask) + tl.store(output_token_ids + token_offsets, global_id, mask=token_mask) + tl.store(output_logsumexp + token_offsets, global_max + tl.log(global_sum), mask=token_mask) + + +@torch.no_grad() +def vocab_parallel_greedy( + local_logits: torch.Tensor, + *, + vocab_size: int, + tp_world_size: int, + group, + alloc_func, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return exact sparse logits, global token ids and full-vocab logsumexp.""" + + assert local_logits.ndim == 2 and local_logits.is_cuda and local_logits.is_contiguous() + local_vocab_size, token_num = local_logits.shape + assert local_vocab_size in { + vocab_size // tp_world_size, + (vocab_size + tp_world_size - 1) // tp_world_size, + } + + transposed_logits = alloc_func( + (token_num, local_vocab_size), + dtype=local_logits.dtype, + device=local_logits.device, + ) + transpose_convert_2d(local_logits, transposed_logits) + local_stats = greedy_sample_local_stats(transposed_logits, alloc_func=alloc_func) + + if tp_world_size == 1: + gathered_stats = local_stats.view(1, 3, token_num) + else: + gathered_stats = alloc_func((tp_world_size, 3, token_num), dtype=torch.float32, device=local_logits.device) + all_gather_into_tensor( + output_=gathered_stats, + input_=local_stats, + group=group, + async_op=False, + ) + + output_logits = alloc_func((token_num, 1), dtype=torch.float32, device=local_logits.device) + output_token_ids = alloc_func((token_num, 1), dtype=torch.int64, device=local_logits.device) + output_logsumexp = alloc_func((token_num,), dtype=torch.float32, device=local_logits.device) + _combine_vocab_parallel_stats_kernel[(triton.cdiv(token_num, 256),)]( + gathered_stats, + gathered_stats.view(torch.int32), + output_logits, + output_token_ids, + output_logsumexp, + token_num, + vocab_size=vocab_size, + tp_world_size=tp_world_size, + BLOCK_SIZE=256, + num_warps=4, + ) + return output_logits, output_token_ids, output_logsumexp diff --git a/lightllm/common/basemodel/triton_kernel/transpose_convert.py b/lightllm/common/basemodel/triton_kernel/transpose_convert.py new file mode 100644 index 0000000000..b618d69526 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/transpose_convert.py @@ -0,0 +1,65 @@ +"""Tiled transpose kernels used by the post-layer logits path.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _transpose_convert_2d_kernel( + input_ptr, + output_ptr, + rows, + cols, + input_stride_0, + input_stride_1, + output_stride_0, + output_stride_1, + BLOCK_ROWS: tl.constexpr, + BLOCK_COLS: tl.constexpr, +): + row_offsets = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + col_offsets = tl.program_id(1) * BLOCK_COLS + tl.arange(0, BLOCK_COLS) + input_offsets = row_offsets[:, None] * input_stride_0 + col_offsets[None, :] * input_stride_1 + mask = (row_offsets[:, None] < rows) & (col_offsets[None, :] < cols) + values = tl.load(input_ptr + input_offsets, mask=mask) + + output_offsets = col_offsets[:, None] * output_stride_0 + row_offsets[None, :] * output_stride_1 + tl.store(output_ptr + output_offsets, tl.trans(values), mask=tl.trans(mask)) + + +@torch.no_grad() +def transpose_convert_2d( + input_tensor: torch.Tensor, + output_tensor: torch.Tensor, + *, + block_rows: int = 64, + block_cols: int = 64, + num_warps: int = 8, + num_stages: int = 1, +) -> torch.Tensor: + """Transpose a contiguous 2-D CUDA tensor while converting its dtype.""" + + assert input_tensor.is_cuda and output_tensor.is_cuda + assert input_tensor.device == output_tensor.device + assert input_tensor.ndim == 2 and output_tensor.ndim == 2 + assert output_tensor.shape == (input_tensor.shape[1], input_tensor.shape[0]) + assert input_tensor.is_contiguous() and output_tensor.is_contiguous() + + rows, cols = input_tensor.shape + grid = (triton.cdiv(rows, block_rows), triton.cdiv(cols, block_cols)) + _transpose_convert_2d_kernel[grid]( + input_tensor, + output_tensor, + rows, + cols, + input_tensor.stride(0), + input_tensor.stride(1), + output_tensor.stride(0), + output_tensor.stride(1), + BLOCK_ROWS=block_rows, + BLOCK_COLS=block_cols, + num_warps=num_warps, + num_stages=num_stages, + ) + return output_tensor diff --git a/lightllm/common/linear_att_cache_manager/config_objs.py b/lightllm/common/linear_att_cache_manager/config_objs.py index f588ec7d5c..adc90f9486 100644 --- a/lightllm/common/linear_att_cache_manager/config_objs.py +++ b/lightllm/common/linear_att_cache_manager/config_objs.py @@ -113,7 +113,14 @@ def load_from_args() -> "LinearAttCacheConfig": model_cfg, _ = PretrainedConfig.get_config_dict(model_path) model_type = model_cfg["model_type"] - assert model_type in ["qwen3_5", "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text"] + assert model_type in [ + "qwen3_5", + "qwen3_5_moe", + "qwen3_5_text", + "qwen3_5_moe_text", + "glm5_next", + "glm5_next_text", + ] llm_config = model_cfg try: llm_config = llm_config["text_config"] @@ -122,6 +129,30 @@ def load_from_args() -> "LinearAttCacheConfig": n_layer = llm_config["num_hidden_layers"] + if model_type in ["glm5_next", "glm5_next_text"]: + linear = llm_config["linear_attn_config"] + tp_world_size = get_env_start_args().tp // get_env_start_args().dp + return LinearAttCacheConfig( + tp_world_size=tp_world_size, + full_att_all_num_kv_heads=1, + full_att_dtype=get_torch_dtype(args.data_type), + full_att_num_kv_heads=1, + full_att_head_dim=llm_config["kv_lora_rank"], + global_linear_k_heads=linear["num_heads"], + global_linear_v_heads=linear["num_heads"], + num_linear_k_heads=linear["num_heads"] // tp_world_size, + num_linear_v_heads=linear["num_heads"] // tp_world_size, + head_linear_k_dim=linear["head_dim"], + head_linear_v_dim=linear["head_dim"], + conv_kernel_size=linear["short_conv_kernel_size"], + linear_layer_num=len(linear["kda_layers"]), + conv_state_dtype=get_torch_dtype(args.data_type), + ssm_state_dtype=get_torch_dtype(args.linear_att_ssm_data_type), + full_attention_interval=4, + all_layer_num=n_layer, + draft_full_att_kv_layer_num=get_added_mtp_kv_layer_num(), + ) + tp_world_size = get_env_start_args().tp // get_env_start_args().dp return LinearAttCacheConfig( tp_world_size=tp_world_size, diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/chunk_gated_delta_rule_fwd_h/{BT=64,H=8,K=128,V=128}_NVIDIA_H100_80GB_HBM3.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/chunk_gated_delta_rule_fwd_h/{BT=64,H=8,K=128,V=128}_NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 0000000000..dfbc9746cc --- /dev/null +++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/chunk_gated_delta_rule_fwd_h/{BT=64,H=8,K=128,V=128}_NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,7 @@ +{ + "8": { + "BV": 32, + "num_stages": 4, + "num_warps": 4 + } +} diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/silu_and_mul_fwd:v1/{N=1536,out_dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/silu_and_mul_fwd:v1/{N=1536,out_dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 0000000000..89c7d19be9 --- /dev/null +++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/silu_and_mul_fwd:v1/{N=1536,out_dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,68 @@ +{ + "1": { + "BLOCK_M": 256, + "BLOCK_N": 256, + "NUM_STAGES": 1, + "num_warps": 8 + }, + "1024": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 1, + "num_warps": 4 + }, + "128": { + "BLOCK_M": 1, + "BLOCK_N": 256, + "NUM_STAGES": 1, + "num_warps": 1 + }, + "16": { + "BLOCK_M": 1, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 4 + }, + "2048": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "256": { + "BLOCK_M": 1, + "BLOCK_N": 256, + "NUM_STAGES": 1, + "num_warps": 4 + }, + "32": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 2, + "num_warps": 1 + }, + "4": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 1, + "num_warps": 1 + }, + "4096": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "64": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 1, + "num_warps": 1 + }, + "8": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 1, + "num_warps": 1 + } +} diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/silu_and_mul_fwd:v1/{N=2048,out_dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/silu_and_mul_fwd:v1/{N=2048,out_dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 0000000000..9c0dd1629a --- /dev/null +++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.7.1/NVIDIA_H100_80GB_HBM3/silu_and_mul_fwd:v1/{N=2048,out_dtype=torch.bfloat16}_NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,278 @@ +{ + "1": { + "BLOCK_M": 128, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 8 + }, + "10112": { + "BLOCK_M": 64, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "1024": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "10240": { + "BLOCK_M": 64, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "10496": { + "BLOCK_M": 64, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "1152": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 1, + "num_warps": 1 + }, + "128": { + "BLOCK_M": 1, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "1280": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "1536": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "16": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 1, + "num_warps": 8 + }, + "1792": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "18304": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "18432": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "18560": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "1920": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "19200": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "2048": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "2176": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "2304": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "256": { + "BLOCK_M": 1, + "BLOCK_N": 256, + "NUM_STAGES": 2, + "num_warps": 1 + }, + "2816": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "2944": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "32": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 1, + "num_warps": 1 + }, + "3328": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "34048": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "34688": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "34944": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "35072": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "35584": { + "BLOCK_M": 1, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "3584": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 2, + "num_warps": 1 + }, + "3712": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "384": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "3968": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "4": { + "BLOCK_M": 1, + "BLOCK_N": 256, + "NUM_STAGES": 2, + "num_warps": 1 + }, + "4096": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 2, + "num_warps": 1 + }, + "4224": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 2, + "num_warps": 1 + }, + "4352": { + "BLOCK_M": 32, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "4480": { + "BLOCK_M": 32, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "4608": { + "BLOCK_M": 32, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "4736": { + "BLOCK_M": 32, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "512": { + "BLOCK_M": 8, + "BLOCK_N": 128, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "64": { + "BLOCK_M": 1, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "640": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 1, + "num_warps": 4 + }, + "768": { + "BLOCK_M": 8, + "BLOCK_N": 256, + "NUM_STAGES": 1, + "num_warps": 4 + }, + "8": { + "BLOCK_M": 1, + "BLOCK_N": 64, + "NUM_STAGES": 4, + "num_warps": 1 + }, + "9856": { + "BLOCK_M": 32, + "BLOCK_N": 256, + "NUM_STAGES": 4, + "num_warps": 1 + } +} diff --git a/lightllm/distributed/communication_op.py b/lightllm/distributed/communication_op.py index 93c603212d..5b67ffa424 100644 --- a/lightllm/distributed/communication_op.py +++ b/lightllm/distributed/communication_op.py @@ -38,7 +38,7 @@ create_new_group_for_current_dp, create_dp_special_inter_group, ) -from lightllm.utils.device_utils import get_device_sm_count, is_sm100_gpu +from lightllm.utils.device_utils import get_device_sm_count, is_sm90_gpu, is_sm100_gpu from lightllm.utils.torch_dtype_utils import get_torch_dtype logger = init_logger(__name__) @@ -57,6 +57,7 @@ class CustomProcessGroup: def __init__(self): self.symm_mem_reduce = None self.flashinfer_reduce = None + self.symm_mem_out_of_place = os.getenv("LIGHTLLM_SYMM_MEM_OUT_OF_PLACE", "0").upper() in {"1", "ON", "TRUE"} self.dp_world_size = get_dp_world_size() self.device_group = create_new_group_for_current_dp("nccl") if get_env_start_args().enable_dp_prefill_balance: @@ -97,7 +98,10 @@ def all_reduce(self, input_: torch.Tensor) -> None: input_.data = self.flashinfer_reduce.all_reduce(input_) return if self.symm_mem_reduce is not None and self.symm_mem_reduce.should_use(input_): - self.symm_mem_reduce.all_reduce(input_) + if self.symm_mem_out_of_place: + input_.data = self.symm_mem_reduce.all_reduce_out_of_place(input_) + else: + self.symm_mem_reduce.all_reduce(input_) return return dist.all_reduce(input_, group=self.device_group) @@ -109,6 +113,7 @@ class DistributeGroupManager: def __init__(self): self.groups = [] self.ep_buffer = None + self.ep_prefill_uses_legacy_buffer = False self.ep_low_latency_buffer = None self.ep_mega_moe_buffer = None self.ep_num_sms = None @@ -175,6 +180,7 @@ def new_deepep_group( decode_num_max_dispatch_tokens_per_rank = get_deepep_num_max_dispatch_tokens_per_rank_decode() if not enable_ep_moe: self.ep_buffer = None + self.ep_prefill_uses_legacy_buffer = False self.ep_low_latency_buffer = None self.ep_mega_moe_buffer = None self.ep_num_sms = None @@ -194,14 +200,38 @@ def new_deepep_group( self.ll_decode_num_tokens = decode_num_max_dispatch_tokens_per_rank self.ll_hidden = hidden_size self.ll_num_experts = n_routed_experts + get_redundancy_expert_num() * global_world_size - self.ep_buffer = deep_ep.ElasticBuffer( - deepep_group, - num_max_tokens_per_rank=self.ll_num_tokens, - hidden=self.ll_hidden, - num_topk=num_experts_per_tok, - use_fp8_dispatch=True, - allow_multiple_reduction=True, - ) + # ElasticBuffer uses NCCL GIN even for an intra-node group. GIN is not + # available on many otherwise fully-connected NVLink hosts. The legacy + # normal DeepEP buffer is the native intra-node transport and avoids that + # unnecessary network dependency; keep ElasticBuffer for multi-node EP. + self.ep_prefill_uses_legacy_buffer = get_env_start_args().nnodes == 1 + if self.ep_prefill_uses_legacy_buffer: + hidden_bytes = self.ll_hidden * get_torch_dtype(get_env_start_args().data_type).itemsize + dispatch_config = deep_ep.Buffer.get_dispatch_config(global_world_size) + combine_config = deep_ep.Buffer.get_combine_config(global_world_size) + num_nvl_bytes = max( + dispatch_config.get_nvl_buffer_size_hint(hidden_bytes, global_world_size), + combine_config.get_nvl_buffer_size_hint(hidden_bytes, global_world_size), + ) + num_rdma_bytes = max( + dispatch_config.get_rdma_buffer_size_hint(hidden_bytes, global_world_size), + combine_config.get_rdma_buffer_size_hint(hidden_bytes, global_world_size), + ) + self.ep_buffer = deep_ep.Buffer( + deepep_group, + num_nvl_bytes=num_nvl_bytes, + num_rdma_bytes=num_rdma_bytes, + low_latency_mode=False, + ) + else: + self.ep_buffer = deep_ep.ElasticBuffer( + deepep_group, + num_max_tokens_per_rank=self.ll_num_tokens, + hidden=self.ll_hidden, + num_topk=num_experts_per_tok, + use_fp8_dispatch=True, + allow_multiple_reduction=True, + ) self.ep_mega_moe_buffer = None self.ep_low_latency_buffer = None @@ -209,7 +239,13 @@ def new_deepep_group( raise ValueError("No valid MoE quant method was found while initializing DeepEP buffers") mega_moe_quant_method = "fp4fp8-b32-deepgemm" + sm90_mega_moe_quant_method = "fp8w8a8-b128-deepgemm" is_sm100 = is_sm100_gpu() + enable_sm90_mega_moe = is_sm90_gpu() and os.getenv("LIGHTLLM_ENABLE_SM90_MEGA_MOE", "0").upper() in { + "1", + "ON", + "TRUE", + } # Buffer 选择规则： # 1. 非 SM100 不支持 Mega MoE，只初始化 legacy low-latency buffer； @@ -225,6 +261,13 @@ def new_deepep_group( ) enable_mega_moe_buffer = has_mega_moe_layer enable_low_latency_buffer = has_legacy_moe_layer + elif enable_sm90_mega_moe: + has_mega_moe_layer = sm90_mega_moe_quant_method in expert_quant_method_names + has_legacy_moe_layer = any( + method_name != sm90_mega_moe_quant_method for method_name in expert_quant_method_names + ) + enable_mega_moe_buffer = has_mega_moe_layer + enable_low_latency_buffer = has_legacy_moe_layer else: enable_mega_moe_buffer = False enable_low_latency_buffer = True @@ -265,6 +308,9 @@ def new_deepep_group( import deep_gemm + mega_moe_kwargs = {} + if enable_sm90_mega_moe: + mega_moe_kwargs.update(use_fp8_dispatch=True, activation="swiglu") self.ep_mega_moe_buffer = deep_gemm.get_symm_buffer_for_mega_moe( deepep_group, self.ll_num_experts, @@ -272,14 +318,25 @@ def new_deepep_group( num_experts_per_tok, self.ll_hidden, moe_intermediate_size, + **mega_moe_kwargs, ) logger.info( - "Initialize DeepEP MoE buffers: low_latency=%s, mega_moe=%s, expert_quant_method_names=%s", + "Initialize DeepEP MoE buffers: legacy_prefill=%s, low_latency=%s, mega_moe=%s, " + "expert_quant_method_names=%s", + self.ep_prefill_uses_legacy_buffer, enable_low_latency_buffer, enable_mega_moe_buffer, sorted(expert_quant_method_names), ) - theoretical_sms = self.ep_buffer.get_theoretical_num_sms(self.ll_num_experts, num_experts_per_tok) + theoretical_sms = ( + 0 + if enable_mega_moe_buffer and not enable_low_latency_buffer + else ( + deep_ep.Buffer.num_sms + if self.ep_prefill_uses_legacy_buffer + else self.ep_buffer.get_theoretical_num_sms(self.ll_num_experts, num_experts_per_tok) + ) + ) self._set_num_sms_for_deep_gemm(theoretical_sms) def _set_num_sms_for_deep_gemm(self, deepep_sms: int): diff --git a/lightllm/distributed/symm_mem_all_reduce.py b/lightllm/distributed/symm_mem_all_reduce.py index 2256a5093f..e23c0a478e 100644 --- a/lightllm/distributed/symm_mem_all_reduce.py +++ b/lightllm/distributed/symm_mem_all_reduce.py @@ -55,7 +55,16 @@ def __init__(self, group: ProcessGroup, device, dtype: torch.dtype = torch.bfloa self.use_multimem = self.world_size in _WORLD_SIZES_MULTIMEM.get(cap_str, []) try: - self.buffer = torch_symm_mem.empty(self.max_size // dtype.itemsize, device=device, dtype=dtype) + # The optional out-of-place path adopts this storage through + # ``Tensor.data``. PyTorch cannot safely mix a normal tensor's + # version-counter state into an inference tensor, so make the + # reusable communication workspace explicitly inference-only. + with torch.inference_mode(): + self.buffer = torch_symm_mem.empty( + self.max_size // dtype.itemsize, + device=device, + dtype=dtype, + ) handle = torch_symm_mem.rendezvous(self.buffer, group.group_name) except RuntimeError as e: logger.warning("SymmMemAllreduce: rendezvous failed (%s). Disabling.", e) @@ -82,11 +91,29 @@ def should_use(self, inp: torch.Tensor) -> bool: # CustomProcessGroup.all_reduce: FlashInfer claims small messages first. return nbytes < self.max_size + def _reduce_to_workspace(self, inp: torch.Tensor) -> torch.Tensor: + # Mutating an inference tensor is only legal inside inference mode. + # Enter it explicitly as capability probes also call this helper from + # ordinary no-grad code. + with torch.inference_mode(): + n = inp.numel() + output = self.buffer[:n] + output.copy_(inp.view(-1)) + if self.use_multimem: + torch.ops.symm_mem.multimem_all_reduce_(output, "sum", self.group.group_name) + else: + torch.ops.symm_mem.two_shot_all_reduce_(output, "sum", self.group.group_name) + return output.view_as(inp) + def all_reduce(self, inp: torch.Tensor) -> None: - n = inp.numel() - self.buffer[:n].copy_(inp.view(-1)) - if self.use_multimem: - torch.ops.symm_mem.multimem_all_reduce_(self.buffer[:n], "sum", self.group.group_name) - else: - torch.ops.symm_mem.two_shot_all_reduce_(self.buffer[:n], "sum", self.group.group_name) - inp.view(-1).copy_(self.buffer[:n]) + inp.copy_(self._reduce_to_workspace(inp)) + + def all_reduce_out_of_place(self, inp: torch.Tensor) -> torch.Tensor: + """Reduce into symmetric memory and return its shaped workspace view. + + The result aliases the single reusable communication workspace. This + is suitable for the model's serialized current-stream all-reduces: + each result is consumed before the next reduction overwrites it. + """ + + return self._reduce_to_workspace(inp) diff --git a/lightllm/models/__init__.py b/lightllm/models/__init__.py index c7e9a59aad..b464f65340 100644 --- a/lightllm/models/__init__.py +++ b/lightllm/models/__init__.py @@ -21,6 +21,8 @@ from lightllm.models.deepseek2.model import Deepseek2TpPartModel from lightllm.models.deepseek3_2.model import Deepseek3_2TpPartModel from lightllm.models.glm4_moe_lite.model import Glm4MoeLiteTpPartModel +from lightllm.models.glm5_next.model import Glm5NextTpPartModel +from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel from lightllm.models.internvl.model import ( InternVLLlamaTpPartModel, InternVLPhi3TpPartModel, diff --git a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py index 3254031056..a1343e4165 100644 --- a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py @@ -8,7 +8,7 @@ from lightllm.models.deepseek2.triton_kernel.rotary_emb import rotary_emb_fwd from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import ( - use_sm100_mega_moe, + use_mega_moe, ) from functools import partial from lightllm.models.llama.yarn_rotary_utils import get_deepseek_mscale @@ -213,6 +213,10 @@ def _get_o( o_tensor = self._tpsp_reduce(input=o_tensor, infer_state=infer_state) return o_tensor + def _shared_ffn_tp(self, input, infer_state, layer_weight): + """Shared-expert FFN hook for model-specific activation semantics.""" + return LlamaTransformerLayerInfer._ffn_tp(self, input, infer_state, layer_weight) + def _moe_ffn_tp( self, input, infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight ) -> torch.Tensor: @@ -222,7 +226,7 @@ def _moe_ffn_tp( # if fused_shared_experts is not enabled, compute shared_output if self.n_shared_experts is not None and layer_weight.num_fused_shared_experts == 0: - shared_output = LlamaTransformerLayerInfer._ffn_tp(self, hidden_states, infer_state, layer_weight) + shared_output = self._shared_ffn_tp(hidden_states, infer_state, layer_weight) moe_gate_dtype = layer_weight.moe_gate.data_type_ router_logits = layer_weight.moe_gate.mm(hidden_states.to(moe_gate_dtype)) @@ -234,6 +238,7 @@ def _moe_ffn_tp( use_grouped_topk=self.n_group, topk_group=self.topk_group, num_expert_group=self.n_group, + is_prefill=infer_state.is_prefill, infer_state=infer_state, ) @@ -249,7 +254,7 @@ def _moe_ffn_edp( hidden_states = input token_num, hidden_dim = hidden_states.shape if self.n_shared_experts is not None: - shared_output = LlamaTransformerLayerInfer._ffn_tp(self, hidden_states, infer_state, layer_weight) + shared_output = self._shared_ffn_tp(hidden_states, infer_state, layer_weight) moe_gate_dtype = layer_weight.moe_gate.data_type_ router_logits = layer_weight.moe_gate.mm(hidden_states.to(moe_gate_dtype)) @@ -300,7 +305,7 @@ def overlap_tpsp_token_forward( infer_state1: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, ): - if not self.is_moe or use_sm100_mega_moe(layer_weight.experts.quant_method): + if not self.is_moe or use_mega_moe(layer_weight.experts.quant_method): return super().overlap_tpsp_token_forward( input_embdings, input_embdings1, infer_state, infer_state1, layer_weight ) @@ -324,7 +329,7 @@ def overlap_tpsp_token_forward( # 0 shared expert if self.n_shared_experts is not None: - _0_shared_output = LlamaTransformerLayerInfer._ffn_tp(self, _0_input1, infer_state, layer_weight) + _0_shared_output = self._shared_ffn_tp(_0_input1, infer_state, layer_weight) # 0 dispatch ( @@ -359,7 +364,7 @@ def overlap_tpsp_token_forward( # 1 shared expert if self.n_shared_experts is not None: - _1_shared_output = LlamaTransformerLayerInfer._ffn_tp(self, _1_input1, infer_state1, layer_weight) + _1_shared_output = self._shared_ffn_tp(_1_input1, infer_state1, layer_weight) # 1 dispatch ( @@ -426,7 +431,7 @@ def overlap_tpsp_context_forward( infer_state1: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, ): - if not self.is_moe or use_sm100_mega_moe(layer_weight.experts.quant_method): + if not self.is_moe or use_mega_moe(layer_weight.experts.quant_method): return super().overlap_tpsp_context_forward( input_embdings, input_embdings1, infer_state, infer_state1, layer_weight ) @@ -495,17 +500,17 @@ def overlap_tpsp_context_forward( # 0 shared expert if self.n_shared_experts is not None: - _0_shared_output = LlamaTransformerLayerInfer._ffn_tp(self, _0_input1, infer_state, layer_weight) + _0_shared_output = self._shared_ffn_tp(_0_input1, infer_state, layer_weight) # 1 shared expert if self.n_shared_experts is not None: - _1_shared_output = LlamaTransformerLayerInfer._ffn_tp(self, _1_input1, infer_state1, layer_weight) + _1_shared_output = self._shared_ffn_tp(_1_input1, infer_state1, layer_weight) # 0 moe calu _0_moe_out = layer_weight.experts.prefilled_group_gemm( _0_num_recv_tokens_per_expert_list, - _0_handle.num_unaligned_recv_tokens_per_expert, - _0_handle.recv_src_metadata, + getattr(_0_handle, "num_unaligned_recv_tokens_per_expert", None), + getattr(_0_handle, "recv_src_metadata", None), _0_recv_x, _0_recv_topk_idx, _0_recv_topk_weight, @@ -536,8 +541,8 @@ def overlap_tpsp_context_forward( # 1 moe calc _1_moe_out = layer_weight.experts.prefilled_group_gemm( _1_num_recv_tokens_per_expert_list, - _1_handle.num_unaligned_recv_tokens_per_expert, - _1_handle.recv_src_metadata, + getattr(_1_handle, "num_unaligned_recv_tokens_per_expert", None), + getattr(_1_handle, "recv_src_metadata", None), _1_recv_x, _1_recv_topk_idx, _1_recv_topk_weight, diff --git a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py index d6eaebe2fd..a31eb3e407 100644 --- a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py @@ -1,14 +1,25 @@ +import os + import torch from typing import Any from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo -from lightllm.models.deepseek2.layer_infer.transformer_layer_infer import Deepseek2TransformerLayerInfer -from lightllm.models.deepseek3_2.layer_weights.transformer_layer_weight import Deepseek3_2TransformerLayerWeight +from lightllm.models.deepseek2.layer_infer.transformer_layer_infer import ( + Deepseek2TransformerLayerInfer, +) +from lightllm.models.deepseek3_2.layer_weights.transformer_layer_weight import ( + Deepseek3_2TransformerLayerWeight, +) from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward from lightllm.models.deepseek2.triton_kernel.rotary_emb import rotary_emb_fwd from lightllm.common.basemodel.attention.base_att import AttControl from lightllm.models.deepseek3_2.triton_kernel.act_quant import act_quant -from lightllm.models.deepseek3_2.triton_kernel.destindex_copy_indexer_ks import destindex_copy_indexer_ks -from lightllm.models.deepseek3_2.triton_kernel.extract_indexer_ks import extract_indexer_ks +from lightllm.models.deepseek3_2.triton_kernel.destindex_copy_indexer_ks import ( + destindex_copy_indexer_ks, +) +from lightllm.models.deepseek3_2.triton_kernel.extract_indexer_ks import ( + extract_indexer_ks, + extract_indexer_ks_dynamic, +) from lightllm.utils.envs_utils import get_env_start_args from lightllm.distributed import all_gather_into_tensor @@ -19,7 +30,9 @@ def __init__(self, layer_num, network_config): super().__init__(layer_num, network_config) self.indexer = NsaInfer( - layer_idx=self.layer_num_, network_config=self.network_config_, tp_world_size=self.tp_world_size_ + layer_idx=self.layer_num_, + network_config=self.network_config_, + tp_world_size=self.tp_world_size_, ) return @@ -115,7 +128,10 @@ def _token_attention_kernel( out=None, ): # Model-specific q projection (uses layer weights) - q_nope, q_rope = q[:, :, : -self.qk_rope_head_dim], q[:, :, -self.qk_rope_head_dim :] + q_nope, q_rope = ( + q[:, :, : -self.qk_rope_head_dim], + q[:, :, -self.qk_rope_head_dim :], + ) q_nope = layer_weight.k_b_proj_.bmm(q_nope.transpose(0, 1)).transpose(0, 1) # 计算 topk mem indices @@ -151,6 +167,12 @@ def _token_attention_kernel( class NsaInfer: + _MQA_LOGITS_BYTES_PER_ELEM = 4 + _MQA_LOGITS_STATIC_SKIP_ELEMS = 8_000_000 + _MQA_LOGITS_FREE_MEM_FRACTION = 0.2 + _FLASHMLA_SPARSE_TOPK_ALIGNMENT = 128 + _mqa_logits_budget_bytes = {} + def __init__(self, layer_idx: int, network_config: dict, tp_world_size: int): super().__init__() self.layer_idx_ = layer_idx @@ -167,6 +189,20 @@ def __init__(self, layer_idx: int, network_config: dict, tp_world_size: int): self.index_n_heads_scale = (self.index_n_heads ** -0.5) * self.softmax_scale self.tp_world_size_ = tp_world_size self.tp_index_n_heads = self.index_n_heads // self.tp_world_size_ + self.index_kpool = network_config.get("index_kpool", 1) + self.index_kpool_compress = network_config.get("index_kpool_compress", False) + self.enable_kpool_decode_fastpath = os.getenv("LIGHTLLM_ENABLE_KPOOL_DECODE_FASTPATH", "0").upper() in { + "1", + "ON", + "TRUE", + } + self._kpool_tail_k = None + self._kpool_tail_score = None + # Most NSA models only instantiate the target model, so their decode + # layout follows the process-wide MTP setting. A model that reuses an + # NSA layer as a recurrent drafter can override this with its own + # physical decode width (normally zero extra rows). + self.decode_mtp_step = None def _get_indices( self, @@ -176,16 +212,25 @@ def _get_indices( att_state: Any, layer_weight: Deepseek3_2TransformerLayerWeight, ): - - q, k = self._get_q_k_bf16(hidden_states, q_lora, infer_state, layer_weight) + q_k = self._get_q_k_bf16(hidden_states, q_lora, infer_state, layer_weight) + if len(q_k) == 2: + q, k = q_k + raw_k = None + else: + q, k, raw_k = q_k if self.tp_world_size_ > 1: q_merge = torch.empty( - size=(self.tp_world_size_ * q.numel()), + size=(self.tp_world_size_ * q.numel(),), dtype=q.dtype, device=q.device, ) - all_gather_into_tensor(output_=q_merge, input_=q.view(-1), group=infer_state.dist_group, async_op=False) + all_gather_into_tensor( + output_=q_merge, + input_=q.view(-1), + group=infer_state.dist_group, + async_op=False, + ) q = ( q_merge.view(self.tp_world_size_, q.shape[0], self.tp_index_n_heads, q.shape[2]) .transpose(0, 1) @@ -193,72 +238,393 @@ def _get_indices( .view(q.shape[0], self.index_n_heads, q.shape[2]) ) - q_fp8, q_scale = act_quant(q, self.block_size, self.scale_fmt) - k_fp8, k_scale = act_quant(k, self.block_size, self.scale_fmt) + q_fp8, q_scale = self._quantize_indexer_activation(q) + k_fp8, k_scale = self._quantize_indexer_activation(k) + indexer_k_buffer = infer_state.mem_manager.get_indexer_k_buffer(self.layer_idx_) destindex_copy_indexer_ks( K_fp8=k_fp8, K_scale=k_scale, DestLoc=infer_state.mem_index, - O_buffer=infer_state.mem_manager.get_indexer_k_buffer(self.layer_idx_), + O_buffer=indexer_k_buffer, ) - weights = layer_weight.weights_proj_.mm(hidden_states) * self.index_n_heads_scale - weights = weights.unsqueeze(-1) * q_scale + weights = self._scale_indexer_weights(layer_weight.weights_proj_.mm(hidden_states), q_scale) ks = att_state.ks ke = att_state.ke lengths = att_state.lengths - if infer_state.is_prefill: + use_kpool_prefill = ( + infer_state.is_prefill + and self.index_kpool > 1 + and self.index_kpool_compress + and raw_k is not None + and infer_state.kpool_prefill_aligned + and infer_state.mem_index.shape[0] == q_fp8.shape[0] + ) + use_kpool_decode = ( + not infer_state.is_prefill + and self.enable_kpool_decode_fastpath + and infer_state.kpool_decode_aligned + and raw_k is not None + and get_env_start_args().mtp_mode is None + ) + use_kpool = use_kpool_prefill or use_kpool_decode + if use_kpool_prefill: + (k_fp8_, k_scale_, score_ks, score_ke, score_lengths,) = self._prepare_kpool_scoring( + raw_k=raw_k, + hidden_states=hidden_states, + q_lora=q_lora, + infer_state=infer_state, + layer_weight=layer_weight, + indexer_k_buffer=indexer_k_buffer, + ragged_mem_index=att_state.ragged_mem_index, + ks=ks, + lengths=lengths, + ) + use_kpool = k_fp8_ is not None + elif use_kpool_decode: + (k_fp8_, k_scale_, score_ks, score_ke, score_lengths,) = self._prepare_kpool_decode_scoring( + raw_k=raw_k, + hidden_states=hidden_states, + q_lora=q_lora, + infer_state=infer_state, + layer_weight=layer_weight, + indexer_k_buffer=indexer_k_buffer, + lengths=lengths, + ) + use_kpool = k_fp8_ is not None + + if use_kpool: + mtp_step = 0 + elif infer_state.is_prefill: mtp_step = 0 else: - mtp_step = get_env_start_args().mtp_step - # Use efficient Triton kernel to extract FP8 keys and scales from buffer - k_fp8_, k_scale_ = extract_indexer_ks( - I_buffer=infer_state.mem_manager.get_indexer_k_buffer(self.layer_idx_), - b_seq_len=infer_state.b_seq_len, - b_req_idx=infer_state.b_req_idx, - req_to_token_indexs=infer_state.req_manager.req_to_token_indexs, - out_token_num=infer_state.b_seq_len.shape[0] * infer_state.max_kv_seq_len, - max_kv_seq_len=infer_state.max_kv_seq_len, - mtp_step=mtp_step, - ) + mtp_step = get_env_start_args().mtp_step if self.decode_mtp_step is None else self.decode_mtp_step + # LightSpec compacts each request to a variable number of contiguous + # verify rows. Its sparse-index K packing must follow request boundaries + # instead of assuming the fixed process-wide MTP width. + use_dynamic_layout = not infer_state.is_prefill and mtp_step > 0 and get_env_start_args().mtp_dynamic_verify + if use_kpool: + pass + elif use_dynamic_layout: + k_fp8_, k_scale_ = extract_indexer_ks_dynamic( + I_buffer=indexer_k_buffer, + b_seq_len=infer_state.b_seq_len, + b_req_idx=infer_state.b_req_idx, + b_mtp_index=infer_state.b_mtp_index, + req_to_token_indexs=infer_state.req_manager.req_to_token_indexs, + max_kv_seq_len=infer_state.max_kv_seq_len, + max_request_num=infer_state.req_manager.max_request_num, + ) + else: + k_fp8_, k_scale_ = extract_indexer_ks( + I_buffer=indexer_k_buffer, + b_seq_len=infer_state.b_seq_len, + b_req_idx=infer_state.b_req_idx, + req_to_token_indexs=infer_state.req_manager.req_to_token_indexs, + out_token_num=infer_state.b_seq_len.shape[0] * infer_state.max_kv_seq_len, + max_kv_seq_len=infer_state.max_kv_seq_len, + mtp_step=mtp_step, + ) + if not use_kpool: + score_ks, score_ke, score_lengths = ks, ke, lengths import deep_gemm - logits = deep_gemm.fp8_mqa_logits( - q_fp8, - (k_fp8_, k_scale_), - weights.squeeze(-1), - ks, - ke, - clean_logits=False, - max_seqlen_k=infer_state.max_kv_seq_len, - ) - from sgl_kernel import fast_topk_v2 - b_topk_index = fast_topk_v2( - score=logits, - lengths=lengths, - topk=self.index_topk, + query_token_num = q_fp8.shape[0] + kv_token_num = k_fp8_.shape[0] + query_chunk_size = self._get_mqa_logits_chunk_size( + query_token_num=query_token_num, + kv_token_num=kv_token_num, + device=q_fp8.device, ) - b_topk_index = torch.where(b_topk_index != -1, b_topk_index + ks.view(-1, 1), -1) + + if use_kpool: + output_topk = ( + (self.index_topk + self.index_kpool - 1 + self._FLASHMLA_SPARSE_TOPK_ALIGNMENT - 1) + // self._FLASHMLA_SPARSE_TOPK_ALIGNMENT + * self._FLASHMLA_SPARSE_TOPK_ALIGNMENT + ) + # K-pool appends at most pool_size - 1 always-selected tail + # tokens. Allocate the FlashMLA-aligned result once and leave the + # alignment gap masked, avoiding two full-size top-k buffers. + b_topk_index = torch.full( + (query_token_num, output_topk), + -1, + dtype=torch.int32, + device=q_fp8.device, + ) + else: + b_topk_index = torch.empty( + (query_token_num, self.index_topk), + dtype=torch.int32, + device=q_fp8.device, + ) + for start in range(0, query_token_num, query_chunk_size): + end = min(start + query_chunk_size, query_token_num) + logits = deep_gemm.fp8_mqa_logits( + q_fp8[start:end], + (k_fp8_, k_scale_), + weights[start:end], + score_ks[start:end], + score_ke[start:end], + # fast top-k already masks columns past each row's valid + # length, while DeepGEMM's small-M path rejects clean_logits. + clean_logits=False, + max_seqlen_k=kv_token_num, + ) + if use_kpool: + from sglang.srt.layers.attention.dsa.kpool_fp8_index import ( + topk_from_pooled_history_logits, + ) + + topk_index = topk_from_pooled_history_logits( + logits=logits, + group_lengths=score_lengths[start:end], + pool_size=self.index_kpool, + topk=self.index_topk, + seq_lens=lengths[start:end], + ) + else: + topk_index = fast_topk_v2( + score=logits, + lengths=score_lengths[start:end], + topk=self.index_topk, + ) + b_topk_index[start:end, : topk_index.shape[1]].copy_(topk_index) + # The long-prefill score matrix can be tens of GiB. Release each + # chunk before computing the next one. + del logits, topk_index # 将 topk index 转化为 mem index - from ..triton_kernel.topk_index_to_mem_index import trans_topk_index_to_mem_index + from ..triton_kernel.topk_index_to_mem_index import ( + trans_topk_index_to_mem_index, + ) b_topk_mem_index = trans_topk_index_to_mem_index( topk_index=b_topk_index, + ragged_start_index=ks, ragged_mem_index=att_state.ragged_mem_index, ) return b_topk_mem_index, b_topk_index + def _quantize_indexer_activation(self, value: torch.Tensor): + return act_quant(value, self.block_size, self.scale_fmt) + + def _scale_indexer_weights(self, weights: torch.Tensor, q_scale: torch.Tensor) -> torch.Tensor: + return (weights.mul(self.index_n_heads_scale).unsqueeze(-1).mul(q_scale)).squeeze(-1) + + def _prepare_kpool_scoring( + self, + raw_k, + hidden_states, + q_lora, + infer_state, + layer_weight, + indexer_k_buffer, + ragged_mem_index, + ks, + lengths, + ): + pool_size = self.index_kpool + query_token_num = raw_k.shape[0] + if not infer_state.kpool_prefill_aligned: + return None, None, None, None, None + # Without decode K-pool enabled the zero-prefix path is intentionally + # transient, so no pooled history exists for a later chunk. + if infer_state.max_cache_len > 0 and not self.enable_kpool_decode_fastpath: + return None, None, None, None, None + + gate_score = layer_weight.index_kpool_compress_gate.mm(hidden_states.to(q_lora.dtype)) + closed_pool_num = query_token_num // pool_size + compressed_k = None + compressed_scale = None + if closed_pool_num: + closed_token_num = closed_pool_num * pool_size + slot_k = raw_k[:closed_token_num].view(closed_pool_num, pool_size, self.index_head_dim) + slot_score = gate_score[:closed_token_num].view(closed_pool_num, pool_size, self.index_head_dim) + write_locs = infer_state.mem_index[:closed_token_num].view(closed_pool_num, pool_size)[:, -1] + compressed_k, compressed_scale = self._compress_kpool_keys( + slot_k=slot_k, + slot_score=slot_score, + write_locs=write_locs, + layer_weight=layer_weight, + output_buffer=indexer_k_buffer, + persist=self.enable_kpool_decode_fastpath, + ) + + # The common serving path has no prefix and each request fits in this + # aligned chunk. DeepGEMM can consume the freshly compressed pools + # directly; allocating and gathering a second max-token-sized cache + # would waste several GiB for GLM-5.3 TP8 and can make c64 OOM. + if infer_state.max_cache_len == 0: + if compressed_k is None: + return None, None, None, None, None + pool_lengths = torch.div(lengths, pool_size, rounding_mode="floor").to(torch.int32) + score_ks = torch.div(ks, pool_size, rounding_mode="floor").to(torch.int32) + score_ke = score_ks + pool_lengths + return compressed_k, compressed_scale, score_ks, score_ke, pool_lengths + + pooled_token_num = infer_state.total_token_num // pool_size + if pooled_token_num == 0: + return None, None, None, None, None + # Every request boundary is pool-aligned, so selecting each pool's + # final token from the ragged request-major layout preserves exactly + # the packed-K order expected by DeepGEMM's per-row ks/ke ranges. + pooled_ragged_positions = torch.arange( + pool_size - 1, + infer_state.total_token_num, + pool_size, + dtype=torch.int64, + device=raw_k.device, + ) + mem_indices = ragged_mem_index[pooled_ragged_positions].to(torch.int64) + packed_k = indexer_k_buffer[mem_indices, 0] + k_fp8 = packed_k[:, : self.index_head_dim].contiguous().view(torch.float8_e4m3fn) + k_scale = packed_k[:, self.index_head_dim : self.index_head_dim + 4].contiguous().view(torch.float32).view(-1) + + pool_lengths = torch.div(lengths, pool_size, rounding_mode="floor").to(torch.int32) + score_ks = torch.div(ks, pool_size, rounding_mode="floor").to(torch.int32) + score_ke = score_ks + pool_lengths + return k_fp8, k_scale, score_ks, score_ke, pool_lengths + + def _prepare_kpool_decode_scoring( + self, + raw_k, + hidden_states, + q_lora, + infer_state, + layer_weight, + indexer_k_buffer, + lengths, + ): + """Update one-token K-pool tails and gather pooled decode history. + + This fast path deliberately targets plain decode. MTP verify has + multiple speculative rows per request and needs acceptance-aware tail + rollback, so it remains on the regular indexer path. + """ + + pool_size = self.index_kpool + batch_size = raw_k.shape[0] + if batch_size == 0: + return None, None, None, None, None + + if self._kpool_tail_k is None: + tail_shape = ( + infer_state.req_manager.max_request_num + 1, + pool_size, + self.index_head_dim, + ) + self._kpool_tail_k = torch.empty(tail_shape, dtype=torch.bfloat16, device=raw_k.device) + self._kpool_tail_score = torch.empty_like(self._kpool_tail_k) + + req_idx = infer_state.b_req_idx.to(torch.int64) + positions = infer_state.b_seq_len.to(torch.int64) - 1 + tail_slots = torch.remainder(positions, pool_size) + gate_score = layer_weight.index_kpool_compress_gate.mm(hidden_states.to(q_lora.dtype)) + self._kpool_tail_k[req_idx, tail_slots] = raw_k + self._kpool_tail_score[req_idx, tail_slots] = gate_score + + # Compress every row's current tail. Only pool-end token locations are + # gathered as history, so writes at intermediate token locations are + # harmless and avoid a dynamic nonzero/host synchronization. + self._compress_kpool_keys( + slot_k=self._kpool_tail_k[req_idx], + slot_score=self._kpool_tail_score[req_idx], + write_locs=infer_state.mem_index, + layer_weight=layer_weight, + output_buffer=indexer_k_buffer, + persist=True, + ) + + max_pool_len = infer_state.max_kv_seq_len // pool_size + if max_pool_len == 0: + return None, None, None, None, None + pool_lengths = torch.div(lengths, pool_size, rounding_mode="floor").to(torch.int32) + endpoint_positions = ( + torch.arange(max_pool_len, dtype=torch.int64, device=raw_k.device) * pool_size + pool_size - 1 + ) + last_endpoint = torch.clamp(pool_lengths.to(torch.int64) * pool_size - 1, min=0) + safe_positions = torch.minimum(endpoint_positions.unsqueeze(0), last_endpoint.unsqueeze(1)) + mem_indices = infer_state.req_manager.req_to_token_indexs[req_idx.unsqueeze(1), safe_positions].to(torch.int64) + packed_k = indexer_k_buffer[mem_indices.reshape(-1), 0] + k_fp8 = packed_k[:, : self.index_head_dim].contiguous().view(torch.float8_e4m3fn) + k_scale = packed_k[:, self.index_head_dim : self.index_head_dim + 4].contiguous().view(torch.float32).view(-1) + score_ks = torch.arange(batch_size, dtype=torch.int32, device=raw_k.device) * max_pool_len + score_ke = score_ks + pool_lengths + return k_fp8, k_scale, score_ks, score_ke, pool_lengths + + def _compress_kpool_keys( + self, + slot_k, + slot_score, + write_locs, + layer_weight, + output_buffer, + persist, + ): + from types import SimpleNamespace + + from sglang.srt.layers.attention.dsa.kpool_fp8_index import ( + kpool_softmax_rotate_write_cache, + ) + + compressed_k, compressed_scale = kpool_softmax_rotate_write_cache( + pool=SimpleNamespace(page_size=64, index_head_dim=self.index_head_dim), + buf=output_buffer, + slot_k=slot_k, + slot_score=slot_score, + ape=layer_weight.index_kpool_compress_ape.weight, + loc=write_locs.to(torch.int64), + round_scale=self.scale_fmt is not None, + return_compressed=True, + write_cache=False, + ) + if persist: + destindex_copy_indexer_ks( + K_fp8=compressed_k, + K_scale=compressed_scale, + DestLoc=write_locs, + O_buffer=output_buffer, + ) + return compressed_k, compressed_scale + + @classmethod + def _get_mqa_logits_chunk_size(cls, query_token_num: int, kv_token_num: int, device: torch.device) -> int: + score_element_num = query_token_num * kv_token_num + if score_element_num < cls._MQA_LOGITS_STATIC_SKIP_ELEMS: + return query_token_num + + device_index = device.index + if device_index is None: + device_index = torch.cuda.current_device() + logits_budget_bytes = cls._mqa_logits_budget_bytes.get(device_index) + if logits_budget_bytes is None: + free_memory_bytes, _ = torch.cuda.mem_get_info(device_index) + free_mem_fraction = float( + os.getenv( + "LIGHTLLM_MQA_LOGITS_FREE_MEM_FRACTION", + cls._MQA_LOGITS_FREE_MEM_FRACTION, + ) + ) + free_mem_fraction = min(max(free_mem_fraction, 0.01), 0.9) + logits_budget_bytes = max(1, int(free_memory_bytes * free_mem_fraction)) + cls._mqa_logits_budget_bytes[device_index] = logits_budget_bytes + + bytes_per_query = kv_token_num * cls._MQA_LOGITS_BYTES_PER_ELEM + return max(1, min(query_token_num, logits_budget_bytes // bytes_per_query)) + @staticmethod def _rotate_activation(x: torch.Tensor) -> torch.Tensor: assert x.dtype == torch.bfloat16 - from lightllm.models.deepseek3_2.triton_kernel.hadamard_transform import hadamard_transform + from lightllm.models.deepseek3_2.triton_kernel.hadamard_transform import ( + hadamard_transform, + ) hidden_size = x.size(-1) assert (hidden_size & (hidden_size - 1)) == 0, "Hidden size must be a power of 2 for Hadamard transform." diff --git a/lightllm/models/deepseek3_2/triton_kernel/extract_indexer_ks.py b/lightllm/models/deepseek3_2/triton_kernel/extract_indexer_ks.py index f02fc30942..06caa7dae6 100644 --- a/lightllm/models/deepseek3_2/triton_kernel/extract_indexer_ks.py +++ b/lightllm/models/deepseek3_2/triton_kernel/extract_indexer_ks.py @@ -43,7 +43,12 @@ def _fwd_kernel_extract_indexer_ks( store_start_index = tl.sum(b_seq_len) for i in range(token_start_index, cur_seq_len, tl.num_programs(1)): - mem_index = tl.load(req_to_token_indexs + cur_req_idx * stride_req_to_token_m + i * stride_req_to_token_n) + # Token slots can exceed INT32_MAX / stride_in_fp8_bs when the KV + # cache packs multiple states into one row. Promote before pointer + # arithmetic so large cache offsets do not wrap around. + mem_index = tl.load(req_to_token_indexs + cur_req_idx * stride_req_to_token_m + i * stride_req_to_token_n).to( + tl.int64 + ) in_fp8_ptrs = in_fp8 + mem_index * stride_in_fp8_bs + 0 * stride_in_fp8_h + stride_in_fp8_d * offs_d kv_fp8 = tl.load(in_fp8_ptrs) @@ -60,6 +65,86 @@ def _fwd_kernel_extract_indexer_ks( return +@triton.jit +def _fwd_kernel_extract_indexer_ks_dynamic( + in_fp8, + stride_in_fp8_bs, + stride_in_fp8_h, + stride_in_fp8_d, + in_fp8_scale, + stride_in_scale_bs, + stride_in_scale_h, + stride_in_scale_d, + req_to_token_indexs, + stride_req_to_token_m, + stride_req_to_token_n, + b_seq_len, + b_req_idx, + b_mtp_index, + batch_size, + O_fp8, + stride_o_fp8_bs, + stride_o_fp8_d, + O_scale, + stride_o_scale_bs, + stride_o_scale_d, + BLOCK_DMODEL: tl.constexpr, + BLOCK_BATCH: tl.constexpr, +): + """Extract one packed K sequence per variable-width MTP request. + + Dynamic verification keeps each request's selected rows contiguous and + starts every request (and every CUDA-graph padding row) at MTP index zero. + A request's last row therefore has either no successor or a successor with + MTP index zero. Using that boundary instead of a process-wide fixed width + lets the sparse indexer consume LightSpec's compacted verify layout. + """ + + cur_row = tl.program_id(0) + token_start_index = tl.program_id(1) + next_mtp_index = tl.load( + b_mtp_index + cur_row + 1, + mask=cur_row + 1 < batch_size, + other=0, + ) + is_group_end = (cur_row == batch_size - 1) | (next_mtp_index == 0) + if not is_group_end: + return + + cur_req_idx = tl.load(b_req_idx + cur_row) + cur_seq_len = tl.load(b_seq_len + cur_row) + + rows = tl.arange(0, BLOCK_BATCH) + next_rows = rows + 1 + next_mtp_indices = tl.load( + b_mtp_index + next_rows, + mask=next_rows < batch_size, + other=0, + ) + prior_group_ends = (rows < cur_row) & (next_mtp_indices == 0) + prior_group_seq_lens = tl.load( + b_seq_len + rows, + mask=prior_group_ends, + other=0, + ) + store_start_index = tl.sum(prior_group_seq_lens) + + offs_d = tl.arange(0, BLOCK_DMODEL) + for i in range(token_start_index, cur_seq_len, tl.num_programs(1)): + mem_index = tl.load(req_to_token_indexs + cur_req_idx * stride_req_to_token_m + i * stride_req_to_token_n).to( + tl.int64 + ) + + in_fp8_ptrs = in_fp8 + mem_index * stride_in_fp8_bs + stride_in_fp8_d * offs_d + kv_fp8 = tl.load(in_fp8_ptrs) + in_scale_ptr = in_fp8_scale + mem_index * stride_in_scale_bs + kv_scale = tl.load(in_scale_ptr) + + o_fp8_ptrs = O_fp8 + (store_start_index + i) * stride_o_fp8_bs + stride_o_fp8_d * offs_d + tl.store(o_fp8_ptrs, kv_fp8) + tl.store(O_scale + (store_start_index + i) * stride_o_scale_bs, kv_scale) + + @torch.no_grad() def extract_indexer_ks( I_buffer: torch.Tensor, @@ -113,3 +198,65 @@ def extract_indexer_ks( ) return O_fp8, O_scale.squeeze(-1) + + +@torch.no_grad() +def extract_indexer_ks_dynamic( + I_buffer: torch.Tensor, + b_seq_len: torch.Tensor, + b_req_idx: torch.Tensor, + b_mtp_index: torch.Tensor, + req_to_token_indexs: torch.Tensor, + max_kv_seq_len: int, + max_request_num: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Extract packed indexer keys for a variable-width MTP verify batch.""" + + head_dim = 128 + batch_size = b_seq_len.shape[0] + assert I_buffer.dtype == torch.uint8, f"Expected I_buffer dtype=uint8, got {I_buffer.dtype}" + assert I_buffer.shape[2] == 132, f"Expected I_buffer last dim=132, got {I_buffer.shape[2]}" + assert b_req_idx.shape == b_seq_len.shape == b_mtp_index.shape + + in_fp8 = I_buffer[:, :, 0:128].view(dtype=torch.float8_e4m3fn) + in_fp8_scale = I_buffer[:, :, 128:132].view(dtype=torch.float32) + + # At most max_request_num rows belong to distinct real requests. Any + # remaining CUDA-graph padding rows have a fixed sequence length of two. + # This static upper bound avoids allocating batch_size * max_kv_seq_len for + # every captured graph while remaining safe for every compacted layout. + max_real_groups = min(batch_size, max_request_num) + output_capacity = max_real_groups * max_kv_seq_len + (batch_size - max_real_groups) * 2 + O_fp8 = torch.empty((output_capacity, head_dim), dtype=torch.float8_e4m3fn, device=I_buffer.device) + O_scale = torch.empty((output_capacity, 1), dtype=torch.float32, device=I_buffer.device) + + grid = (batch_size, min(256, max_kv_seq_len)) + _fwd_kernel_extract_indexer_ks_dynamic[grid]( + in_fp8, + stride_in_fp8_bs=in_fp8.stride(0), + stride_in_fp8_h=in_fp8.stride(1), + stride_in_fp8_d=in_fp8.stride(2), + in_fp8_scale=in_fp8_scale, + stride_in_scale_bs=in_fp8_scale.stride(0), + stride_in_scale_h=in_fp8_scale.stride(1), + stride_in_scale_d=in_fp8_scale.stride(2), + req_to_token_indexs=req_to_token_indexs, + stride_req_to_token_m=req_to_token_indexs.stride(0), + stride_req_to_token_n=req_to_token_indexs.stride(1), + b_seq_len=b_seq_len, + b_req_idx=b_req_idx, + b_mtp_index=b_mtp_index, + batch_size=batch_size, + O_fp8=O_fp8, + stride_o_fp8_bs=O_fp8.stride(0), + stride_o_fp8_d=O_fp8.stride(1), + O_scale=O_scale, + stride_o_scale_bs=O_scale.stride(0), + stride_o_scale_d=O_scale.stride(1), + BLOCK_DMODEL=head_dim, + BLOCK_BATCH=triton.next_power_of_2(batch_size), + num_warps=1, + num_stages=1, + ) + + return O_fp8, O_scale.squeeze(-1) diff --git a/lightllm/models/deepseek3_2/triton_kernel/hadamard_transform.py b/lightllm/models/deepseek3_2/triton_kernel/hadamard_transform.py index eabf703f56..5e6e1941f2 100644 --- a/lightllm/models/deepseek3_2/triton_kernel/hadamard_transform.py +++ b/lightllm/models/deepseek3_2/triton_kernel/hadamard_transform.py @@ -41,6 +41,41 @@ def _hadamard_transform_kernel( tl.store(Y + offsets, x * scale, mask=mask) +@triton.jit +def _hadamard_transform_quant_fp8_kernel( + X, + Y, + S, + n_rows, + scale: tl.constexpr, + BLOCK_R: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid = tl.program_id(0) + rows = pid * BLOCK_R + tl.arange(0, BLOCK_R) + row_mask = rows < n_rows + cols = tl.arange(0, BLOCK_N) + offsets = rows[:, None] * BLOCK_N + cols[None, :] + x = tl.load(X + offsets, mask=row_mask[:, None], other=0.0).to(tl.float32) + + x = _butterfly_stage(x, 64, 1, BLOCK_R, BLOCK_N) + x = _butterfly_stage(x, 32, 2, BLOCK_R, BLOCK_N) + x = _butterfly_stage(x, 16, 4, BLOCK_R, BLOCK_N) + x = _butterfly_stage(x, 8, 8, BLOCK_R, BLOCK_N) + x = _butterfly_stage(x, 4, 16, BLOCK_R, BLOCK_N) + x = _butterfly_stage(x, 2, 32, BLOCK_R, BLOCK_N) + x = _butterfly_stage(x, 1, 64, BLOCK_R, BLOCK_N) + + # Match the unfused path's bf16 Hadamard output before FP8 quantization. + x = (x * scale).to(tl.bfloat16).to(tl.float32) + absmax = tl.maximum(tl.max(tl.abs(x), axis=1), 1e-4) + quant_scale = tl.exp2(tl.ceil(tl.log2(absmax * (1.0 / 448.0)))) + y = tl.minimum(tl.maximum(x / quant_scale[:, None], -448.0), 448.0) + + tl.store(Y + offsets, y, mask=row_mask[:, None]) + tl.store(S + rows, quant_scale, mask=row_mask) + + @functools.lru_cache(maxsize=None) def _target_programs(device_index: int) -> int: return torch.cuda.get_device_properties(device_index).multi_processor_count * 2 @@ -78,3 +113,33 @@ def hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: assert x.size(-1) == 128, "DeepSeek-V3.2 Hadamard transform expects hidden size 128" return _hadamard_transform_triton(x, scale) + + +def hadamard_transform_quant_fp8(x: torch.Tensor, scale: float = 1.0) -> tuple[torch.Tensor, torch.Tensor]: + """Fuse Hadamard-128 with the following ue8m0 FP8 quantization.""" + + assert x.is_cuda, "hadamard_transform_quant_fp8 only supports CUDA tensors" + assert x.dtype == torch.bfloat16, "Hadamard transform expects bfloat16 input" + assert x.size(-1) == 128, "Hadamard transform expects hidden size 128" + if not x.is_contiguous(): + x = x.contiguous() + + original_shape = x.shape + rows = x.numel() // 128 + output = torch.empty_like(x, dtype=torch.float8_e4m3fn) + output_scale = torch.empty((*original_shape[:-1], 1), dtype=torch.float32, device=x.device) + if rows == 0: + return output, output_scale + + block_r = 32 + _hadamard_transform_quant_fp8_kernel[(triton.cdiv(rows, block_r),)]( + x, + output, + output_scale, + rows, + scale, + BLOCK_R=block_r, + BLOCK_N=128, + num_warps=2, + ) + return output.view(original_shape), output_scale diff --git a/lightllm/models/deepseek3_2/triton_kernel/indexer_weight_scale.py b/lightllm/models/deepseek3_2/triton_kernel/indexer_weight_scale.py new file mode 100644 index 0000000000..391d826f21 --- /dev/null +++ b/lightllm/models/deepseek3_2/triton_kernel/indexer_weight_scale.py @@ -0,0 +1,39 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _indexer_weight_scale_kernel( + weights, + q_scale, + size, + scale: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < size + value = tl.load(weights + offsets, mask=mask).to(tl.float32) + query_scale = tl.load(q_scale + offsets, mask=mask).to(tl.float32) + tl.store(weights + offsets, value * scale * query_scale, mask=mask) + + +def scale_indexer_weights_(weights: torch.Tensor, q_scale: torch.Tensor, scale: float) -> torch.Tensor: + """Apply the indexer query/head scales in one in-place kernel.""" + + assert weights.dtype == torch.float32 + assert weights.is_contiguous() and q_scale.is_contiguous() + assert weights.numel() == q_scale.numel() + size = weights.numel() + if size == 0: + return weights + block_size = 256 + _indexer_weight_scale_kernel[(triton.cdiv(size, block_size),)]( + weights, + q_scale, + size, + scale, + BLOCK_SIZE=block_size, + num_warps=4, + ) + return weights diff --git a/lightllm/models/deepseek3_2/triton_kernel/topk_index_to_mem_index.py b/lightllm/models/deepseek3_2/triton_kernel/topk_index_to_mem_index.py index 12786c6619..6c0996aaeb 100644 --- a/lightllm/models/deepseek3_2/triton_kernel/topk_index_to_mem_index.py +++ b/lightllm/models/deepseek3_2/triton_kernel/topk_index_to_mem_index.py @@ -9,25 +9,39 @@ def _trans_topk_index_to_mem_index( topk_index, topk_index_stride_b, topk_index_stride_k, + ragged_start_index, ragged_mem_index, topk_mem_index, topk_mem_index_stride_b, topk_mem_index_stride_k, + topk_width, BLOCK_DMODEL: tl.constexpr, ): cur_index = tl.program_id(0) offs_d = tl.arange(0, BLOCK_DMODEL) + mask = offs_d < topk_width topk_index_ptrs = topk_index + cur_index * topk_index_stride_b + offs_d * topk_index_stride_k - topk_indices = tl.load(topk_index_ptrs) + topk_indices = tl.load(topk_index_ptrs, mask=mask, other=-1) + ragged_start = tl.load(ragged_start_index + cur_index) + topk_indices = tl.where(topk_indices != -1, topk_indices + ragged_start, -1) + tl.store(topk_index_ptrs, topk_indices, mask=mask) dest_mem_index = ragged_mem_index + topk_indices - mem_index = tl.load(dest_mem_index, mask=topk_indices != -1, other=-1) - tl.store(topk_mem_index + cur_index * topk_mem_index_stride_b + offs_d * topk_mem_index_stride_k, mem_index) + mem_index = tl.load(dest_mem_index, mask=mask & (topk_indices != -1), other=-1) + tl.store( + topk_mem_index + cur_index * topk_mem_index_stride_b + offs_d * topk_mem_index_stride_k, + mem_index, + mask=mask, + ) @torch.no_grad() -def trans_topk_index_to_mem_index(topk_index: torch.Tensor, ragged_mem_index: torch.Tensor): - assert topk_index.shape[1] == 2048, f"Expected topk_index shape[1]=2048, got {topk_index.shape[1]}" +def trans_topk_index_to_mem_index( + topk_index: torch.Tensor, + ragged_start_index: torch.Tensor, + ragged_mem_index: torch.Tensor, +): + assert ragged_start_index.shape == (topk_index.shape[0],) grid = (topk_index.shape[0],) @@ -37,11 +51,13 @@ def trans_topk_index_to_mem_index(topk_index: torch.Tensor, ragged_mem_index: to topk_index=topk_index, topk_index_stride_b=topk_index.stride(0), topk_index_stride_k=topk_index.stride(1), + ragged_start_index=ragged_start_index, ragged_mem_index=ragged_mem_index, topk_mem_index=topk_mem_index, topk_mem_index_stride_b=topk_mem_index.stride(0), topk_mem_index_stride_k=topk_mem_index.stride(1), - BLOCK_DMODEL=2048, + topk_width=topk_index.shape[1], + BLOCK_DMODEL=triton.next_power_of_2(topk_index.shape[1]), num_warps=8, ) return topk_mem_index diff --git a/lightllm/models/gemma4/layer_infer/post_layer_infer.py b/lightllm/models/gemma4/layer_infer/post_layer_infer.py index b736a2d6c1..98158bd79c 100644 --- a/lightllm/models/gemma4/layer_infer/post_layer_infer.py +++ b/lightllm/models/gemma4/layer_infer/post_layer_infer.py @@ -5,18 +5,18 @@ class Gemma4PostLayerInfer(LlamaPostLayerInfer): """ Same final RMSNorm + tied lm_head path as Llama, with an extra tanh-based - logit softcap at the end: logits = softcap * tanh(logits / softcap). + transform before sampling: logits = softcap * tanh(logits / softcap). """ def __init__(self, network_config): super().__init__(network_config) self.final_logit_softcapping = float(network_config.get("final_logit_softcapping")) - def token_forward(self, input_embdings, infer_state, layer_weight): - logits = super().token_forward(input_embdings, infer_state, layer_weight) + def _apply_logit_postprocessing(self, logits: torch.Tensor) -> torch.Tensor: if self.final_logit_softcapping is not None and self.final_logit_softcapping > 0: cap = self.final_logit_softcapping - logits = torch.tanh(logits / cap) * cap - if infer_state.prompt_logics is not None: - infer_state.prompt_logics = torch.tanh(infer_state.prompt_logics / cap) * cap + # The historical dense path first materializes FP32 logits and then + # applies the softcap. Preserve that numerical contract when the + # transform runs on a local BF16/FP16 vocabulary shard. + return torch.tanh(logits.float() / cap) * cap return logits diff --git a/lightllm/models/glm5_next/__init__.py b/lightllm/models/glm5_next/__init__.py new file mode 100644 index 0000000000..728282aa05 --- /dev/null +++ b/lightllm/models/glm5_next/__init__.py @@ -0,0 +1 @@ +"""GLM-5-Next model support.""" diff --git a/lightllm/models/glm5_next/glm5_next_visual.py b/lightllm/models/glm5_next/glm5_next_visual.py new file mode 100644 index 0000000000..6825d49697 --- /dev/null +++ b/lightllm/models/glm5_next/glm5_next_visual.py @@ -0,0 +1,115 @@ +from __future__ import annotations + +import json +import os +from io import BytesIO +from typing import List + +import torch +from PIL import Image +from safetensors import safe_open +from transformers import AutoConfig, AutoProcessor + +from lightllm.server.embed_cache.utils import get_shm_name_data, read_shm +from lightllm.server.multimodal_params import ImageItem +from lightllm.utils.torch_dtype_utils import get_torch_dtype + + +class Glm5NextVisionModel: + """LightLLM visual-server adapter around the official GLM-5 vision tower.""" + + def __init__(self, data_type="bfloat16"): + self.data_type = data_type if isinstance(data_type, torch.dtype) else get_torch_dtype(data_type) + self.device = torch.device("cpu") + self.vision_tower = None + self.image_processor = None + + @staticmethod + def _weight_files(weight_dir): + index_path = os.path.join(weight_dir, "model.safetensors.index.json") + if os.path.exists(index_path): + with open(index_path, "r", encoding="utf-8") as stream: + weight_map = json.load(stream)["weight_map"] + return sorted(set(weight_map.values())) + return sorted(name for name in os.listdir(weight_dir) if name.endswith(".safetensors")) + + @classmethod + def _load_prefix_state_dict(cls, weight_dir, prefix): + state_dict = {} + for file_name in cls._weight_files(weight_dir): + with safe_open(os.path.join(weight_dir, file_name), framework="pt", device="cpu") as stream: + for key in stream.keys(): + if key.startswith(prefix): + state_dict[key[len(prefix) :]] = stream.get_tensor(key) + return state_dict + + def load_model(self, weight_dir): + try: + from transformers.models.glm5_next.modeling_glm5_next import ( + Glm5NextVisionModel as HFGlm5NextVisionModel, + ) + except ImportError as exc: + raise ImportError("GLM-5 vision requires a Transformers build with glm5_next support") from exc + + config = AutoConfig.from_pretrained(weight_dir, trust_remote_code=True) + if config.vision_config is None: + raise ValueError("GLM-5 checkpoint does not contain vision_config") + # Direct submodel construction bypasses AutoModel's normal attention + # implementation selection. SDPA keeps large-image attention + # memory-efficient and is available in the pinned PyTorch runtime. + config.vision_config._attn_implementation = "sdpa" + self.vision_tower = HFGlm5NextVisionModel(config.vision_config).eval() + self.image_processor = AutoProcessor.from_pretrained(weight_dir).image_processor + + state_dict = self._load_prefix_state_dict(weight_dir, "model.visual.") + missing, unexpected = self.vision_tower.load_state_dict(state_dict, strict=False) + if missing or unexpected: + raise RuntimeError(f"GLM-5 vision weight mismatch: missing={missing}, unexpected={unexpected}") + return self + + def cuda(self): + self.device = torch.device("cuda") + self.vision_tower = self.vision_tower.to(device=self.device, dtype=self.data_type) + return self + + @torch.inference_mode() + def forward(self, pixel_values, image_grid_thw): + output = self.vision_tower( + hidden_states=pixel_values.to(self.device, dtype=self.data_type, non_blocking=True), + grid_thw=image_grid_thw.to(self.device, non_blocking=True), + ) + return output.pooler_output.to(self.data_type) + + @torch.inference_mode() + def encode(self, images: List[ImageItem]): + pil_images = [] + uuids = [] + for image_item in images: + if not isinstance(image_item, ImageItem): + raise TypeError(f"Unsupported GLM-5 image input type: {type(image_item)}") + uuids.append(image_item.uuid) + image_data = read_shm(get_shm_name_data(image_item.uuid)) + with Image.open(BytesIO(image_data)) as image: + pil_images.append(image.convert("RGB")) + + if not pil_images: + return None + + image_inputs = self.image_processor(pil_images, return_tensors="pt") + pixel_values = image_inputs["pixel_values"] + image_grid_thw = image_inputs["image_grid_thw"] + + merge_area = self.image_processor.merge_size ** 2 + token_nums = [int(grid.prod().item() // merge_area) for grid in image_grid_thw] + valid_ids = [] + valid_start = 0 + for image_item, token_num in zip(images, token_nums): + if image_item.token_num is not None and image_item.token_num != token_num: + raise ValueError(f"GLM-5 image token mismatch: allocated={image_item.token_num}, encoded={token_num}") + valid_ids.append([valid_start, valid_start + token_num]) + valid_start += token_num + + image_embeds = self.forward(pixel_values, image_grid_thw) + if image_embeds.shape[0] != valid_start: + raise ValueError(f"GLM-5 image embed length mismatch: embeds={image_embeds.shape[0]}, tokens={valid_start}") + return image_embeds, uuids, valid_ids diff --git a/lightllm/models/glm5_next/infer_struct.py b/lightllm/models/glm5_next/infer_struct.py new file mode 100644 index 0000000000..74d7ae824c --- /dev/null +++ b/lightllm/models/glm5_next/infer_struct.py @@ -0,0 +1,5 @@ +from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo + + +class Glm5NextInferStateInfo(Deepseek2InferStateInfo): + pass diff --git a/lightllm/models/glm5_next/layer_infer/__init__.py b/lightllm/models/glm5_next/layer_infer/__init__.py new file mode 100644 index 0000000000..1235114fbb --- /dev/null +++ b/lightllm/models/glm5_next/layer_infer/__init__.py @@ -0,0 +1,3 @@ +from .transformer_layer_infer import Glm5NextTransformerLayerInfer + +__all__ = ["Glm5NextTransformerLayerInfer"] diff --git a/lightllm/models/glm5_next/layer_infer/pre_layer_infer.py b/lightllm/models/glm5_next/layer_infer/pre_layer_infer.py new file mode 100644 index 0000000000..ff0e48c71a --- /dev/null +++ b/lightllm/models/glm5_next/layer_infer/pre_layer_infer.py @@ -0,0 +1,7 @@ +from lightllm.models.qwen_vl.layer_infer.pre_layer_infer import LlamaMultimodalPreLayerInfer + + +class Glm5NextMultimodalPreLayerInfer(LlamaMultimodalPreLayerInfer): + """Merge cached GLM-5 vision embeddings into the text-token stream.""" + + pass diff --git a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py new file mode 100644 index 0000000000..b71ad99947 --- /dev/null +++ b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py @@ -0,0 +1,301 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import torch + +from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward +from lightllm.distributed.communication_op import all_reduce +from lightllm.models.deepseek3_2.layer_infer.transformer_layer_infer import ( + Deepseek3_2TransformerLayerInfer, + NsaInfer, +) +from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import ( + silu_and_mul_fwd, +) +from lightllm.models.glm5_next.triton_kernel.mhc import ( + hc_contract, + hc_expand, + hc_post, + hc_pre_norm, +) +from lightllm.common.triton_utils.autotuner import Autotuner +from lightllm.utils.envs_utils import get_env_start_args + + +class Glm5NextNsaInfer(NsaInfer): + """GLM indexer projection without rotary dimensions.""" + + def _get_q_k_bf16(self, hidden_states, q_lora, infer_state, layer_weight): + q = layer_weight.wq_b_proj_.mm(q_lora).view(-1, self.tp_index_n_heads, self.index_head_dim) + k = layer_weight.wk_proj_.mm(hidden_states.to(q_lora.dtype)) + k = layer_weight.k_norm_(k, eps=self.eps) + return q, k, k + + def _quantize_indexer_activation(self, value: torch.Tensor): + from lightllm.models.deepseek3_2.triton_kernel.hadamard_transform import ( + hadamard_transform_quant_fp8, + ) + + assert self.block_size == 128 and self.scale_fmt == "ue8m0" + return hadamard_transform_quant_fp8(value, scale=self.index_head_dim ** -0.5) + + def _scale_indexer_weights(self, weights: torch.Tensor, q_scale: torch.Tensor) -> torch.Tensor: + from lightllm.models.deepseek3_2.triton_kernel.indexer_weight_scale import ( + scale_indexer_weights_, + ) + + return scale_indexer_weights_(weights, q_scale, self.index_n_heads_scale) + + def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weight): + # GLM stores weights_proj in FP32, so its activation must match before + # delegating to the shared NSA scoring and top-k implementation. + return super()._get_indices(hidden_states.float(), q_lora, infer_state, att_state, layer_weight) + + +class Glm5NextTransformerLayerInfer(Deepseek3_2TransformerLayerInfer): + def __init__(self, layer_num, network_config): + super().__init__(layer_num, network_config) + self.num_hidden_layers = network_config["num_hidden_layers"] + self.autotune_layer_num = network_config.get("autotune_layer_num", self.num_hidden_layers) + self.is_mtp_layer = layer_num >= self.num_hidden_layers + self.is_linear_attention_layer = ( + not self.is_mtp_layer and network_config["layer_types"][layer_num] == "linear_attention" + ) + self.mhc_streams = network_config.get("hc_mult", 4) + self.hc_eps = network_config.get("hc_eps", 1e-6) + self.hc_sinkhorn_iters = network_config.get("hc_sinkhorn_iters", 20) + self.swiglu_limit = network_config["swiglu_limit"] + linear = network_config["linear_attn_config"] + self.linear_num_heads = linear["num_heads"] + self.linear_head_dim = linear["head_dim"] + self.tp_linear_num_heads = self.linear_num_heads // self.tp_world_size_ + self.tp_linear_projection_size = self.tp_linear_num_heads * self.linear_head_dim + if not self.is_linear_attention_layer: + self.indexer = Glm5NextNsaInfer( + layer_idx=self.layer_num_, + network_config=self.network_config_, + tp_world_size=self.tp_world_size_, + ) + # GLM's recurrent EAGLE drafter processes one row per logical + # request. Only target-model decode uses the widened verification + # layout of mtp_step + 1 rows. + self.indexer.decode_mtp_step = 0 if self.is_mtp_layer else get_env_start_args().mtp_step + + def _ffn_tp(self, input, infer_state, layer_weight): + """Dense/shared GLM FFN with the checkpoint's clamp semantics.""" + + input = input.view(-1, self.embed_dim_) + up_gate_out = layer_weight.gate_up_proj.mm(input) + ffn1_out = self.alloc_tensor((input.size(0), up_gate_out.size(1) // 2), input.dtype) + silu_and_mul_fwd( + up_gate_out, + ffn1_out, + limit=self.swiglu_limit, + alpha=1.0, + clamp_up_add_one=False, + ) + return layer_weight.down_proj.mm(ffn1_out) + + def _shared_ffn_tp(self, input, infer_state, layer_weight): + return self._ffn_tp(input, infer_state, layer_weight) + + def _get_qkv(self, input, infer_state, layer_weight): + if self.is_linear_attention_layer: + raise AssertionError("KDA projections use _kda_projections") + + input = input.view(-1, self.embed_dim_) + if not infer_state.use_replicated_attention_ep: + input = self._tpsp_allgather(input=input, infer_state=infer_state) + if infer_state.need_dp_prefill_balance: + input = infer_state._all_to_all_unbalance_get(data=input) + + q, cache_kv = layer_weight.qkv_a_proj_with_mqa_.mm(input).split([self.q_lora_rank, self.kv_lora_rank], dim=-1) + q = rmsnorm_forward(q, weight=layer_weight.q_a_layernorm_.weight, eps=self.eps_) + infer_state.get_topk_indices_params = {"hidden_states": input, "q_lora": q} + q = layer_weight.q_b_proj_.mm(q).view(-1, self.tp_q_head_num_, self.qk_nope_head_dim) + cache_kv = cache_kv.view(-1, 1, self.kv_lora_rank) + rmsnorm_forward( + cache_kv[:, :, : self.kv_lora_rank], + weight=layer_weight.kv_a_layernorm_.weight, + eps=self.eps_, + out=cache_kv[:, :, : self.kv_lora_rank], + ) + return q, cache_kv + + def _get_o(self, input, infer_state, layer_weight): + if not infer_state.use_replicated_attention_ep: + return super()._get_o(input, infer_state, layer_weight) + + if infer_state.need_dp_prefill_balance: + input = infer_state._all_to_all_balance_get(data=input) + if input.shape[2] == self.kv_lora_rank: + input = layer_weight.v_b_proj_.bmm(input.transpose(0, 1)).transpose(0, 1) + output = layer_weight.o_weight_.mm(input.reshape(-1, self.tp_q_head_num_ * self.v_head_dim)) + all_reduce(output, group=infer_state.dist_group) + return output + + def _token_attention_kernel(self, q, infer_state, layer_weight, out=None): + if self.is_linear_attention_layer: + raise AssertionError("KDA uses its dedicated backend") + q_nope = layer_weight.k_b_proj_.bmm(q.transpose(0, 1)).transpose(0, 1) + topk_mem_indices, _ = self.indexer._get_indices( + hidden_states=infer_state.get_topk_indices_params["hidden_states"], + q_lora=infer_state.get_topk_indices_params["q_lora"], + infer_state=infer_state, + att_state=infer_state.decode_att_state, + layer_weight=layer_weight, + ) + del infer_state.get_topk_indices_params + q_rope = q_nope[..., :0] + return infer_state.decode_att_state.decode_att( + q=(q_nope, q_rope), + k=infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_), + v=None, + att_control=AttControl( + nsa_decode=True, + nsa_decode_dict={ + "layer_index": self.layer_num_, + "topk_mem_indices": topk_mem_indices, + "softmax_scale": self.softmax_scale, + "kv_lora_rank": self.kv_lora_rank, + "qk_rope_head_dim": 0, + }, + ), + ) + + def _kda_projections(self, input, infer_state, layer_weight): + # KDA shards heads across TP ranks, so every rank still needs every + # token before updating its recurrent head state. In TP/SP mode the + # layer input is sequence-sharded; gather it here just like the MLA + # projection path and reduce-scatter the output in _kda_post. + input = input.view(-1, self.embed_dim_) + if not infer_state.use_replicated_attention_ep: + input = self._tpsp_allgather(input=input, infer_state=infer_state) + projected = layer_weight.linear_qkvbfg_a_proj.mm(input) + qkv_size = 3 * self.tp_linear_projection_size + mixed_qkv, raw_beta, f_a, g_a = projected.split( + [ + qkv_size, + self.tp_linear_num_heads, + self.linear_head_dim, + self.linear_head_dim, + ], + dim=-1, + ) + raw_gate, norm_gate = layer_weight.project_kda_fg_b(f_a, g_a) + return mixed_qkv, raw_gate, raw_beta, norm_gate + + def _kda_post(self, core_output, norm_gate, infer_state, layer_weight): + tokens = norm_gate.shape[0] + core_output = core_output.view(-1, self.linear_head_dim) + norm_gate = norm_gate.view(tokens, self.tp_linear_num_heads, self.linear_head_dim) + output = layer_weight.linear_o_norm( + input=core_output, + gate_value=norm_gate, + eps=self.eps_, + alloc_func=self.alloc_tensor, + ) + output = layer_weight.linear_o_proj.mm(output.view(tokens, -1)) + if infer_state.use_replicated_attention_ep: + all_reduce(output, group=infer_state.dist_group) + return output + return self._tpsp_reduce(input=output, infer_state=infer_state) + + def context_attention_forward(self, input_embeddings, infer_state, layer_weight): + if not self.is_linear_attention_layer: + return super().context_attention_forward(input_embeddings, infer_state, layer_weight) + mixed_qkv, raw_gate, raw_beta, norm_gate = self._kda_projections(input_embeddings, infer_state, layer_weight) + core_output = infer_state.prefill_att_state1.prefill_att( + q=None, + k=None, + v=None, + att_control=AttControl( + linear_att_prefill=True, + linear_att_prefill_dict={ + "mixed_qkv": mixed_qkv, + "raw_gate": raw_gate, + "raw_beta": raw_beta, + "layer_weight": layer_weight, + "layer_num": self.layer_num_, + }, + ), + alloc_func=self.alloc_tensor, + ) + return self._kda_post(core_output, norm_gate, infer_state, layer_weight) + + def token_attention_forward(self, input_embeddings, infer_state, layer_weight): + if not self.is_linear_attention_layer: + return super().token_attention_forward(input_embeddings, infer_state, layer_weight) + mixed_qkv, raw_gate, raw_beta, norm_gate = self._kda_projections(input_embeddings, infer_state, layer_weight) + core_output = infer_state.decode_att_state1.decode_att( + q=None, + k=None, + v=None, + att_control=AttControl( + linear_att_decode=True, + linear_att_decode_dict={ + "mixed_qkv": mixed_qkv, + "raw_gate": raw_gate, + "raw_beta": raw_beta, + "layer_weight": layer_weight, + "layer_num": self.layer_num_, + }, + ), + alloc_func=self.alloc_tensor, + ) + return self._kda_post(core_output, norm_gate, infer_state, layer_weight) + + def _hc_pre(self, streams, layer_weight, prefix, norm_weight): + return hc_pre_norm( + x=streams, + fn=getattr(layer_weight, f"hc_{prefix}_fn").weight, + scale=getattr(layer_weight, f"hc_{prefix}_scale").weight, + base=getattr(layer_weight, f"hc_{prefix}_base").weight, + norm_weight=norm_weight.weight, + streams=self.mhc_streams, + rms_eps=self.eps_, + norm_eps=self.eps_, + hc_eps=self.hc_eps, + sinkhorn_iters=self.hc_sinkhorn_iters, + ) + + def _forward_mhc(self, input_embeddings, infer_state, layer_weight, *, prefill): + streams = input_embeddings + if self.layer_num_ == 0: + streams = hc_expand(streams.view(-1, self.embed_dim_), self.mhc_streams) + + layer_input, residual_mix, post_mix = self._hc_pre(streams, layer_weight, "attn", layer_weight.att_norm_weight_) + if prefill: + layer_output = self.context_attention_forward(layer_input, infer_state, layer_weight) + else: + layer_output = self.token_attention_forward(layer_input, infer_state, layer_weight) + streams = hc_post(layer_output, streams, residual_mix, post_mix, self.mhc_streams) + + layer_input, residual_mix, post_mix = self._hc_pre(streams, layer_weight, "ffn", layer_weight.ffn_norm_weight_) + if infer_state.use_replicated_attention_ep: + if self.is_moe: + local_input = self._tpsp_sp_split(input=layer_input, infer_state=infer_state) + local_output = self._ffn(local_input, infer_state, layer_weight) + layer_output = self._tpsp_allgather(input=local_output, infer_state=infer_state) + else: + layer_output = self._ffn_tp(layer_input, infer_state, layer_weight) + all_reduce(layer_output, group=infer_state.dist_group) + else: + layer_output = self._ffn(layer_input, infer_state, layer_weight) + streams = hc_post(layer_output, streams, residual_mix, post_mix, self.mhc_streams) + is_autotune_last_layer = Autotuner.is_autotune_warmup() and self.layer_num_ == self.autotune_layer_num - 1 + if self.layer_num_ == self.num_hidden_layers - 1 or is_autotune_last_layer: + return hc_contract(streams, self.mhc_streams) + return streams + + def context_forward(self, input_embeddings, infer_state, layer_weight): + if self.is_mtp_layer: + return super().context_forward(input_embeddings, infer_state, layer_weight) + return self._forward_mhc(input_embeddings, infer_state, layer_weight, prefill=True) + + def token_forward(self, input_embeddings, infer_state, layer_weight): + if self.is_mtp_layer: + return super().token_forward(input_embeddings, infer_state, layer_weight) + return self._forward_mhc(input_embeddings, infer_state, layer_weight, prefill=False) diff --git a/lightllm/models/glm5_next/layer_weights/__init__.py b/lightllm/models/glm5_next/layer_weights/__init__.py new file mode 100644 index 0000000000..4136e1883d --- /dev/null +++ b/lightllm/models/glm5_next/layer_weights/__init__.py @@ -0,0 +1,4 @@ +from .pre_and_post_layer_weight import Glm5NextPreAndPostLayerWeight +from .transformer_layer_weight import Glm5NextTransformerLayerWeight + +__all__ = ["Glm5NextPreAndPostLayerWeight", "Glm5NextTransformerLayerWeight"] diff --git a/lightllm/models/glm5_next/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/glm5_next/layer_weights/pre_and_post_layer_weight.py new file mode 100644 index 0000000000..b336f00632 --- /dev/null +++ b/lightllm/models/glm5_next/layer_weights/pre_and_post_layer_weight.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 + +from lightllm.models.llama.layer_weights.pre_and_post_layer_weight import ( + LlamaPreAndPostLayerWeight, +) + + +def add_language_model_aliases(weights: dict) -> None: + """Expose GLM's nested language-model keys under LightLLM names.""" + + prefix = "model.language_model." + for name in list(weights): + if name.startswith(prefix): + weights.setdefault("model." + name[len(prefix) :], weights[name]) + + +class Glm5NextPreAndPostLayerWeight(LlamaPreAndPostLayerWeight): + def load_hf_weights(self, weights): + add_language_model_aliases(weights) + return super().load_hf_weights(weights) diff --git a/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py b/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py new file mode 100644 index 0000000000..32e712f010 --- /dev/null +++ b/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py @@ -0,0 +1,328 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import torch + +from lightllm.common.basemodel.layer_weights.transformer_layer_weight import ( + TransformerLayerWeight, +) +from lightllm.common.basemodel.layer_weights.meta_weights import ( + COLMMWeight, + GatedRMSNormWeight, + LayerNormWeight, + ParameterWeight, + RMSNormWeight, + ROWMMWeight, + TpParameterWeight, +) +from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_slicer import ( + get_row_slice_mixin, +) +from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_weight import ( + MMWeightTpl, +) +from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size +from lightllm.models.deepseek2.layer_weights.transformer_layer_weight import ( + Deepseek2TransformerLayerWeight, +) +from lightllm.models.deepseek3_2.layer_weights.transformer_layer_weight import ( + Deepseek3_2TransformerLayerWeight, +) +from .pre_and_post_layer_weight import add_language_model_aliases + + +class Glm5NextMergedKdaProjection(MMWeightTpl): + """One KDA input GEMM with TP-sharded q/k/v/b and replicated f_a/g_a.""" + + def __init__( + self, + in_dim: int, + projection: int, + head_count: int, + head_dim: int, + weight_names: list[str], + data_type: torch.dtype, + tp_rank: int | None = None, + tp_world_size: int | None = None, + ): + tp_rank = get_current_rank_in_dp() if tp_rank is None else tp_rank + tp_world_size = get_dp_world_size() if tp_world_size is None else tp_world_size + assert projection % tp_world_size == 0 + assert head_count % tp_world_size == 0 + super().__init__( + in_dim=in_dim, + out_dims=[ + projection // tp_world_size, + projection // tp_world_size, + projection // tp_world_size, + head_count // tp_world_size, + head_dim, + head_dim, + ], + weight_names=weight_names, + bias_names=None, + data_type=data_type, + quant_method=None, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ) + self.sharded_slicer = get_row_slice_mixin("none", tp_rank=tp_rank, tp_world_size=tp_world_size) + self.replicated_slicer = get_row_slice_mixin("none", tp_rank=0, tp_world_size=1) + + def _get_param_slicer(self, sub_child_index: int): + return self.replicated_slicer if sub_child_index >= 4 else self.sharded_slicer + + +class Glm5NextTransformerLayerWeight(Deepseek3_2TransformerLayerWeight): + def _parse_config(self): + super()._parse_config() + # The released sparse MLA keeps kv_b_proj in BF16 even though the + # surrounding projections are native FP8. Its compressed-context + # shortcut assumes a quantized kv_b matrix, so GLM uses the BMM split. + self.enable_cc_method = False + self.is_mtp_layer = self.layer_num_ >= self.network_config_["num_hidden_layers"] + self.is_linear_attention_layer = ( + not self.is_mtp_layer and self.network_config_["layer_types"][self.layer_num_] == "linear_attention" + ) + linear = self.network_config_["linear_attn_config"] + self.linear_num_heads = linear["num_heads"] + self.linear_head_dim = linear["head_dim"] + self.linear_projection_size = self.linear_num_heads * self.linear_head_dim + self.linear_conv_kernel_size = linear["short_conv_kernel_size"] + self.mhc_streams = self.network_config_.get("hc_mult", 4) + + def _init_weight(self): + if self.is_linear_attention_layer: + self._init_kda() + else: + Deepseek2TransformerLayerWeight._init_qkvo(self) + self._init_indexer_weight() + + if self.is_moe: + self._init_moe() + else: + self._init_ffn() + self._init_glm_norms() + if not self.is_mtp_layer: + self._init_mhc() + + def _init_kda(self): + prefix = f"model.layers.{self.layer_num_}.self_attn" + projection = self.linear_projection_size + head_count = self.linear_num_heads + head_dim = self.linear_head_dim + + self.linear_qkvbfg_a_proj = Glm5NextMergedKdaProjection( + in_dim=self.n_embed, + projection=projection, + head_count=head_count, + head_dim=head_dim, + weight_names=[ + f"{prefix}.q_proj.weight", + f"{prefix}.k_proj.weight", + f"{prefix}.v_proj.weight", + f"{prefix}.b_proj.weight", + f"{prefix}.f_a_proj.weight", + f"{prefix}.g_a_proj.weight", + ], + data_type=self.data_type_, + ) + self.linear_fg_b_proj = ROWMMWeight( + in_dim=head_dim, + out_dims=[projection, projection], + weight_names=[f"{prefix}.f_b_proj.weight", f"{prefix}.g_b_proj.weight"], + data_type=self.data_type_, + quant_method=None, + ) + self.linear_qkv_conv1d = ROWMMWeight( + in_dim=self.linear_conv_kernel_size, + out_dims=[projection, projection, projection], + weight_names=[ + f"{prefix}.q_conv1d.weight", + f"{prefix}.k_conv1d.weight", + f"{prefix}.v_conv1d.weight", + ], + data_type=self.data_type_, + quant_method=None, + ) + self.linear_A_log = TpParameterWeight( + weight_name=f"{prefix}.A_log", + data_type=torch.float32, + weight_shape=(head_count,), + ) + self.linear_dt_bias = TpParameterWeight( + weight_name=f"{prefix}.dt_bias", + data_type=torch.float32, + weight_shape=(projection,), + ) + self.linear_o_norm = GatedRMSNormWeight( + dim=head_dim, + weight_name=f"{prefix}.o_norm.weight", + data_type=self.data_type_, + activation="sigmoid", + ) + self.linear_o_proj = COLMMWeight( + in_dim=projection, + out_dims=[self.n_embed], + weight_names=f"{prefix}.o_proj.weight", + data_type=self.data_type_, + quant_method=None, + ) + + def _init_indexer_weight(self): + """Initialize GLM's NoPE, K-pool indexer parameters. + + The head-weight projection intentionally accumulates in fp32. Both + reference engines do this because bf16 head weights can change close + K-pool rankings on difficult long-context prompts. + """ + + prefix = f"model.layers.{self.layer_num_}.self_attn.indexer" + self.wq_b_proj_ = ROWMMWeight( + in_dim=self.q_lora_rank, + out_dims=[self.index_n_heads * self.index_head_dim], + weight_names=f"{prefix}.wq_b.weight", + data_type=self.data_type_, + quant_method=None, + ) + self.wk_proj_ = ROWMMWeight( + in_dim=self.hidden_size, + out_dims=[self.index_head_dim], + weight_names=f"{prefix}.wk.weight", + data_type=self.data_type_, + quant_method=None, + tp_rank=0, + tp_world_size=1, + ) + self.k_norm_ = LayerNormWeight( + dim=self.index_head_dim, + weight_name=f"{prefix}.k_norm.weight", + data_type=self.data_type_, + bias_name=f"{prefix}.k_norm.bias", + ) + self.weights_proj_ = ROWMMWeight( + in_dim=self.hidden_size, + out_dims=[self.index_n_heads], + weight_names=f"{prefix}.weights_proj.weight", + data_type=torch.float32, + quant_method=None, + tp_rank=0, + tp_world_size=1, + ) + self.index_kpool_compress_gate = ROWMMWeight( + in_dim=self.hidden_size, + out_dims=[self.index_head_dim], + weight_names=f"{prefix}.index_kpool_compress_gate", + data_type=self.data_type_, + quant_method=None, + tp_rank=0, + tp_world_size=1, + ) + self.index_kpool_compress_ape = ParameterWeight( + weight_name=f"{prefix}.index_kpool_compress_ape", + data_type=torch.float32, + weight_shape=(self.network_config_["index_kpool"], self.index_head_dim), + ) + + def _init_glm_norms(self): + prefix = f"model.layers.{self.layer_num_}" + self.att_norm_weight_ = RMSNormWeight( + dim=self.n_embed, + weight_name=f"{prefix}.input_layernorm.weight", + data_type=self.data_type_, + ) + self.ffn_norm_weight_ = RMSNormWeight( + dim=self.n_embed, + weight_name=f"{prefix}.post_attention_layernorm.weight", + data_type=self.data_type_, + ) + if not self.is_linear_attention_layer: + self.kv_a_layernorm_ = RMSNormWeight( + dim=self.kv_lora_rank, + weight_name=f"{prefix}.self_attn.kv_a_layernorm.weight", + data_type=self.data_type_, + ) + self.q_a_layernorm_ = RMSNormWeight( + dim=self.q_lora_rank, + weight_name=f"{prefix}.self_attn.q_a_layernorm.weight", + data_type=self.data_type_, + ) + + def _init_mhc(self): + prefix = f"model.layers.{self.layer_num_}" + streams = self.mhc_streams + mix_size = (2 + streams) * streams + flattened_hidden = streams * self.n_embed + self.hc_attn_fn = ParameterWeight( + weight_name=f"{prefix}.hc_attn_fn", + data_type=torch.float32, + weight_shape=(mix_size, flattened_hidden), + ) + self.hc_attn_base = ParameterWeight( + weight_name=f"{prefix}.hc_attn_base", + data_type=torch.float32, + weight_shape=(mix_size,), + ) + self.hc_attn_scale = ParameterWeight( + weight_name=f"{prefix}.hc_attn_scale", + data_type=torch.float32, + weight_shape=(3,), + ) + self.hc_ffn_fn = ParameterWeight( + weight_name=f"{prefix}.hc_ffn_fn", + data_type=torch.float32, + weight_shape=(mix_size, flattened_hidden), + ) + self.hc_ffn_base = ParameterWeight( + weight_name=f"{prefix}.hc_ffn_base", + data_type=torch.float32, + weight_shape=(mix_size,), + ) + self.hc_ffn_scale = ParameterWeight( + weight_name=f"{prefix}.hc_ffn_scale", + data_type=torch.float32, + weight_shape=(3,), + ) + + def get_merged_kda_conv_weight(self): + return self.linear_qkv_conv1d.mm_param.weight + + def project_kda_fg_b(self, f_a: torch.Tensor, g_a: torch.Tensor): + method = self.linear_fg_b_proj.quant_method + f = method.apply(f_a, self.linear_fg_b_proj.mm_param_list[0]) + g = method.apply(g_a, self.linear_fg_b_proj.mm_param_list[1]) + return f, g + + def _preprocess_kda_weights(self, weights): + prefix = f"model.layers.{self.layer_num_}.self_attn" + for projection in ("q", "k", "v"): + name = f"{prefix}.{projection}_conv1d.weight" + if name in weights and weights[name].ndim == 3: + weights[name] = weights[name].squeeze(1) + + def load_hf_weights(self, weights): + add_language_model_aliases(weights) + + # GLM checkpoints nest the shared expert under + # ``mlp.shared_experts``. This class deliberately bypasses + # Deepseek2TransformerLayerWeight.load_hf_weights below, so perform + # the fused-shared remap here before the generic loader consumes the + # expert tensors. + if self.num_fused_shared_experts > 0: + self._rename_shared_experts( + weights, + self.experts.quant_method.weight_scale_suffix, + ) + + if self.is_linear_attention_layer: + self._preprocess_kda_weights(weights) + return TransformerLayerWeight.load_hf_weights(self, weights) + + kv_b_name = f"model.layers.{self.layer_num_}.self_attn.kv_b_proj.weight" + if kv_b_name in weights: + k_b_proj, v_b_proj = self._split_kv_b_proj(weights[kv_b_name]) + weights[f"model.layers.{self.layer_num_}.self_attn.k_b_proj.weight"] = k_b_proj + weights[f"model.layers.{self.layer_num_}.self_attn.v_b_proj.weight"] = v_b_proj + + return TransformerLayerWeight.load_hf_weights(self, weights) diff --git a/lightllm/models/glm5_next/mem_manager.py b/lightllm/models/glm5_next/mem_manager.py new file mode 100644 index 0000000000..20ab8e445f --- /dev/null +++ b/lightllm/models/glm5_next/mem_manager.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from lightllm.common.kv_cache_mem_manager.operator import LinearAttMemOperator +from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import ( + Qwen3NextMemManager, +) + + +class Glm5NextMemOperator(LinearAttMemOperator): + def copy_kv_to_mem_manager(self, layer_index, mem_index, kv): + layer_index = self.linear_config.get_full_att_kv_layer_index(layer_index) + from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import ( + destindex_copy_kv, + ) + + output = self.mem_manager.kv_buffer[layer_index][:, :, : self.mem_manager.mla_head_dim] + destindex_copy_kv(kv, mem_index, output) + + +class Glm5NextMemManager(Qwen3NextMemManager): + """Packed sparse-MLA KV, DSA index keys, and KDA recurrent states.""" + + operator_class = Glm5NextMemOperator + indexer_padding_bytes = 144 + indexer_payload_bytes = 132 + + def __init__(self, *args, **kwargs): + self.mla_head_dim = kwargs.get("head_dim", args[3] if len(args) > 3 else None) + super().__init__(*args, **kwargs) + + def get_cell_size(self): + bytes_per_token = self.mla_head_dim * self.dtype.itemsize + self.indexer_padding_bytes + return bytes_per_token * self.layer_num + + def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): + assert head_num == 1 and dtype in (torch.bfloat16, torch.float16) + padding_elements = self.indexer_padding_bytes // dtype.itemsize + self.kv_buffer = torch.empty( + (layer_num, size + 1, head_num, head_dim + padding_elements), + dtype=dtype, + device="cuda", + ) + self._init_linear_att_buffers() + + def get_att_input_params(self, layer_index): + packed_index = self.linear_config.get_full_att_kv_layer_index(layer_index) + return self.kv_buffer[packed_index][:, :, : self.mla_head_dim] + + def get_indexer_k_buffer(self, layer_index): + packed_index = self.linear_config.get_full_att_kv_layer_index(layer_index) + return self.kv_buffer[packed_index].view(torch.uint8)[:, :, -self.indexer_payload_bytes :] diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py new file mode 100644 index 0000000000..f74ced0f78 --- /dev/null +++ b/lightllm/models/glm5_next/model.py @@ -0,0 +1,206 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import os + +import torch + +from lightllm.common.basemodel.hidden_collector import FinalHiddenCollector +from lightllm.common.basemodel.attention.linear import KDALinearAttBackend +from lightllm.common.build_utils import repair_config +from lightllm.common.linear_att_cache_manager.config_objs import LinearAttCacheConfig +from lightllm.common.req_manager import ReqManagerForMamba +from lightllm.distributed.communication_op import dist_group_manager +from lightllm.models.deepseek3_2.model import Deepseek3_2TpPartModel +from lightllm.models.glm5_next.infer_struct import Glm5NextInferStateInfo +from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import ( + Glm5NextTransformerLayerInfer, +) +from lightllm.models.glm5_next.layer_infer.pre_layer_infer import ( + Glm5NextMultimodalPreLayerInfer, +) +from lightllm.models.glm5_next.layer_weights.pre_and_post_layer_weight import ( + Glm5NextPreAndPostLayerWeight, +) +from lightllm.models.glm5_next.layer_weights.transformer_layer_weight import ( + Glm5NextTransformerLayerWeight, +) +from lightllm.models.glm5_next.mem_manager import Glm5NextMemManager +from lightllm.models.registry import ModelRegistry +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args + + +class Glm5NextPostNormHiddenCollector(FinalHiddenCollector): + """Expose the post-final-norm hidden state expected by GLM NextN. + + The generic EAGLE collector returns the decoder output before the model's + final RMSNorm. GLM's NextN block was trained against the normalized target + hidden and also recycles its own normalized hidden between recurrent draft + steps (matching the vLLM and SGLang implementations). + """ + + def __init__(self, norm_weight, eps: float): + super().__init__() + self.norm_weight = norm_weight + self.eps = eps + self.draft_token_ids = None + self.draft_token_probs = None + + def new_instance(self): + return Glm5NextPostNormHiddenCollector(self.norm_weight, self.eps) + + def add_final_hidden(self, final_hidden: torch.Tensor) -> None: + self.final_hidden = self.norm_weight(input=final_hidden, eps=self.eps) + + def add_mtp_outputs( + self, + draft_token_ids: torch.Tensor | None, + confidence_logits: torch.Tensor | None, + draft_token_probs: torch.Tensor | None = None, + ) -> None: + assert confidence_logits is None + self.draft_token_ids = draft_token_ids + self.draft_token_probs = draft_token_probs + + def finish_output(self, infer_state): + output = super().finish_output(infer_state) + output.draft_token_ids = self.draft_token_ids + output.draft_token_probs = self.draft_token_probs + self.draft_token_ids = None + self.draft_token_probs = None + return output + + +@ModelRegistry(["glm5_next", "glm5_next_text"]) +class Glm5NextTpPartModel(Deepseek3_2TpPartModel): + # Keep attention/mHC rows replicated so their row-parallel projections can + # use the fast custom all-reduce. Only shard rows around EP MoE, then + # all-gather once. The recurrent draft model opts out below. + replicated_attention_ep = True + pre_and_post_weight_class = Glm5NextPreAndPostLayerWeight + transformer_weight_class = Glm5NextTransformerLayerWeight + transformer_layer_infer_class = Glm5NextTransformerLayerInfer + infer_state_class = Glm5NextInferStateInfo + + def _init_config(self): + with open(os.path.join(self.weight_dir_, "config.json"), "r") as config_file: + outer_config = json.load(config_file) + self.outer_config = outer_config + self.config = dict(outer_config.get("text_config", outer_config)) + self.vision_config = outer_config.get("vision_config") + if "quantization_config" in outer_config: + self.config["quantization_config"] = outer_config["quantization_config"] + # GLM-5.3's index cache follows the official ue8m0-scale path; + # the released HF config does not spell this implementation detail + # out in quantization_config. + self.config["quantization_config"].setdefault("scale_fmt", "ue8m0") + # The checkpoint uses the standard clamped SwiGLU, not GPT-OSS's + # clamped (up + 1) variant supported by the shared kernel. + self.config["swiglu_clamp_up_add_one"] = False + # The generic autotune warmup executes only a representative prefix. + # Let the last layer in that prefix contract mHC's residual streams + # before the shared LM head consumes its output. + self.config["autotune_layer_num"] = 4 + repair_config(self.config, same_names=["num_attention_heads", "n_head"]) + repair_config(self.config, same_names=["hidden_size", "n_embd", "n_embed"]) + repair_config(self.config, same_names=["num_hidden_layers", "n_layer"]) + if self.finetune_config: + self.config["vocab_size"] = self.finetune_config.vocab_size + + def autotune_layers(self): + return 4 + + def _init_hidden_collector(self): + collector = self.mtp_manager.create_hidden_collector(model=self) + if isinstance(collector, FinalHiddenCollector): + collector = Glm5NextPostNormHiddenCollector( + norm_weight=self.pre_post_weight.final_norm_weight_, + eps=self.config["rms_norm_eps"], + ) + self.hidden_collector_prototype = collector + + def _make_linear_config(self): + linear = self.config["linear_attn_config"] + start_args: StartArgs = get_env_start_args() + state_dtypes = {"bfloat16": torch.bfloat16, "float32": torch.float32} + return LinearAttCacheConfig( + tp_world_size=self.tp_world_size_, + full_att_all_num_kv_heads=1, + full_att_dtype=self.data_type, + full_att_num_kv_heads=1, + full_att_head_dim=self.config["kv_lora_rank"], + global_linear_k_heads=linear["num_heads"], + global_linear_v_heads=linear["num_heads"], + num_linear_k_heads=linear["num_heads"] // self.tp_world_size_, + num_linear_v_heads=linear["num_heads"] // self.tp_world_size_, + head_linear_k_dim=linear["head_dim"], + head_linear_v_dim=linear["head_dim"], + conv_kernel_size=linear["short_conv_kernel_size"], + linear_layer_num=len(linear["kda_layers"]), + conv_state_dtype=self.data_type, + ssm_state_dtype=state_dtypes[start_args.linear_att_ssm_data_type], + full_attention_interval=4, + all_layer_num=self.config["n_layer"], + draft_full_att_kv_layer_num=get_added_mtp_kv_layer_num(), + ) + + def _init_req_manager(self): + max_sequence_length = 0 + if self.batch_max_tokens is not None: + max_sequence_length = max(max_sequence_length, self.batch_max_tokens) + if self.max_seq_length is not None: + max_sequence_length = max(max_sequence_length, self.max_seq_length) + self.linear_config = self._make_linear_config() + self.req_manager = ReqManagerForMamba( + self.max_req_num, + max_sequence_length, + None, + linear_config=self.linear_config, + ) + + def _init_mem_manager(self): + self.linear_config = getattr(self, "linear_config", self._make_linear_config()) + self.mem_manager = Glm5NextMemManager( + size=self.max_total_token_num, + dtype=self.data_type, + num_kv_heads=1, + head_dim=self.config["kv_lora_rank"], + full_att_layer_num=self.linear_config.get_full_att_kv_layer_num_with_draft_model(), + linear_config=self.linear_config, + mem_fraction=self.mem_fraction, + ) + + def _init_att_backend1(self): + if getattr(self, "is_mtp_draft_model", False): + self.prefill_att_backend1 = None + self.decode_att_backend1 = None + return + self.prefill_att_backend1 = KDALinearAttBackend(model=self) + self.decode_att_backend1 = self.prefill_att_backend1 + + def _init_custom(self): + # GLM-5 sparse MLA is entirely NoPE. Keep zero-width tables so the + # generic infer-state position setup remains valid without allocating + # a million-token rotary cache. + max_length = max(self.config["max_position_embeddings"], self.max_seq_length or 0) + self._cos_cached = torch.empty((max_length, 0), dtype=self.data_type, device="cuda") + self._sin_cached = torch.empty_like(self._cos_cached) + dist_group_manager.new_deepep_group( + n_routed_experts=self.config["n_routed_experts"], + hidden_size=self.config["hidden_size"], + expert_quant_method_names=dist_group_manager.get_moe_quant_methods(self.trans_layers_weight), + num_experts_per_tok=self.config["num_experts_per_tok"], + moe_intermediate_size=self.config["moe_intermediate_size"], + ) + + +@ModelRegistry( + "glm5_next", + is_multimodal=True, + condition=lambda model_cfg: model_cfg.get("vision_config") is not None, +) +class Glm5NextMultimodalTpPartModel(Glm5NextTpPartModel): + pre_layer_infer_class = Glm5NextMultimodalPreLayerInfer diff --git a/lightllm/models/glm5_next/tokenizer.py b/lightllm/models/glm5_next/tokenizer.py new file mode 100644 index 0000000000..fbd2843ba7 --- /dev/null +++ b/lightllm/models/glm5_next/tokenizer.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +import copy +from typing import List, Union + +from lightllm.common.basemodel.multimodal_tokenizer import BaseMultiModalTokenizer +from lightllm.models.qwen2_vl.model import QWen2VLTokenizer +from lightllm.server.multimodal_params import ImageItem, MultimodalParams + + +class Glm5NextTokenizer(QWen2VLTokenizer): + """Multimodal tokenizer adapter for GLM-5 Next checkpoints. + + The released GLM-5.3-Flash chat template deliberately renders OpenAI + image parts as a text-only reminder. LightLLM carries image bytes through + ``MultimodalParams``, so template-facing image parts must instead become + one GLM image placeholder. ``encode`` (inherited from QWen2VLTokenizer) + replaces that placeholder with the virtual token range allocated by the + embedding cache. + """ + + image_placeholder = "<|begin_of_image|><|image|><|end_of_image|>" + + def __init__(self, tokenizer=None, image_processor=None, **kwargs): + BaseMultiModalTokenizer.__init__(self, tokenizer) + self.image_processor = image_processor + model_cfg = kwargs["model_cfg"] + self.image_start_id = model_cfg["image_start_token_id"] + self.image_end_id = model_cfg["image_end_token_id"] + self.image_token_id = model_cfg["image_token_id"] + self.patch_size = image_processor.patch_size + self.merge_size = image_processor.merge_size + self.min_image_tokens = image_processor.min_image_tokens + self.max_image_tokens = image_processor.max_image_tokens + + def get_image_token_length(self, img: ImageItem): + if img.image_w <= 0 or img.image_h <= 0: + raise ValueError(f"invalid GLM-5 image size: {img.image_w}x{img.image_h}") + patch_num = self.image_processor.get_number_of_image_patches(img.image_h, img.image_w) + token_num = patch_num // (self.merge_size ** 2) + if token_num <= 0: + raise ValueError(f"GLM-5 image produced no visual tokens: {img.image_w}x{img.image_h}") + return token_num + + def apply_chat_template(self, conversation=None, messages=None, **kwargs): + source = conversation if conversation is not None else messages + if source is None: + return self.tokenizer.apply_chat_template(conversation=conversation, messages=messages, **kwargs) + + normalized = copy.deepcopy(source) + for message in normalized: + content = message.get("content") + if not isinstance(content, list): + continue + rendered_parts = [] + for part in content: + if isinstance(part, dict) and part.get("type") in ("image", "image_url"): + rendered_parts.append({"type": "text", "text": self.image_placeholder}) + else: + rendered_parts.append(part) + message["content"] = rendered_parts + + if conversation is not None: + return self.tokenizer.apply_chat_template(conversation=normalized, **kwargs) + return self.tokenizer.apply_chat_template(messages=normalized, **kwargs) + + def encode( + self, + prompt: Union[str, List[int]], + multimodal_params: MultimodalParams = None, + **kwargs, + ): + return super().encode(prompt, multimodal_params=multimodal_params, **kwargs) diff --git a/lightllm/models/glm5_next/triton_kernel/__init__.py b/lightllm/models/glm5_next/triton_kernel/__init__.py new file mode 100644 index 0000000000..7db51ba063 --- /dev/null +++ b/lightllm/models/glm5_next/triton_kernel/__init__.py @@ -0,0 +1,3 @@ +from .mhc import hc_contract, hc_expand, hc_post, hc_pre, hc_pre_norm + +__all__ = ["hc_contract", "hc_expand", "hc_post", "hc_pre", "hc_pre_norm"] diff --git a/lightllm/models/glm5_next/triton_kernel/mhc.py b/lightllm/models/glm5_next/triton_kernel/mhc.py new file mode 100644 index 0000000000..87eadf7b50 --- /dev/null +++ b/lightllm/models/glm5_next/triton_kernel/mhc.py @@ -0,0 +1,547 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""mHC operators used by GLM-5-Next. + +The public entry points use fused Triton kernels for the small, launch-bound +mixing operations. The explicit PyTorch implementations remain available as +correctness oracles: all mixing math is accumulated in fp32 and only the +collapsed layer input / expanded residual streams are cast back to the +activation dtype. +""" + +from __future__ import annotations + +from typing import Tuple + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + + +@triton.jit +def _hc_prepare_kernel( + mixes, + scale, + base, + pre, + post, + residual_mix, + mix_stride_m: tl.constexpr, + pre_stride_m: tl.constexpr, + residual_stride_m: tl.constexpr, + STREAMS: tl.constexpr, + HC_EPS: tl.constexpr, + POST_MULTIPLIER: tl.constexpr, + SINKHORN_ITERS: tl.constexpr, +): + token = tl.program_id(0) + stream_offsets = tl.arange(0, STREAMS) + matrix_offsets = tl.arange(0, STREAMS * STREAMS) + + pre_raw = tl.load(mixes + token * mix_stride_m + stream_offsets) + post_raw = tl.load(mixes + token * mix_stride_m + STREAMS + stream_offsets) + pre_values = tl.sigmoid(pre_raw * tl.load(scale) + tl.load(base + stream_offsets)) + HC_EPS + post_values = POST_MULTIPLIER * tl.sigmoid(post_raw * tl.load(scale + 1) + tl.load(base + STREAMS + stream_offsets)) + + logits = tl.load(mixes + token * mix_stride_m + 2 * STREAMS + matrix_offsets) + logits = logits * tl.load(scale + 2) + tl.load(base + 2 * STREAMS + matrix_offsets) + logits = tl.reshape(logits, (STREAMS, STREAMS)) + logits = logits - tl.max(logits, axis=1)[:, None] + matrix = tl.exp(logits) + matrix = matrix / tl.sum(matrix, axis=1)[:, None] + matrix += HC_EPS + + # The checkpoint definition starts with a column normalization, then + # alternates row and column normalizations for the remaining iterations. + matrix = matrix / (tl.sum(matrix, axis=0)[None, :] + HC_EPS) + for _ in tl.static_range(1, SINKHORN_ITERS): + matrix = matrix / (tl.sum(matrix, axis=1)[:, None] + HC_EPS) + matrix = matrix / (tl.sum(matrix, axis=0)[None, :] + HC_EPS) + + tl.store(pre + token * pre_stride_m + stream_offsets, pre_values) + tl.store(post + token * pre_stride_m + stream_offsets, post_values) + tl.store( + residual_mix + token * residual_stride_m + matrix_offsets, + tl.reshape(matrix, (STREAMS * STREAMS,)), + ) + + +@triton.jit +def _hc_prepare_prenorm_kernel( + gemm_partial, + sqrsum_partial, + scale, + base, + pre, + post, + residual_mix, + gemm_stride_s: tl.constexpr, + gemm_stride_m: tl.constexpr, + sqrsum_stride_s: tl.constexpr, + sqrsum_stride_m: tl.constexpr, + pre_stride_m: tl.constexpr, + residual_stride_m: tl.constexpr, + FLATTENED_HIDDEN: tl.constexpr, + RMS_EPS: tl.constexpr, + STREAMS: tl.constexpr, + HC_EPS: tl.constexpr, + POST_MULTIPLIER: tl.constexpr, + SINKHORN_ITERS: tl.constexpr, + N_SPLITS: tl.constexpr, +): + token = tl.program_id(0) + stream_offsets = tl.arange(0, STREAMS) + matrix_offsets = tl.arange(0, STREAMS * STREAMS) + pre_raw = tl.zeros((STREAMS,), dtype=tl.float32) + post_raw = tl.zeros((STREAMS,), dtype=tl.float32) + matrix_raw = tl.zeros((STREAMS * STREAMS,), dtype=tl.float32) + sqrsum = 0.0 + for split in tl.static_range(N_SPLITS): + partial_base = gemm_partial + split * gemm_stride_s + token * gemm_stride_m + pre_raw += tl.load(partial_base + stream_offsets) + post_raw += tl.load(partial_base + STREAMS + stream_offsets) + matrix_raw += tl.load(partial_base + 2 * STREAMS + matrix_offsets) + sqrsum += tl.load(sqrsum_partial + split * sqrsum_stride_s + token * sqrsum_stride_m) + inv_rms = tl.rsqrt(sqrsum / FLATTENED_HIDDEN + RMS_EPS) + pre_raw *= inv_rms + post_raw *= inv_rms + matrix_raw *= inv_rms + + pre_values = tl.sigmoid(pre_raw * tl.load(scale) + tl.load(base + stream_offsets)) + HC_EPS + post_values = POST_MULTIPLIER * tl.sigmoid(post_raw * tl.load(scale + 1) + tl.load(base + STREAMS + stream_offsets)) + + logits = matrix_raw * tl.load(scale + 2) + tl.load(base + 2 * STREAMS + matrix_offsets) + logits = tl.reshape(logits, (STREAMS, STREAMS)) + logits = logits - tl.max(logits, axis=1)[:, None] + matrix = tl.exp(logits) + matrix = matrix / tl.sum(matrix, axis=1)[:, None] + matrix += HC_EPS + matrix = matrix / (tl.sum(matrix, axis=0)[None, :] + HC_EPS) + for _ in tl.static_range(1, SINKHORN_ITERS): + matrix = matrix / (tl.sum(matrix, axis=1)[:, None] + HC_EPS) + matrix = matrix / (tl.sum(matrix, axis=0)[None, :] + HC_EPS) + + tl.store(pre + token * pre_stride_m + stream_offsets, pre_values) + tl.store(post + token * pre_stride_m + stream_offsets, post_values) + tl.store( + residual_mix + token * residual_stride_m + matrix_offsets, + tl.reshape(matrix, (STREAMS * STREAMS,)), + ) + + +@triton.jit +def _hc_pre_combine_kernel( + x, + pre, + output, + hidden: tl.constexpr, + x_stride_m: tl.constexpr, + pre_stride_m: tl.constexpr, + out_stride_m: tl.constexpr, + STREAMS: tl.constexpr, + BLOCK_H: tl.constexpr, +): + token = tl.program_id(0) + hidden_block = tl.program_id(1) + hidden_offsets = hidden_block * BLOCK_H + tl.arange(0, BLOCK_H) + hidden_mask = hidden_offsets < hidden + accumulator = tl.zeros((BLOCK_H,), dtype=tl.float32) + for stream in tl.static_range(STREAMS): + residual = tl.load( + x + token * x_stride_m + stream * hidden + hidden_offsets, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + pre_value = tl.load(pre + token * pre_stride_m + stream) + accumulator += residual * pre_value + tl.store( + output + token * out_stride_m + hidden_offsets, + accumulator, + mask=hidden_mask, + ) + + +@triton.jit +def _hc_pre_combine_norm_kernel( + x, + pre, + norm_weight, + output, + hidden: tl.constexpr, + x_stride_m: tl.constexpr, + pre_stride_m: tl.constexpr, + out_stride_m: tl.constexpr, + STREAMS: tl.constexpr, + NORM_EPS: tl.constexpr, + BLOCK_H: tl.constexpr, +): + token = tl.program_id(0) + hidden_offsets = tl.arange(0, BLOCK_H) + hidden_mask = hidden_offsets < hidden + accumulator = tl.zeros((BLOCK_H,), dtype=tl.float32) + for stream in tl.static_range(STREAMS): + residual = tl.load( + x + token * x_stride_m + stream * hidden + hidden_offsets, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + pre_value = tl.load(pre + token * pre_stride_m + stream) + accumulator += residual * pre_value + + # hc_pre returns bf16 before the following RMSNorm in the checkpoint + # definition. Preserve that rounding point while keeping both operations + # in one kernel. + rounded = accumulator.to(tl.bfloat16).to(tl.float32) + variance = tl.sum(rounded * rounded, axis=0) / hidden + inv_rms = tl.rsqrt(variance + NORM_EPS) + weight = tl.load(norm_weight + hidden_offsets, mask=hidden_mask, other=0.0) + tl.store( + output + token * out_stride_m + hidden_offsets, + rounded * inv_rms * weight, + mask=hidden_mask, + ) + + +@triton.jit +def _hc_post_4stream_kernel( + layer_output, + residual, + residual_mix, + post_mix, + output, + hidden: tl.constexpr, + layer_stride_m: tl.constexpr, + residual_stride_m: tl.constexpr, + mix_stride_m: tl.constexpr, + post_stride_m: tl.constexpr, + out_stride_m: tl.constexpr, + BLOCK_H: tl.constexpr, +): + token = tl.program_id(0) + hidden_block = tl.program_id(1) + hidden_offsets = hidden_block * BLOCK_H + tl.arange(0, BLOCK_H) + hidden_mask = hidden_offsets < hidden + + # GLM-5 always uses four mHC streams. Compute all four outputs in one + # program so the layer output and residual streams are read only once. + # The previous output-stream grid reread each residual stream four times; + # that becomes bandwidth-bound for large prefills. + layer_value = tl.load( + layer_output + token * layer_stride_m + hidden_offsets, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + residual_base = residual + token * residual_stride_m + hidden_offsets + residual_0 = tl.load(residual_base, mask=hidden_mask, other=0.0).to(tl.float32) + residual_1 = tl.load(residual_base + hidden, mask=hidden_mask, other=0.0).to(tl.float32) + residual_2 = tl.load(residual_base + 2 * hidden, mask=hidden_mask, other=0.0).to(tl.float32) + residual_3 = tl.load(residual_base + 3 * hidden, mask=hidden_mask, other=0.0).to(tl.float32) + + post_base = post_mix + token * post_stride_m + mix_base = residual_mix + token * mix_stride_m + accumulator_0 = layer_value * tl.load(post_base) + accumulator_1 = layer_value * tl.load(post_base + 1) + accumulator_2 = layer_value * tl.load(post_base + 2) + accumulator_3 = layer_value * tl.load(post_base + 3) + + accumulator_0 += residual_0 * tl.load(mix_base) + accumulator_0 += residual_1 * tl.load(mix_base + 4) + accumulator_0 += residual_2 * tl.load(mix_base + 8) + accumulator_0 += residual_3 * tl.load(mix_base + 12) + accumulator_1 += residual_0 * tl.load(mix_base + 1) + accumulator_1 += residual_1 * tl.load(mix_base + 5) + accumulator_1 += residual_2 * tl.load(mix_base + 9) + accumulator_1 += residual_3 * tl.load(mix_base + 13) + accumulator_2 += residual_0 * tl.load(mix_base + 2) + accumulator_2 += residual_1 * tl.load(mix_base + 6) + accumulator_2 += residual_2 * tl.load(mix_base + 10) + accumulator_2 += residual_3 * tl.load(mix_base + 14) + accumulator_3 += residual_0 * tl.load(mix_base + 3) + accumulator_3 += residual_1 * tl.load(mix_base + 7) + accumulator_3 += residual_2 * tl.load(mix_base + 11) + accumulator_3 += residual_3 * tl.load(mix_base + 15) + + output_base = output + token * out_stride_m + hidden_offsets + tl.store(output_base, accumulator_0, mask=hidden_mask) + tl.store(output_base + hidden, accumulator_1, mask=hidden_mask) + tl.store(output_base + 2 * hidden, accumulator_2, mask=hidden_mask) + tl.store(output_base + 3 * hidden, accumulator_3, mask=hidden_mask) + + +def hc_expand(x: torch.Tensor, streams: int) -> torch.Tensor: + """Expand ``[tokens, hidden]`` into flattened residual streams.""" + + assert x.ndim == 2 + return x.unsqueeze(1).expand(-1, streams, -1).reshape(x.shape[0], -1) + + +def hc_contract(x: torch.Tensor, streams: int) -> torch.Tensor: + """Contract flattened residual streams by taking their mean.""" + + assert x.ndim == 2 and x.shape[-1] % streams == 0 + return x.view(x.shape[0], streams, -1).mean(dim=1) + + +def hc_pre_reference( + x: torch.Tensor, + fn: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + streams: int, + rms_eps: float, + hc_eps: float, + sinkhorn_iters: int, + post_multiplier: float = 2.0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute mHC pre-mixes. + + Returns ``(layer_input, residual_mix, post_mix)``. ``x`` and + ``layer_input`` use the activation dtype; both mix tensors are fp32. + """ + + assert x.ndim == 2 and x.shape[-1] % streams == 0 + tokens, flattened_hidden = x.shape + hidden = flattened_hidden // streams + residual = x.view(tokens, streams, hidden) + + x_fp32 = x.float() + inv_rms = torch.rsqrt(x_fp32.square().mean(dim=-1, keepdim=True) + rms_eps) + mixes = F.linear(x_fp32, fn) * inv_rms + + pre_raw = mixes[:, :streams] + post_raw = mixes[:, streams : 2 * streams] + residual_raw = mixes[:, 2 * streams :].view(tokens, streams, streams) + + pre = torch.sigmoid(pre_raw * scale[0] + base[:streams]) + hc_eps + post = post_multiplier * torch.sigmoid(post_raw * scale[1] + base[streams : 2 * streams]) + residual_mix = (residual_raw * scale[2] + base[2 * streams :].view(streams, streams)).softmax(dim=-1) + residual_mix = residual_mix + hc_eps + residual_mix = residual_mix / (residual_mix.sum(dim=-2, keepdim=True) + hc_eps) + for _ in range(sinkhorn_iters - 1): + residual_mix = residual_mix / (residual_mix.sum(dim=-1, keepdim=True) + hc_eps) + residual_mix = residual_mix / (residual_mix.sum(dim=-2, keepdim=True) + hc_eps) + + layer_input = (pre.unsqueeze(-1) * residual.float()).sum(dim=1).to(x.dtype) + return layer_input, residual_mix, post + + +def hc_post_reference( + layer_output: torch.Tensor, + residual: torch.Tensor, + residual_mix: torch.Tensor, + post_mix: torch.Tensor, + streams: int, +) -> torch.Tensor: + """Mix a sublayer output back into the flattened residual streams.""" + + tokens, hidden = layer_output.shape + residual_3d = residual.view(tokens, streams, hidden) + mixed_residual = (residual_mix.unsqueeze(-1) * residual_3d.float().unsqueeze(2)).sum(dim=1) + out = post_mix.unsqueeze(-1) * layer_output.float().unsqueeze(1) + mixed_residual + return out.to(layer_output.dtype).reshape(tokens, streams * hidden) + + +def hc_pre( + x: torch.Tensor, + fn: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + streams: int, + rms_eps: float, + hc_eps: float, + sinkhorn_iters: int, + post_multiplier: float = 2.0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute mHC pre-mixes with fused Sinkhorn and residual combining.""" + + assert x.ndim == 2 and x.shape[-1] % streams == 0 + assert streams == 4, "the fused GLM-5 mHC kernel is specialized for four streams" + assert x.is_contiguous() and fn.is_contiguous() + tokens, flattened_hidden = x.shape + hidden = flattened_hidden // streams + + x_fp32 = x.float() + inv_rms = torch.rsqrt(x_fp32.square().mean(dim=-1, keepdim=True) + rms_eps) + mixes = F.linear(x_fp32, fn) * inv_rms + + pre = torch.empty((tokens, streams), dtype=torch.float32, device=x.device) + post = torch.empty_like(pre) + residual_mix = torch.empty((tokens, streams, streams), dtype=torch.float32, device=x.device) + _hc_prepare_kernel[(tokens,)]( + mixes, + scale, + base, + pre, + post, + residual_mix, + mixes.stride(0), + pre.stride(0), + residual_mix.stride(0), + STREAMS=streams, + HC_EPS=hc_eps, + POST_MULTIPLIER=post_multiplier, + SINKHORN_ITERS=sinkhorn_iters, + num_warps=1, + ) + + layer_input = torch.empty((tokens, hidden), dtype=x.dtype, device=x.device) + block_h = min(triton.next_power_of_2(hidden), 1024) + _hc_pre_combine_kernel[(tokens, triton.cdiv(hidden, block_h))]( + x, + pre, + layer_input, + hidden=hidden, + x_stride_m=x.stride(0), + pre_stride_m=pre.stride(0), + out_stride_m=layer_input.stride(0), + STREAMS=streams, + BLOCK_H=block_h, + num_warps=8, + ) + return layer_input, residual_mix, post + + +def _compute_prenorm_splits(tokens: int, flattened_hidden: int, device: torch.device) -> int: + grid_size = triton.cdiv(tokens, 64) + k_blocks = triton.cdiv(flattened_hidden, 64) + sms = torch.cuda.get_device_properties(device).multi_processor_count + return max(1, min(sms // max(grid_size, 1), k_blocks // 4)) + + +def hc_pre_norm( + x: torch.Tensor, + fn: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + norm_weight: torch.Tensor, + streams: int, + rms_eps: float, + norm_eps: float, + hc_eps: float, + sinkhorn_iters: int, + post_multiplier: float = 2.0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Fuse mHC pre-mixing with its immediately following RMSNorm. + + DeepGEMM computes the fp32 projection and residual square sum in one + split-K kernel. Triton then reduces those partials, runs Sinkhorn, forms + the bf16 residual collapse, and applies RMSNorm. + """ + + assert x.ndim == 2 and x.shape[-1] % streams == 0 + assert streams == 4, "the fused GLM-5 mHC kernel is specialized for four streams" + assert x.dtype == torch.bfloat16 and fn.dtype == torch.float32 + assert x.is_contiguous() and fn.is_contiguous() and norm_weight.is_contiguous() + tokens, flattened_hidden = x.shape + hidden = flattened_hidden // streams + + try: + import deep_gemm + + prenorm_gemm = deep_gemm.tf32_hc_prenorm_gemm + except (AttributeError, ImportError): + from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import ( + rmsnorm_forward, + ) + + layer_input, residual_mix, post = hc_pre( + x, + fn, + scale, + base, + streams, + rms_eps, + hc_eps, + sinkhorn_iters, + post_multiplier, + ) + layer_input = rmsnorm_forward(layer_input, weight=norm_weight, eps=norm_eps) + return layer_input, residual_mix, post + + mix_size = (2 + streams) * streams + n_splits = _compute_prenorm_splits(tokens, flattened_hidden, x.device) + gemm_partial = torch.empty((n_splits, tokens, mix_size), dtype=torch.float32, device=x.device) + sqrsum_partial = torch.empty((n_splits, tokens), dtype=torch.float32, device=x.device) + prenorm_gemm(x, fn, gemm_partial, sqrsum_partial, n_splits) + + pre = torch.empty((tokens, streams), dtype=torch.float32, device=x.device) + post = torch.empty_like(pre) + residual_mix = torch.empty((tokens, streams, streams), dtype=torch.float32, device=x.device) + _hc_prepare_prenorm_kernel[(tokens,)]( + gemm_partial, + sqrsum_partial, + scale, + base, + pre, + post, + residual_mix, + gemm_partial.stride(0), + gemm_partial.stride(1), + sqrsum_partial.stride(0), + sqrsum_partial.stride(1), + pre.stride(0), + residual_mix.stride(0), + FLATTENED_HIDDEN=flattened_hidden, + RMS_EPS=rms_eps, + STREAMS=streams, + HC_EPS=hc_eps, + POST_MULTIPLIER=post_multiplier, + SINKHORN_ITERS=sinkhorn_iters, + N_SPLITS=n_splits, + num_warps=1, + ) + + layer_input = torch.empty((tokens, hidden), dtype=x.dtype, device=x.device) + block_h = triton.next_power_of_2(hidden) + _hc_pre_combine_norm_kernel[(tokens,)]( + x, + pre, + norm_weight, + layer_input, + hidden=hidden, + x_stride_m=x.stride(0), + pre_stride_m=pre.stride(0), + out_stride_m=layer_input.stride(0), + STREAMS=streams, + NORM_EPS=norm_eps, + BLOCK_H=block_h, + num_warps=8, + ) + return layer_input, residual_mix, post + + +def hc_post( + layer_output: torch.Tensor, + residual: torch.Tensor, + residual_mix: torch.Tensor, + post_mix: torch.Tensor, + streams: int, +) -> torch.Tensor: + """Mix a sublayer output into residual streams with one Triton launch.""" + + tokens, hidden = layer_output.shape + assert streams == 4, "the fused GLM-5 mHC kernel is specialized for four streams" + assert layer_output.is_contiguous() and residual.is_contiguous() + output = torch.empty( + (tokens, streams * hidden), + dtype=layer_output.dtype, + device=layer_output.device, + ) + block_h = min(triton.next_power_of_2(hidden), 1024) + _hc_post_4stream_kernel[(tokens, triton.cdiv(hidden, block_h))]( + layer_output, + residual, + residual_mix, + post_mix, + output, + hidden=hidden, + layer_stride_m=layer_output.stride(0), + residual_stride_m=residual.stride(0), + mix_stride_m=residual_mix.stride(0), + post_stride_m=post_mix.stride(0), + out_stride_m=output.stride(0), + BLOCK_H=block_h, + num_warps=8, + ) + return output diff --git a/lightllm/models/glm5_next_mtp/__init__.py b/lightllm/models/glm5_next_mtp/__init__.py new file mode 100644 index 0000000000..d55bc8b398 --- /dev/null +++ b/lightllm/models/glm5_next_mtp/__init__.py @@ -0,0 +1,3 @@ +from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel + +__all__ = ["Glm5NextMTPModel"] diff --git a/lightllm/models/glm5_next_mtp/layer_infer/__init__.py b/lightllm/models/glm5_next_mtp/layer_infer/__init__.py new file mode 100644 index 0000000000..0ec4b5a01b --- /dev/null +++ b/lightllm/models/glm5_next_mtp/layer_infer/__init__.py @@ -0,0 +1,5 @@ +from lightllm.models.glm5_next_mtp.layer_infer.pre_layer_infer import ( + Glm5NextMTPPreLayerInfer, +) + +__all__ = ["Glm5NextMTPPreLayerInfer"] diff --git a/lightllm/models/glm5_next_mtp/layer_infer/post_layer_infer.py b/lightllm/models/glm5_next_mtp/layer_infer/post_layer_infer.py new file mode 100644 index 0000000000..5708ce30ce --- /dev/null +++ b/lightllm/models/glm5_next_mtp/layer_infer/post_layer_infer.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 + +from typing import Callable + +import torch + +from lightllm.distributed.communication_op import all_gather_into_tensor +from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer + + +def vocab_parallel_top1( + local_logits: torch.Tensor, + local_vocab_start_id: int, + tp_world_size: int, + dist_group, + alloc_func: Callable, +) -> torch.Tensor: + """Return global greedy token ids without materializing global logits.""" + + assert local_logits.ndim == 2 + local_max_values, local_max_indexes = torch.max(local_logits, dim=0) + local_token_ids = local_max_indexes + local_vocab_start_id + if tp_world_size == 1: + return local_token_ids + + local_winners = torch.stack( + [local_max_values.float(), local_token_ids.float()], + dim=-1, + ).contiguous() + token_num = local_winners.shape[0] + gathered_winners = alloc_func( + (tp_world_size * token_num, 2), + dtype=torch.float32, + device=local_logits.device, + ) + all_gather_into_tensor( + gathered_winners, + local_winners, + group=dist_group, + async_op=False, + ) + gathered_winners = gathered_winners.view(tp_world_size, token_num, 2) + winning_ranks = torch.argmax(gathered_winners[:, :, 0], dim=0) + token_rows = torch.arange(token_num, dtype=torch.long, device=local_logits.device) + return gathered_winners[winning_ranks, token_rows, 1].long() + + +def vocab_parallel_top1_and_prob( + local_logits: torch.Tensor, + local_vocab_start_id: int, + tp_world_size: int, + dist_group, + alloc_func: Callable, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return global Top-1 ids and exact softmax probabilities. + + Only three FP32 values per token and rank are gathered: the local maximum, + its global token id, and the local log-sum-exp. This replaces the normal + all-gather of every vocabulary logit while preserving the result of a + global FP32 argmax/softmax reduction. + """ + + assert local_logits.ndim == 2 + local_logits_fp32 = local_logits.float() + local_max_values, local_max_indexes = torch.max(local_logits_fp32, dim=0) + local_logsumexp = torch.logsumexp(local_logits_fp32, dim=0) + local_token_ids = local_max_indexes + local_vocab_start_id + + if tp_world_size == 1: + return local_token_ids, torch.exp(local_max_values - local_logsumexp) + + local_stats = torch.stack( + [local_max_values, local_token_ids.float(), local_logsumexp], + dim=-1, + ).contiguous() + token_num = local_stats.shape[0] + gathered_stats = alloc_func( + (tp_world_size * token_num, 3), + dtype=torch.float32, + device=local_logits.device, + ) + all_gather_into_tensor( + gathered_stats, + local_stats, + group=dist_group, + async_op=False, + ) + gathered_stats = gathered_stats.view(tp_world_size, token_num, 3) + winning_ranks = torch.argmax(gathered_stats[:, :, 0], dim=0) + token_rows = torch.arange(token_num, dtype=torch.long, device=local_logits.device) + winning_stats = gathered_stats[winning_ranks, token_rows] + global_logsumexp = torch.logsumexp(gathered_stats[:, :, 2], dim=0) + return winning_stats[:, 1].long(), torch.exp(winning_stats[:, 0] - global_logsumexp) + + +class Glm5NextMTPPostLayerInfer(LlamaPostLayerInfer): + """GLM NextN uses the common exact vocab-parallel greedy output path.""" diff --git a/lightllm/models/glm5_next_mtp/layer_infer/pre_layer_infer.py b/lightllm/models/glm5_next_mtp/layer_infer/pre_layer_infer.py new file mode 100644 index 0000000000..6fc7826a75 --- /dev/null +++ b/lightllm/models/glm5_next_mtp/layer_infer/pre_layer_infer.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: Apache-2.0 + +from lightllm.models.deepseek_mtp.layer_infer.pre_layer_infer import ( + Deepseek3MTPPreLayerInfer, +) +from lightllm.models.glm5_next_mtp.triton_kernel.zero_position_embedding import ( + zero_position_embedding_, +) + + +class Glm5NextMTPPreLayerInfer(Deepseek3MTPPreLayerInfer): + """GLM NextN input fusion with the trained position-zero convention.""" + + def _mtp_context_forward(self, input_embdings, infer_state, layer_weight): + # GLM's fused_eh_norm reference zeros the token embedding at absolute + # position zero before applying enorm. The first target hidden remains + # intact; only the missing previous-token embedding is suppressed. + zero_position_embedding_(input_embdings, infer_state.position_ids) + return super()._mtp_context_forward(input_embdings, infer_state, layer_weight) diff --git a/lightllm/models/glm5_next_mtp/layer_weights/__init__.py b/lightllm/models/glm5_next_mtp/layer_weights/__init__.py new file mode 100644 index 0000000000..8507c67bd8 --- /dev/null +++ b/lightllm/models/glm5_next_mtp/layer_weights/__init__.py @@ -0,0 +1,5 @@ +from lightllm.models.glm5_next_mtp.layer_weights.pre_and_post_layer_weight import ( + Glm5NextMTPPreAndPostLayerWeight, +) + +__all__ = ["Glm5NextMTPPreAndPostLayerWeight"] diff --git a/lightllm/models/glm5_next_mtp/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/glm5_next_mtp/layer_weights/pre_and_post_layer_weight.py new file mode 100644 index 0000000000..f697114829 --- /dev/null +++ b/lightllm/models/glm5_next_mtp/layer_weights/pre_and_post_layer_weight.py @@ -0,0 +1,49 @@ +from lightllm.common.basemodel import PreAndPostLayerWeight +from lightllm.common.basemodel.layer_weights.meta_weights import ( + EmbeddingWeight, + LMHeadWeight, + RMSNormWeight, + ROWMMWeight, +) +from lightllm.common.quantization import Quantcfg +from lightllm.models.glm5_next.layer_weights.pre_and_post_layer_weight import ( + add_language_model_aliases, +) + + +class Glm5NextMTPPreAndPostLayerWeight(PreAndPostLayerWeight): + def __init__(self, data_type, network_config, quant_cfg: Quantcfg): + super().__init__(data_type, network_config) + layer_idx = network_config["num_hidden_layers"] + hidden_size = network_config["hidden_size"] + prefix = f"model.layers.{layer_idx}" + self.eh_proj_weight_ = ROWMMWeight( + in_dim=hidden_size * 2, + out_dims=[hidden_size], + weight_names=f"{prefix}.eh_proj.weight", + data_type=self.data_type_, + quant_method=None, + tp_rank=0, + tp_world_size=1, + ) + self.enorm_weight_ = RMSNormWeight( + dim=hidden_size, + weight_name=f"{prefix}.enorm.weight", + data_type=self.data_type_, + ) + self.hnorm_weight_ = RMSNormWeight( + dim=hidden_size, + weight_name=f"{prefix}.hnorm.weight", + data_type=self.data_type_, + ) + self.final_norm_weight_ = RMSNormWeight( + dim=hidden_size, + weight_name=f"{prefix}.shared_head.norm.weight", + data_type=self.data_type_, + ) + self.wte_weight_: EmbeddingWeight = None + self.lm_head_weight_: LMHeadWeight = None + + def load_hf_weights(self, weights): + add_language_model_aliases(weights) + return super().load_hf_weights(weights) diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py new file mode 100644 index 0000000000..7c10ade411 --- /dev/null +++ b/lightllm/models/glm5_next_mtp/model.py @@ -0,0 +1,84 @@ +from typing import List + +from lightllm.common.basemodel import TpPartBaseModel +from lightllm.models.glm5_next_mtp.layer_infer.pre_layer_infer import ( + Glm5NextMTPPreLayerInfer, +) +from lightllm.models.glm5_next_mtp.layer_infer.post_layer_infer import ( + Glm5NextMTPPostLayerInfer, +) +from lightllm.models.draft_registry import DraftModelRegistry +from lightllm.models.glm5_next_mtp.layer_weights.pre_and_post_layer_weight import ( + Glm5NextMTPPreAndPostLayerWeight, +) +from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import ( + Glm5NextTransformerLayerInfer, +) +from lightllm.models.glm5_next.layer_weights.transformer_layer_weight import ( + Glm5NextTransformerLayerWeight, +) +from lightllm.models.glm5_next.model import Glm5NextTpPartModel + + +@DraftModelRegistry( + model_type=("glm5_next", "glm5_next_text"), + spec_modes=("vanilla_with_att", "eagle_with_att"), +) +class Glm5NextMTPModel(Glm5NextTpPartModel): + """GLM-5.3's shared one-layer NextN draft model.""" + + is_mtp_draft_model = True + replicated_attention_ep = False + pre_and_post_weight_class = Glm5NextMTPPreAndPostLayerWeight + pre_layer_infer_class = Glm5NextMTPPreLayerInfer + post_layer_infer_class = Glm5NextMTPPostLayerInfer + transformer_weight_class = Glm5NextTransformerLayerWeight + transformer_layer_infer_class = Glm5NextTransformerLayerInfer + + def __init__(self, kvargs: dict): + self.main_model: TpPartBaseModel = kvargs.pop("main_model") + self.mtp_previous_draft_models: List[TpPartBaseModel] = kvargs.pop("mtp_previous_draft_models") + super().__init__(kvargs) + + def _init_custom(self): + self._cos_cached = self.main_model._cos_cached + self._sin_cached = self.main_model._sin_cached + + def _init_req_manager(self): + self.req_manager = self.main_model.req_manager + + def _init_mem_manager(self): + self.mem_manager = self.main_model.mem_manager + + def _init_weights(self, start_layer_index=None): + assert start_layer_index is None + mtp_layer = self.config["num_hidden_layers"] + self.pre_post_weight = self.pre_and_post_weight_class( + self.data_type, network_config=self.config, quant_cfg=self.quant_cfg + ) + self.trans_layers_weight = [ + self.transformer_weight_class( + mtp_layer, + self.data_type, + network_config=self.config, + quant_cfg=self.quant_cfg, + ) + ] + self.pre_post_weight.wte_weight_ = self.main_model.pre_post_weight.wte_weight_ + self.pre_post_weight.lm_head_weight_ = self.main_model.pre_post_weight.lm_head_weight_ + + def _init_infer_layer(self, start_layer_index=None): + assert start_layer_index is None + self.pre_infer = self.pre_layer_infer_class(network_config=self.config) + self.post_infer = self.post_layer_infer_class(network_config=self.config) + logical_layer = len(self.main_model.layers_infer) + sum( + len(model.layers_infer) for model in self.mtp_previous_draft_models + ) + self.layers_infer = [self.transformer_layer_infer_class(logical_layer, network_config=self.config)] + + def _init_some_value(self): + super()._init_some_value() + self.layers_num = 1 + + def autotune_layers(self): + return 1 diff --git a/lightllm/models/glm5_next_mtp/triton_kernel/__init__.py b/lightllm/models/glm5_next_mtp/triton_kernel/__init__.py new file mode 100644 index 0000000000..daeefbafca --- /dev/null +++ b/lightllm/models/glm5_next_mtp/triton_kernel/__init__.py @@ -0,0 +1,5 @@ +from lightllm.models.glm5_next_mtp.triton_kernel.zero_position_embedding import ( + zero_position_embedding_, +) + +__all__ = ["zero_position_embedding_"] diff --git a/lightllm/models/glm5_next_mtp/triton_kernel/zero_position_embedding.py b/lightllm/models/glm5_next_mtp/triton_kernel/zero_position_embedding.py new file mode 100644 index 0000000000..ea0e48dbe7 --- /dev/null +++ b/lightllm/models/glm5_next_mtp/triton_kernel/zero_position_embedding.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 + +import torch + +import triton +import triton.language as tl + + +@triton.jit +def _zero_position_embedding_kernel( + embeddings, + stride_embeddings_m, + stride_embeddings_n, + position_ids, + hidden_size: tl.constexpr, + BLOCK_N: tl.constexpr, +): + row = tl.program_id(0) + if tl.load(position_ids + row) != 0: + return + + offsets = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) + tl.store( + embeddings + row * stride_embeddings_m + offsets * stride_embeddings_n, + 0.0, + mask=offsets < hidden_size, + ) + + +@torch.no_grad() +def zero_position_embedding_(embeddings: torch.Tensor, position_ids: torch.Tensor) -> None: + """Zero MTP token embeddings at absolute position zero in place.""" + + assert embeddings.is_cuda and position_ids.is_cuda + assert embeddings.ndim == 2 and position_ids.shape == embeddings.shape[:1] + block_n = 256 + grid = (embeddings.shape[0], triton.cdiv(embeddings.shape[1], block_n)) + _zero_position_embedding_kernel[grid]( + embeddings=embeddings, + stride_embeddings_m=embeddings.stride(0), + stride_embeddings_n=embeddings.stride(1), + position_ids=position_ids, + hidden_size=embeddings.shape[1], + BLOCK_N=block_n, + num_warps=4, + num_stages=1, + ) diff --git a/lightllm/models/llama/layer_infer/post_layer_infer.py b/lightllm/models/llama/layer_infer/post_layer_infer.py index 6e4b15a55d..701dbd00ed 100644 --- a/lightllm/models/llama/layer_infer/post_layer_infer.py +++ b/lightllm/models/llama/layer_infer/post_layer_infer.py @@ -7,6 +7,9 @@ from lightllm.models.llama.layer_weights.pre_and_post_layer_weight import LlamaPreAndPostLayerWeight from lightllm.models.llama.infer_struct import LlamaInferStateInfo from lightllm.common.basemodel import PostLayerInferTpl +from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_greedy import ( + vocab_parallel_greedy, +) from lightllm.distributed.communication_op import all_gather @@ -21,6 +24,10 @@ def __init__(self, network_config): def _norm(self, input, infer_state, layer_weight: LlamaPreAndPostLayerWeight) -> torch.Tensor: return layer_weight.final_norm_weight_(input=input, eps=self.eps_, alloc_func=self.alloc_tensor) + def _apply_logit_postprocessing(self, logits: torch.Tensor) -> torch.Tensor: + """Apply model-specific elementwise transforms before greedy reduction.""" + return logits + def _slice_get_last_input(self, input_embdings: torch.Tensor, infer_state: LlamaInferStateInfo): embed_dim_ = input_embdings.shape[1] if infer_state.is_prefill: @@ -64,7 +71,7 @@ def _token_forward( if prompt_logics_hiddens is not None: prompt_token_num = prompt_logics_hiddens.shape[0] infer_state.prompt_logics = self._lm_head_and_gather( - prompt_logics_hiddens, prompt_token_num, layer_weight, infer_state + prompt_logics_hiddens, prompt_token_num, layer_weight, infer_state, force_full_logits=True ) return ans_logics @@ -75,6 +82,7 @@ def _lm_head_and_gather( token_num: int, layer_weight: LlamaPreAndPostLayerWeight, infer_state: LlamaInferStateInfo, + force_full_logits: bool = False, ) -> torch.Tensor: normed = self._norm(hidden, infer_state, layer_weight) normed = normed.permute(1, 0).view(-1, token_num) @@ -82,6 +90,19 @@ def _lm_head_and_gather( normed = None vocab_size = layer_weight.lm_head_weight_.vocab_size + if infer_state.use_vocab_parallel_greedy and not force_full_logits: + logic_batch = self._apply_logit_postprocessing(logic_batch) + logits, token_ids, logsumexp = vocab_parallel_greedy( + logic_batch, + vocab_size=vocab_size, + tp_world_size=self.tp_world_size_, + group=infer_state.dist_group, + alloc_func=self.alloc_tensor, + ) + infer_state.logits_token_ids = token_ids + infer_state.logits_logsumexp = logsumexp + return logits + if self.tp_world_size_ == 1: gather_data = logic_batch else: @@ -98,7 +119,7 @@ def _lm_head_and_gather( ans_logics = self.alloc_tensor((token_num, vocab_size), dtype=torch.float32) ans_logics[:, :] = gather_data.permute(1, 0) gather_data = None - return ans_logics + return self._apply_logit_postprocessing(ans_logics) def token_forward( self, input_embdings: torch.Tensor, infer_state: LlamaInferStateInfo, layer_weight: LlamaPreAndPostLayerWeight diff --git a/lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py b/lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py index 5a74cd988e..eb4481fd6b 100644 --- a/lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py +++ b/lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py @@ -181,7 +181,11 @@ def token_forward( logits = self._lm_head_and_gather(last_input, token_num, layer_weight, infer_state) block_logits = logits.reshape(num_reqs, self.block_size_, -1) - sampled_tokens = torch.argmax(block_logits, dim=-1) + if infer_state.logits_token_ids is None: + sampled_tokens = torch.argmax(block_logits, dim=-1) + else: + assert block_logits.shape[-1] == 1 + sampled_tokens = infer_state.logits_token_ids.reshape(num_reqs, self.block_size_) confidence_logits = self.predict_confidence_logits( block_hidden, anchor_token_ids=anchor_token_ids, diff --git a/lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py index 7311c4d141..e317f08f2d 100644 --- a/lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py @@ -7,7 +7,7 @@ from lightllm.models.llama.infer_struct import LlamaInferStateInfo from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import ( - use_sm100_mega_moe, + use_mega_moe, ) from lightllm.utils.dist_utils import get_global_world_size from lightllm.utils.envs_utils import get_env_start_args @@ -138,7 +138,7 @@ def overlap_tpsp_token_forward( infer_state1: LlamaInferStateInfo, layer_weight: Qwen3MOETransformerLayerWeight, ): - if not self.is_moe or use_sm100_mega_moe(layer_weight.experts.quant_method): + if not self.is_moe or use_mega_moe(layer_weight.experts.quant_method): return super().overlap_tpsp_token_forward( input_embdings, input_embdings1, infer_state, infer_state1, layer_weight ) @@ -250,7 +250,7 @@ def overlap_tpsp_context_forward( infer_state1: LlamaInferStateInfo, layer_weight: Qwen3MOETransformerLayerWeight, ): - if not self.is_moe or use_sm100_mega_moe(layer_weight.experts.quant_method): + if not self.is_moe or use_mega_moe(layer_weight.experts.quant_method): return super().overlap_tpsp_context_forward( input_embdings, input_embdings1, infer_state, infer_state1, layer_weight ) @@ -318,8 +318,8 @@ def overlap_tpsp_context_forward( # 0 moe calu _0_moe_out = layer_weight.experts.prefilled_group_gemm( _0_num_recv_tokens_per_expert_list, - _0_handle.num_unaligned_recv_tokens_per_expert, - _0_handle.recv_src_metadata, + getattr(_0_handle, "num_unaligned_recv_tokens_per_expert", None), + getattr(_0_handle, "recv_src_metadata", None), _0_recv_x, _0_recv_topk_idx, _0_recv_topk_weight, @@ -350,8 +350,8 @@ def overlap_tpsp_context_forward( # 1 moe calc _1_moe_out = layer_weight.experts.prefilled_group_gemm( _1_num_recv_tokens_per_expert_list, - _1_handle.num_unaligned_recv_tokens_per_expert, - _1_handle.recv_src_metadata, + getattr(_1_handle, "num_unaligned_recv_tokens_per_expert", None), + getattr(_1_handle, "recv_src_metadata", None), _1_recv_x, _1_recv_topk_idx, _1_recv_topk_weight, diff --git a/lightllm/models/qwen3_omni_moe_thinker/model.py b/lightllm/models/qwen3_omni_moe_thinker/model.py index e4ea3ccc34..32b959675e 100644 --- a/lightllm/models/qwen3_omni_moe_thinker/model.py +++ b/lightllm/models/qwen3_omni_moe_thinker/model.py @@ -1,6 +1,5 @@ import os import json -import librosa import copy from functools import lru_cache from io import BytesIO diff --git a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py index 92d68c9fd2..2b7b6a3ae4 100644 --- a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py @@ -319,7 +319,6 @@ def _linear_post( ) -> torch.Tensor: num_tokens = z.shape[0] core_attn_out = core_attn_out.view(-1, core_attn_out.shape[-1]) - z = z.contiguous().view(-1, z.shape[-1]) norm_out = layer_weight.linear_norm(core_attn_out, z, self.eps_) core_attn_out = norm_out.view(num_tokens, -1) output = layer_weight.linear_out_proj.mm(core_attn_out) diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index d72e63d724..e5f45f6963 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -233,6 +233,16 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: parser.add_argument( "--running_max_req_size", type=int, default=256, help="the max size for forward requests in the same time" ) + parser.add_argument( + "--per_dp_max_req_size", + type=int, + default=None, + help=( + "Optional request-state capacity allocated by each local DP replica. " + "Defaults to running_max_req_size for backward compatibility; lowering it can " + "substantially reduce hybrid linear-attention state memory in balanced DP workloads." + ), + ) parser.add_argument("--nnodes", type=int, default=1, help="the number of nodes") parser.add_argument("--node_rank", type=int, default=0, help="the rank of the current node") parser.add_argument( @@ -423,11 +433,12 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: "--llm_prefill_att_backend", type=str, nargs="+", - choices=["auto", "triton", "fa3", "flashinfer", "flashqla"], + choices=["auto", "triton", "fa3", "flashinfer", "flashqla", "tilelang"], default=["auto"], help="""prefill attention kernel used in llm. auto: automatically select best backend based on GPU and available packages (priority: fa3 > flashinfer > triton) + for NSA/DSA models, tilelang selects SGLang's sparse prefill kernel for hybrid linear-attention models, the second value selects the linear-attention backend (priority: flashqla > triton); when omitted, it defaults to auto""", ) @@ -600,6 +611,36 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: default=8192, help="max handle token num for prefill cudagraph", ) + parser.add_argument( + "--prefill_cudagraph_token_nums", + nargs="+", + type=int, + default=None, + help=( + "Optional exact prefill token counts to capture. Requests with other token counts " + "run eagerly instead of being padded to one of these graphs." + ), + ) + parser.add_argument( + "--prefill_cudagraph_batch_sizes", + nargs="+", + type=int, + default=None, + help=( + "Batch size paired with each --prefill_cudagraph_token_nums entry. " + "Exact-layout graphs require uniform uncached sequences, so each token count " + "must be divisible by its paired batch size." + ), + ) + parser.add_argument( + "--prefill_cudagraph_capture_attention", + action="store_true", + help=( + "Capture prefill attention inside the main CUDA Graph instead of running it " + "between graph segments. This experimental mode requires exact token-count and " + "batch-size layouts." + ), + ) parser.add_argument( "--graph_max_batch_size", @@ -791,6 +832,13 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: default=0.03, help="""The interval of the schedule time, default is 30ms.""", ) + parser.add_argument( + "--prefill_coalesce_interval", + type=float, + default=0.0, + help="""Maximum time in seconds to collect a burst of waiting requests before scheduling + a prefill batch. Disabled by default. A full runnable batch is scheduled immediately.""", + ) parser.add_argument( "--afs_image_embed_dir", type=str, diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index 37fe837ad1..33e12c2b03 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -1,8 +1,9 @@ import multiprocessing as mp import os -import uuid -import subprocess import math +import subprocess +import sys +import uuid from lightllm.utils.start_utils import process_manager from .metrics.manager import start_metric_manager from .embed_cache.manager import start_cache_manager @@ -36,10 +37,10 @@ def _set_envs_and_config(args: StartArgs): def _launch_subprocesses(args: StartArgs): _set_envs_and_config(args) - if args.mtp_mode is not None: + if args.mtp_mode is not None and args.mtp_dynamic_verify: assert ( not args.disable_cudagraph or args.run_mode == "prefill" - ), "--disable_cudagraph is only supported on Prefill nodes when --mtp_mode is enabled" + ), "--disable_cudagraph is only supported on Prefill nodes when --mtp_dynamic_verify is enabled" auto_set_max_req_total_len(args) auto_set_fused_shared_experts(args) @@ -413,12 +414,19 @@ def _hypercorn_config_args(args: StartArgs): return ["--keep-alive", "10"] +def _hypercorn_entrypoint(): + # Launch from the active Python environment instead of relying on a + # console-script wrapper being present on PATH (for example, packages + # installed with pip --target do not install that wrapper). + return [sys.executable, "-m", "hypercorn"] + + def normal_or_p_d_start(args: StartArgs): process_manager = _launch_subprocesses(args) # 启动 Hypercorn command = [ - "hypercorn", + *_hypercorn_entrypoint(), *_hypercorn_config_args(args), "--workers", f"{args.httpserver_workers}", @@ -481,7 +489,7 @@ def pd_master_start(args: StartArgs): ) command = [ - "hypercorn", + *_hypercorn_entrypoint(), *_hypercorn_config_args(args), "--workers", "1", @@ -569,7 +577,7 @@ def config_server_start(args): start_redis_service(args) command = [ - "hypercorn", + *_hypercorn_entrypoint(), *_hypercorn_config_args(args), "--workers", "1", diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 5ccaa5a401..b898dec028 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -81,6 +81,7 @@ class StartArgs: ) chat_template: Optional[str] = field(default=None) running_max_req_size: int = field(default=256) + per_dp_max_req_size: Optional[int] = field(default=None) tp: int = field(default=1) dp: int = field(default=1) nnodes: int = field(default=1) @@ -150,6 +151,9 @@ class StartArgs: disable_cudagraph: bool = field(default=False) enable_prefill_cudagraph: bool = field(default=False) prefill_cudagraph_max_handle_token: int = field(default=8192) + prefill_cudagraph_token_nums: Optional[List[int]] = field(default=None) + prefill_cudagraph_batch_sizes: Optional[List[int]] = field(default=None) + prefill_cudagraph_capture_attention: bool = field(default=False) graph_max_batch_size: int = field(default=256) graph_split_batch_size: int = field(default=32) graph_grow_step_size: int = field(default=16) @@ -161,7 +165,7 @@ class StartArgs: expert_dtype: Optional[str] = field(default=None, metadata={"choices": ["fp8", "fp4"]}) llm_prefill_att_backend: List[str] = field( default_factory=lambda: ["auto"], - metadata={"choices": ["auto", "triton", "fa3", "flashinfer", "flashqla"]}, + metadata={"choices": ["auto", "triton", "fa3", "flashinfer", "flashqla", "tilelang"]}, ) llm_decode_att_backend: List[str] = field( default_factory=lambda: ["auto"], metadata={"choices": ["auto", "triton", "fa3", "flashinfer"]} @@ -219,6 +223,7 @@ class StartArgs: enable_mps: bool = field(default=False) multinode_router_gloo_port: int = field(default=20001) schedule_time_interval: float = field(default=0.03) + prefill_coalesce_interval: float = field(default=0.0) use_dynamic_prompt_cache: bool = field(default=False) enable_rl: bool = field(default=False) enable_torch_memory_saver: bool = field(default=False) diff --git a/lightllm/server/httpserver_for_pd_master/manager.py b/lightllm/server/httpserver_for_pd_master/manager.py index b422bf7703..db1d923fe3 100644 --- a/lightllm/server/httpserver_for_pd_master/manager.py +++ b/lightllm/server/httpserver_for_pd_master/manager.py @@ -41,13 +41,16 @@ def __init__( self.metric_client = MetricClient(get_shm_port_args().metric_port) self.id_gen = ReqIDGenerator() - self.pd_manager = PDManager(args) + self.pd_manager = PDManager(args, self.metric_client) self.req_id_to_out_inf: Dict[int, ReqStatus] = {} self.infos_queues = None # 这个需要延迟初始化，否则使用的loop不对 self.health_timeout = int(os.getenv("HEALTH_TIMEOUT", "200")) self.latest_success_infer_time = time.time() self.running_request_count = 0 + self.pd_stage_waiting_request_counts = {"prefill": 0, "decode": 0} + for stage in self.pd_stage_waiting_request_counts: + self.metric_client.gauge_set("lightllm_pd_master_stage_waiting_requests", 0, labels={"stage": stage}) self.tokenizer = get_tokenizer(args.model_dir, args.tokenizer_mode, trust_remote_code=args.trust_remote_code) @@ -359,6 +362,28 @@ async def raise_if_disconnected() -> None: except asyncio.TimeoutError: continue + def _change_pd_stage_waiting_requests(self, stage: str, delta: int) -> None: + self.pd_stage_waiting_request_counts[stage] += delta + self.metric_client.gauge_set( + "lightllm_pd_master_stage_waiting_requests", + self.pd_stage_waiting_request_counts[stage], + labels={"stage": stage}, + ) + + async def _wait_for_pd_stage( + self, + event: asyncio.Event, + request: Request, + timeout: float, + group_request_id: int, + stage: str, + ) -> None: + self._change_pd_stage_waiting_requests(stage, 1) + try: + await self._wait_for_event_or_disconnect(event, request, timeout, group_request_id, stage) + finally: + self._change_pd_stage_waiting_requests(stage, -1) + async def _log_req_header(self, request: Request, group_request_id: int): x_request_id = request.headers.get("X-Request-Id", "") x_session_id = request.headers.get("X-Session-Id", "") @@ -393,7 +418,7 @@ async def fetch_pd_stream( await p_node.websocket.send_bytes(pickle.dumps((ObjType.REQ, (prompt, sampling_params, multimodal_params)))) try: - await self._wait_for_event_or_disconnect( + await self._wait_for_pd_stage( prefill_prompt_ids_event, request, timeout=60, @@ -414,7 +439,7 @@ async def fetch_pd_stream( ) try: - await self._wait_for_event_or_disconnect( + await self._wait_for_pd_stage( up_status_event, request, timeout=180, @@ -436,13 +461,19 @@ async def fetch_pd_stream( first_token_gen = False needs_prefill_first_token = decode_node_info.ready_kv_len != len(prompt_ids) - 1 - prompt_cache_len_from_prefill = await self._wait_for_prefill_token_if_needed( - req_status=req_status, - request=request, - group_request_id=group_request_id, - needs_prefill_first_token=needs_prefill_first_token, - ready_kv_len=decode_node_info.ready_kv_len, - ) + if needs_prefill_first_token: + self._change_pd_stage_waiting_requests("prefill", 1) + try: + prompt_cache_len_from_prefill = await self._wait_for_prefill_token_if_needed( + req_status=req_status, + request=request, + group_request_id=group_request_id, + needs_prefill_first_token=needs_prefill_first_token, + ready_kv_len=decode_node_info.ready_kv_len, + ) + finally: + if needs_prefill_first_token: + self._change_pd_stage_waiting_requests("prefill", -1) while True: await req_status.wait_to_ready() @@ -754,8 +785,9 @@ async def put_tokens_to_front(self, token_list: List[Tuple[int, str, dict, Finis class PDManager: - def __init__(self, args: StartArgs): + def __init__(self, args: StartArgs, metric_client=None): self.args: StartArgs = args + self.metric_client = metric_client self.prefill_nodes: List[PD_Client_Obj] = [] self.decode_nodes: List[PD_Client_Obj] = [] self.url_to_pd_nodes: Dict[str, PD_Client_Obj] = {} @@ -879,6 +911,12 @@ def update_node_load_info(self, load_info: Optional[dict]): total_token_usage_rate = load_info["total_token_usage_rate"] pd_client = self.url_to_pd_nodes.get(client_ip_port) pd_client.run_status.total_token_usage_rate = total_token_usage_rate + if self.metric_client is not None: + self.metric_client.gauge_set( + "lightllm_pd_node_token_usage_ratio", + total_token_usage_rate, + labels={"role": pd_client.mode, "endpoint": client_ip_port}, + ) except BaseException as e: logger.warning(f"udpate node load info failed, load_info: {load_info} error: {str(e)}") return diff --git a/lightllm/server/metrics/manager.py b/lightllm/server/metrics/manager.py index 22f6426a77..f8f56dc8d3 100644 --- a/lightllm/server/metrics/manager.py +++ b/lightllm/server/metrics/manager.py @@ -55,8 +55,9 @@ def exposed_counter_inc_by(self, name: str, amount: float) -> None: def exposed_histogram_observe(self, name: str, value: float, label: str = None) -> None: return self.monitor.histogram_observe(name, value, label) - def exposed_gauge_set(self, name: str, value: float) -> None: - return self.monitor.gauge_set(name, value) + def exposed_gauge_set(self, name: str, value: float, labels: dict = None) -> None: + local_labels = None if labels is None else {key: labels[key] for key in labels} + return self.monitor.gauge_set(name, value, local_labels) def exposed_generate_latest(self) -> bytes: data = generate_latest(self.monitor.registry) diff --git a/lightllm/server/metrics/metrics.py b/lightllm/server/metrics/metrics.py index 0d42462c3f..b58b587b26 100644 --- a/lightllm/server/metrics/metrics.py +++ b/lightllm/server/metrics/metrics.py @@ -32,6 +32,8 @@ "lightllm_cache_hit_rate": "Prefix cache hit rate of latest completed request", "lightllm_gen_throughput": "Generation throughput of latest completed request (tokens/s)", "lightllm_num_running_reqs": "Number of running requests", + "lightllm_pd_node_token_usage_ratio": "Token capacity usage ratio reported by a PD node", + "lightllm_pd_master_stage_waiting_requests": "Requests waiting for a PD stage to become ready", } @@ -111,6 +113,8 @@ def init_metrics(self, args): self.create_gauge("lightllm_cache_hit_rate") self.create_gauge("lightllm_gen_throughput") self.create_gauge("lightllm_num_running_reqs") + self.create_gauge("lightllm_pd_node_token_usage_ratio", labelnames=["role", "endpoint"]) + self.create_gauge("lightllm_pd_master_stage_waiting_requests", labelnames=["stage"]) def create_histogram(self, name, buckets, labelnames=None): all_labels = ["model_name"] + (labelnames or []) @@ -122,8 +126,9 @@ def create_counter(self, name, labelnames=None): counter = Counter(name, MONITOR_INFO[name], labelnames=all_labels, registry=self.registry) self.monitor_registry[name] = counter - def create_gauge(self, name): - gauge = Gauge(name, MONITOR_INFO[name], labelnames=["model_name"], registry=self.registry) + def create_gauge(self, name, labelnames=None): + all_labels = ["model_name"] + (labelnames or []) + gauge = Gauge(name, MONITOR_INFO[name], labelnames=all_labels, registry=self.registry) self.monitor_registry[name] = gauge def counter_inc(self, name, label=None): @@ -141,8 +146,11 @@ def histogram_observe(self, name, value, label=None): else: self.monitor_registry[name].labels(model_name=self.model_name, method=label).observe(value) - def gauge_set(self, name, value): - self.monitor_registry[name].labels(model_name=self.model_name).set(value) + def gauge_set(self, name, value, labels=None): + metric_labels = {"model_name": self.model_name} + if labels: + metric_labels.update(labels) + self.monitor_registry[name].labels(**metric_labels).set(value) def push_metrices(self): if self.gateway_url is not None: diff --git a/lightllm/server/multimodal_params.py b/lightllm/server/multimodal_params.py index ed3535a69f..6a417fd1d1 100644 --- a/lightllm/server/multimodal_params.py +++ b/lightllm/server/multimodal_params.py @@ -1,7 +1,6 @@ """Multimodal parameters for text generation.""" import asyncio import os -import librosa import base64 import numpy as np from typing import List, Tuple, Optional @@ -54,6 +53,8 @@ async def preload(self, request: Request): raise ValueError(f"cannot read audio which type is {self._type}!") # check if valid audio bytes + import librosa + audio_values, _ = await asyncio.to_thread(librosa.load, BytesIO(audio_data), sr=16000) audio_values = np.asarray(audio_values, dtype=np.float32) diff --git a/lightllm/server/router/manager.py b/lightllm/server/router/manager.py index 11e4e5f688..232d205b6f 100644 --- a/lightllm/server/router/manager.py +++ b/lightllm/server/router/manager.py @@ -1,4 +1,5 @@ import time +import math import uvloop import asyncio import pickle @@ -42,6 +43,23 @@ logger = init_logger(__name__) +def resolve_model_max_req_num(args: StartArgs) -> int: + """Resolve the request-state capacity allocated by each model process.""" + + configured = getattr(args, "per_dp_max_req_size", None) + if configured is None: + return args.running_max_req_size + + local_dp_size = max(1, args.dp // args.nnodes) + min_balanced_capacity = math.ceil(args.running_max_req_size / local_dp_size) + if configured < min_balanced_capacity: + raise ValueError( + "per_dp_max_req_size must be at least ceil(running_max_req_size / local_dp_size): " + f"got {configured}, minimum {min_balanced_capacity}" + ) + return configured + + class RouterManager(RouterMultiNodeTpHelper, RouterRlOpHelper, object): def __init__(self, args: StartArgs): self.args = args @@ -52,6 +70,8 @@ def __init__(self, args: StartArgs): self.node_rank = args.node_rank self.dp_size = args.dp self.schedule_time_interval = args.schedule_time_interval # 默认30ms 的调度周期 + self.prefill_coalesce_interval = max(0.0, args.prefill_coalesce_interval) + self._prefill_coalesce_deadline = None # 兼容多机纯tp的运行模式，这时候 1 // 2 == 0, 需要兼容 self.dp_size_in_node = max(1, args.dp // self.nnodes) self.dp_world_size = self.world_size // self.dp_size @@ -148,7 +168,7 @@ async def wait_to_model_ready(self): "weight_dir": self.model_weightdir, "load_way": self.load_way, "max_total_token_num": self.max_total_token_num, - "max_req_num": self.args.running_max_req_size, + "max_req_num": resolve_model_max_req_num(self.args), # MTP length stopping is asynchronous, so up to mtp_step accepted # positions may already be committed when FINISHED_LENGTH is observed. # The overlapped iteration then needs mtp_step positions for target @@ -444,6 +464,43 @@ def _generate_new_batch(self): self.schedule_new_batch = Batch.merge_two_batch(self.schedule_new_batch, new_batch) return + def _should_defer_prefill_batch(self): + """Hold a partial request burst briefly so prefill can run as one batch.""" + if self.prefill_coalesce_interval <= 0.0: + return False + + waiting_req_num = self.req_queue.get_wait_req_num() + if waiting_req_num <= 0: + self._prefill_coalesce_deadline = None + return False + + scheduled_req_num = 0 + if self.running_batch is not None: + scheduled_req_num += len(self.running_batch.reqs) + if self.schedule_new_batch is not None: + scheduled_req_num += len(self.schedule_new_batch.reqs) + runnable_slots = max(0, self.args.running_max_req_size - scheduled_req_num) + + now = time.monotonic() + if self._prefill_coalesce_deadline is None: + self._prefill_coalesce_deadline = now + self.prefill_coalesce_interval + + # There is no reason to wait once all currently runnable slots can be filled. + if runnable_slots > 0 and waiting_req_num >= runnable_slots: + self._prefill_coalesce_deadline = None + return False + + # Preserve the original deadline while the running batch is full. This lets a + # queued burst launch immediately when capacity becomes available. + if runnable_slots == 0: + return True + + if now < self._prefill_coalesce_deadline: + return True + + self._prefill_coalesce_deadline = None + return False + async def _recv_new_reqs_and_schedule(self): if not hasattr(self, "recv_max_count"): self.recv_max_count = 64 @@ -470,7 +527,7 @@ async def _recv_new_reqs_and_schedule(self): if self.is_multinode_tp: self.multinode_tp_generate_new_batch() else: - if self._get_paused_req_num() == 0: + if self._get_paused_req_num() == 0 and not self._should_defer_prefill_batch(): self._generate_new_batch() return diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 28f2abf74b..820ed778e4 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -375,12 +375,20 @@ def _async_copy_next_token_infos_to_pin_mem( ) return next_token_ids_cpu, next_token_logprobs_cpu, next_token_ranks_cpu - def _get_next_token_ranks(self, logits: torch.Tensor, next_token_ids: torch.Tensor) -> torch.Tensor: + def _get_next_token_ranks(self, model_output: ModelOutput, next_token_ids: torch.Tensor) -> torch.Tensor: """计算（或占位）每个 next token 在 vocab 上的 1-based rank（GPU tensor）。 仅 ``--enable_rl`` 时做真实 rank；否则返回 GPU 常量 ``-1``，避免 O(batch * vocab) 比较。 下游 async_copy 在同样条件下会忽略该返回值。 """ + if model_output.has_vocab_parallel_logits: + assert model_output.logits.shape[1] == 1 + return g_pin_mem_manager.get_const_gpu_tensor( + key="next_token_ranks", + shape=next_token_ids.shape, + fill_value=1 if self.args.enable_rl else -1, + dtype=torch.int32, + ) if not self.args.enable_rl: return g_pin_mem_manager.get_const_gpu_tensor( key="next_token_ranks", @@ -388,8 +396,8 @@ def _get_next_token_ranks(self, logits: torch.Tensor, next_token_ids: torch.Tens fill_value=-1, dtype=torch.int32, ) - selected_logits = logits.gather(1, next_token_ids.long().view(-1, 1)) - return (logits > selected_logits).sum(dim=-1, dtype=torch.int32) + 1 + selected_logits = model_output.logits.gather(1, next_token_ids.long().view(-1, 1)) + return (model_output.logits > selected_logits).sum(dim=-1, dtype=torch.int32) + 1 def _capture_prompt_logprobs_if_needed( self, @@ -860,18 +868,35 @@ def _trans_req_ids_to_req_objs(self, req_ids: List[int]) -> List[InferReq]: return [g_infer_context.requests_mapping[req_id] for req_id in req_ids] def _gen_argmax_token_ids(self, model_output: ModelOutput): + if model_output.mtp_collector.draft_token_ids is not None: + return model_output.mtp_collector.draft_token_ids logits = model_output.logits - return torch.argmax(logits, dim=-1) + candidate_indexes = torch.argmax(logits, dim=-1) + return self._map_logits_indexes_to_token_ids(model_output, candidate_indexes) def _gen_argmax_token_ids_and_prob(self, model_output: ModelOutput): + if model_output.mtp_collector.draft_token_ids is not None: + draft_token_probs = model_output.mtp_collector.draft_token_probs + if draft_token_probs is None: + raise RuntimeError("draft head returned token ids without token probabilities") + return model_output.mtp_collector.draft_token_ids, draft_token_probs logits = model_output.logits - probs = torch.softmax(logits, dim=-1) - max_probs, draft_next_token_ids_gpu = torch.max(probs, dim=-1) - return draft_next_token_ids_gpu, max_probs + if model_output.has_vocab_parallel_logits: + max_logits, candidate_indexes = torch.max(logits, dim=-1) + token_ids = self._map_logits_indexes_to_token_ids(model_output, candidate_indexes) + return token_ids, torch.exp(max_logits - model_output.logits_logsumexp) + max_probs, token_ids = torch.max(torch.softmax(logits, dim=-1), dim=-1) + return token_ids, max_probs + + @staticmethod + def _map_logits_indexes_to_token_ids(model_output: ModelOutput, candidate_indexes: torch.Tensor): + if not model_output.has_vocab_parallel_logits: + return candidate_indexes + return model_output.logits_token_ids.gather(1, candidate_indexes.long().view(-1, 1)).view(-1).long() def _sample_and_scatter_token( self, - logits: torch.Tensor, + model_output: ModelOutput, b_req_idx: torch.Tensor, b_mtp_index: torch.Tensor, run_reqs: List[InferReq], @@ -880,12 +905,15 @@ def _sample_and_scatter_token( mask_func: Optional[Callable] = None, ): + logits = model_output.logits + if mask_func is not None: + assert not model_output.has_vocab_parallel_logits, "constrained sampling requires dense logits" assert len(run_reqs) == logits.shape[0] mask_func(run_reqs, logits) - next_token_ids, next_token_logprobs = sample(logits, run_reqs, self.eos_id) - next_token_ranks = self._get_next_token_ranks(logits, next_token_ids) + next_token_ids, next_token_logprobs = sample(model_output, run_reqs, self.eos_id) + next_token_ranks = self._get_next_token_ranks(model_output, next_token_ids) b_has_out = None if is_prefill: b_has_out = g_pin_mem_manager.gen_from_list( diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index 4d09476849..c639c8c5c9 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py @@ -1,3 +1,5 @@ +import os + import torch import time from typing import List @@ -27,6 +29,15 @@ class ChunkedPrefillBackend(ModeBackend): def __init__(self) -> None: super().__init__() + # Mega-MoE owns one symmetric communication workspace per rank. Keep + # the CPU pre/post pipeline enabled, but do not let its two host threads + # enqueue a second model forward while that workspace is still live. + self._serialize_sm90_mega_moe_forwards = os.getenv("LIGHTLLM_ENABLE_SM90_MEGA_MOE", "0").upper() in { + "1", + "ON", + "TRUE", + } + # 用于控制每一步是执行prefill 和 decode 还是跳过 self.control_state_machine = ControlState() @@ -41,6 +52,24 @@ def __init__(self) -> None: self.classed_req_strict_prefill = False return + def _record_forward_completion(self) -> torch.cuda.Event: + sync_event = torch.cuda.Event() + sync_event.record() + return sync_event + + def _notify_next_forward_when_safe(self, event_pack: OverlapEventPack, sync_event: torch.cuda.Event): + if self._serialize_sm90_mega_moe_forwards: + # Mega-MoE owns one symmetric communication workspace per rank. Let + # this iteration's CPU bookkeeping overlap its GPU work, but wait for + # the workspace to be released before waking the next forward thread. + sync_event.synchronize() + event_pack.notify_forward_and_wait_post_handle() + else: + # Preserve the regular two-forward overlap path. + event_pack.notify_forward_and_wait_post_handle() + sync_event.synchronize() + return + def init_spec_engine(self): self.spec_engine = SpecEngine( backend=self, @@ -107,11 +136,13 @@ def prefill_normal( ): # 第一阶段: 模型推理 model_input, run_reqs = prepare_prefill_inputs(prefill_reqs, is_chuncked_mode=not self.disable_chunked_prefill) + if self.prefill_mask_func is not None: + model_input.use_vocab_parallel_greedy = False with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output = self.model.forward(model_input) self._capture_prompt_logprobs_if_needed(model_input, run_reqs, model_output.prompt_logics) (_, next_token_ids_cpu, next_token_logprobs_cpu, next_token_ranks_cpu,) = self._sample_and_scatter_token( - logits=model_output.logits, + model_output=model_output, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, run_reqs=run_reqs, @@ -123,16 +154,14 @@ def prefill_normal( b_req_idx=model_input.b_req_idx, reqs=run_reqs, ) - sync_event = torch.cuda.Event() - sync_event.record() + sync_event = self._record_forward_completion() # 第二阶段 event_pack.notify_post_handle_and_wait_pre_post_handle() update_packs = self._pre_post_handle(run_reqs, is_chuncked_mode=not self.disable_chunked_prefill) # 第三阶段 - event_pack.notify_forward_and_wait_post_handle() - sync_event.synchronize() + self._notify_next_forward_when_safe(event_pack, sync_event) self._post_handle( run_reqs=run_reqs, next_token_ids=next_token_ids_cpu, @@ -152,26 +181,26 @@ def decode_normal( decode_reqs: List[InferReq], ): model_input, run_reqs = prepare_decode_inputs(decode_reqs) + if self.decode_mask_func is not None: + model_input.use_vocab_parallel_greedy = False with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output = self.model.forward(model_input) (_, next_token_ids_cpu, next_token_logprobs_cpu, next_token_ranks_cpu,) = self._sample_and_scatter_token( - logits=model_output.logits, + model_output=model_output, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, run_reqs=run_reqs, is_prefill=False, mask_func=self.decode_mask_func, ) - sync_event = torch.cuda.Event() - sync_event.record() + sync_event = self._record_forward_completion() # 第二阶段 event_pack.notify_post_handle_and_wait_pre_post_handle() update_packs = self._pre_post_handle(run_reqs, is_chuncked_mode=False) # 第三阶段 - event_pack.notify_forward_and_wait_post_handle() - sync_event.synchronize() + self._notify_next_forward_when_safe(event_pack, sync_event) self._post_handle( run_reqs=run_reqs, next_token_ids=next_token_ids_cpu, @@ -191,6 +220,8 @@ def prefill_mtp( prefill_reqs: List[InferReq], ): model_input, run_reqs = prepare_prefill_inputs(prefill_reqs, is_chuncked_mode=not self.disable_chunked_prefill) + if self.prefill_mask_func is not None: + model_input.use_vocab_parallel_greedy = False with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output = self.model.forward(model_input) self._capture_prompt_logprobs_if_needed(model_input, run_reqs, model_output.prompt_logics) @@ -200,7 +231,7 @@ def prefill_mtp( next_token_logprobs_cpu, next_token_ranks_cpu, ) = self._sample_and_scatter_token( - logits=model_output.logits, + model_output=model_output, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, run_reqs=run_reqs, @@ -219,16 +250,14 @@ def prefill_mtp( b_req_idx=model_input.b_req_idx, reqs=run_reqs, ) - sync_event = torch.cuda.Event() - sync_event.record() + sync_event = self._record_forward_completion() # 第二阶段 event_pack.notify_post_handle_and_wait_pre_post_handle() update_packs = self._pre_post_handle(run_reqs, is_chuncked_mode=not self.disable_chunked_prefill) # 第三阶段 - event_pack.notify_forward_and_wait_post_handle() - sync_event.synchronize() + self._notify_next_forward_when_safe(event_pack, sync_event) self._post_handle( run_reqs=run_reqs, @@ -251,6 +280,8 @@ def decode_mtp( ): """Run the speculative draft-and-verify decode flow.""" model_input, run_reqs = prepare_decode_inputs(decode_reqs) + if self.decode_mask_func is not None: + model_input.use_vocab_parallel_greedy = False spec_engine = self.spec_engine req_num = len(decode_reqs) @@ -272,11 +303,11 @@ def decode_mtp( selected_rows = async_selected_row_mask_cpu.tensor.tolist() run_reqs = [req for req, selected in zip(run_reqs, selected_rows) if selected] next_token_ids, next_token_logprobs = sample( - model_output.logits, + model_output, run_reqs, self.eos_id, ) - next_token_ranks = self._get_next_token_ranks(model_output.logits, next_token_ids) + next_token_ranks = self._get_next_token_ranks(model_output, next_token_ids) b_req_mtp_start_loc = gen_b_req_mtp_start_loc(model_input.b_mtp_index, num_reqs=req_num) mtp_accept_len, accepted_index = mtp_utils.verify_mtp_tokens( @@ -331,8 +362,7 @@ def decode_mtp( next_token_ranks=next_token_ranks, ) - sync_event = torch.cuda.Event() - sync_event.record() + sync_event = self._record_forward_completion() # 第二阶段 event_pack.notify_post_handle_and_wait_pre_post_handle() @@ -350,8 +380,7 @@ def decode_mtp( update_packs = self._pre_post_handle(verify_ok_reqs, is_chuncked_mode=False) # 第三阶段 - event_pack.notify_forward_and_wait_post_handle() - sync_event.synchronize() + self._notify_next_forward_when_safe(event_pack, sync_event) spec_engine.update_planner_statics( plan=spec_plan, diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl_for_reward_model.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl_for_reward_model.py index dfb1020820..65161ca8fa 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl_for_reward_model.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl_for_reward_model.py @@ -17,6 +17,7 @@ def reward_prefill(self, event_pack: OverlapEventPack, prefill_reqs: List[InferR assert self.disable_chunked_prefill is True model_input, run_reqs = prepare_prefill_inputs(prefill_reqs, is_chuncked_mode=not self.disable_chunked_prefill) + model_input.use_vocab_parallel_greedy = False model_output = self.model.forward(model_input) scores: torch.Tensor = model_output.logits diff --git a/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py index 21979cbef0..2b6eb0b9d3 100644 --- a/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py @@ -38,12 +38,11 @@ def beam_prefill(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq model_input, group_run_reqs = prepare_prefill_inputs( group_reqs, is_chuncked_mode=not self.disable_chunked_prefill ) + model_input.use_vocab_parallel_greedy = False with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output = self.model.forward(model_input) - logits = model_output.logits - batch_idx, run_reqs = self._diverse_copy( master_reqs=group_reqs, b_prefill_has_out=model_input.b_prefill_has_output_cpu ) @@ -60,11 +59,11 @@ def beam_prefill(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq non_blocking=True ) - logits = logits[batch_idx] + sampled_output = model_output.index_select_logits_rows(batch_idx) b_mtp_index = model_input.b_mtp_index[batch_idx] - next_token_ids, next_token_logprobs = sample(logits, run_reqs, self.eos_id) - next_token_ranks = self._get_next_token_ranks(logits, next_token_ids) + next_token_ids, next_token_logprobs = sample(sampled_output, run_reqs, self.eos_id) + next_token_ranks = self._get_next_token_ranks(sampled_output, next_token_ids) scatter_token( next_token_ids=next_token_ids, diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py index 9a81927bc1..a8c1bd8a40 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py @@ -187,7 +187,7 @@ def prefill_normal( next_token_logprobs_cpu, next_token_ranks_cpu, ) = self._sample_and_scatter_token( - logits=model_output.logits, + model_output=model_output, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, run_reqs=run_reqs, @@ -240,7 +240,7 @@ def decode_normal(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq next_token_logprobs_cpu, next_token_ranks_cpu, ) = self._sample_and_scatter_token( - logits=model_output.logits, + model_output=model_output, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, run_reqs=run_reqs, @@ -287,13 +287,8 @@ def prefill_overlap(self, event_pack: OverlapEventPack, prefill_reqs: List[Infer model_output0, model_output1 = self.model.microbatch_overlap_prefill(model_input0, model_input1) self._capture_prompt_logprobs_if_needed(model_input0, run_reqs0, model_output0.prompt_logics) self._capture_prompt_logprobs_if_needed(model_input1, run_reqs1, model_output1.prompt_logics) - logits0 = model_output0.logits - logits1 = model_output1.logits - req_num0, req_num1 = len(run_reqs0), len(run_reqs1) - logits = torch.empty((req_num0 + req_num1, logits0.shape[1]), dtype=logits0.dtype, device=logits0.device) - logits[0:req_num0, :].copy_(logits0, non_blocking=True) - logits[req_num0 : req_num0 + req_num1, :].copy_(logits1, non_blocking=True) + sampled_output = ModelOutput.concat_logits_rows([model_output0, model_output1]) run_reqs = run_reqs0 + run_reqs1 b_has_out_cpu = model_input0.b_prefill_has_output_cpu + model_input1.b_prefill_has_output_cpu @@ -307,7 +302,7 @@ def prefill_overlap(self, event_pack: OverlapEventPack, prefill_reqs: List[Infer next_token_logprobs_cpu, next_token_ranks_cpu, ) = self._sample_and_scatter_token( - logits=logits, + model_output=sampled_output, b_req_idx=b_req_idx, b_mtp_index=b_mtp_index, run_reqs=run_reqs, @@ -356,7 +351,7 @@ def decode_overlap(self, event_pack: OverlapEventPack, decode_reqs: List[InferRe with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output0, model_output1 = self.model.microbatch_overlap_decode(model_input0, model_input1) if req_num0 + req_num1 > 0: - logits = torch.cat((model_output0.logits, model_output1.logits), dim=0) + sampled_output = ModelOutput.concat_logits_rows([model_output0, model_output1]) b_req_idx = torch.cat((model_input0.b_req_idx, model_input1.b_req_idx), dim=0) b_mtp_index = torch.cat((model_input0.b_mtp_index, model_input1.b_mtp_index), dim=0) ( @@ -365,7 +360,7 @@ def decode_overlap(self, event_pack: OverlapEventPack, decode_reqs: List[InferRe next_token_logprobs_cpu, next_token_ranks_cpu, ) = self._sample_and_scatter_token( - logits=logits, + model_output=sampled_output, b_req_idx=b_req_idx, b_mtp_index=b_mtp_index, run_reqs=run_reqs, @@ -421,7 +416,7 @@ def prefill_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq] next_token_logprobs_cpu, next_token_ranks_cpu, ) = self._sample_and_scatter_token( - logits=model_output.logits, + model_output=model_output, b_req_idx=b_req_idx, b_mtp_index=b_mtp_index, run_reqs=run_reqs, @@ -499,11 +494,11 @@ def decode_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq]): if req_num > 0: next_token_ids, next_token_logprobs = sample( - model_output.logits, + model_output, run_reqs, self.eos_id, ) - next_token_ranks = self._get_next_token_ranks(model_output.logits, next_token_ids) + next_token_ranks = self._get_next_token_ranks(model_output, next_token_ids) b_req_mtp_start_loc = gen_b_req_mtp_start_loc( b_mtp_index=model_input.b_mtp_index, @@ -649,17 +644,9 @@ def prefill_overlap_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[I model_output0, model_output1 = self.model.microbatch_overlap_prefill(model_input0, model_input1) self._capture_prompt_logprobs_if_needed(model_input0, run_reqs0, model_output0.prompt_logics) self._capture_prompt_logprobs_if_needed(model_input1, run_reqs1, model_output1.prompt_logics) - logits0 = model_output0.logits - logits1 = model_output1.logits req_num0, req_num1 = len(run_reqs0), len(run_reqs1) req_num = req_num0 + req_num1 - logits = torch.empty( - (req_num0 + req_num1, logits0.shape[1]), - dtype=logits0.dtype, - device=logits0.device, - ) - logits[0:req_num0, :].copy_(logits0, non_blocking=True) - logits[req_num0 : (req_num0 + req_num1), :].copy_(logits1, non_blocking=True) + sampled_output = ModelOutput.concat_logits_rows([model_output0, model_output1]) run_reqs = run_reqs0 + run_reqs1 b_has_out_cpu = model_input0.b_prefill_has_output_cpu + model_input1.b_prefill_has_output_cpu @@ -673,7 +660,7 @@ def prefill_overlap_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[I next_token_logprobs_cpu, next_token_ranks_cpu, ) = self._sample_and_scatter_token( - logits=logits, + model_output=sampled_output, run_reqs=run_reqs, b_req_idx=b_req_idx, b_mtp_index=b_mtp_index, @@ -681,7 +668,7 @@ def prefill_overlap_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[I b_prefill_has_output_cpu=b_has_out_cpu, ) else: - next_token_ids = torch.empty((0,), dtype=torch.int64, device=logits.device) + next_token_ids = torch.empty((0,), dtype=torch.int64, device=sampled_output.logits.device) target_next_token_ids_gpu0 = next_token_ids[:req_num0] target_next_token_ids_gpu1 = next_token_ids[req_num0:] @@ -769,20 +756,12 @@ def decode_overlap_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[Inf verify_row_num0 = model_input0.batch_size verify_row_num1 = model_input1.batch_size verify_row_num = verify_row_num0 + verify_row_num1 - logits0 = model_output0.logits - logits1 = model_output1.logits run_reqs = run_reqs0 + run_reqs1 if req_num > 0: assert len(run_reqs) == verify_row_num - logits = torch.empty( - (verify_row_num, logits0.shape[1]), - dtype=logits0.dtype, - device=logits0.device, - ) - logits[:verify_row_num0, :].copy_(logits0, non_blocking=True) - logits[verify_row_num0:, :].copy_(logits1, non_blocking=True) - next_token_ids, next_token_logprobs = sample(logits, run_reqs, self.eos_id) - next_token_ranks = self._get_next_token_ranks(logits, next_token_ids) + sampled_output = ModelOutput.concat_logits_rows([model_output0, model_output1]) + next_token_ids, next_token_logprobs = sample(sampled_output, run_reqs, self.eos_id) + next_token_ranks = self._get_next_token_ranks(sampled_output, next_token_ids) ( next_token_ids_cpu, next_token_logprobs_cpu, diff --git a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py index 5b29ea0510..ad428d32af 100644 --- a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py +++ b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py @@ -1,5 +1,9 @@ import torch from typing import List, Tuple +from lightllm.common.basemodel.batch_objs import ModelOutput +from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_greedy import ( + is_vocab_parallel_greedy_enabled, +) from lightllm.common.basemodel.triton_kernel.post_process.apply_penalty import apply_penalty from lightllm.common.basemodel.triton_kernel.post_process.apply_penalty_gpu_cache import apply_penalty_gpu_cache from lightllm.common.basemodel.triton_kernel.post_process.apply_invalid_token import apply_invalid_token_ids @@ -8,7 +12,44 @@ from lightllm.utils.envs_utils import get_env_start_args -def sample(logits: torch.Tensor, reqs: List[InferReq], eos_id: List[int] = [2]): +def _can_use_unmodified_greedy_logits(reqs: List[InferReq]) -> bool: + """Whether sampling is exactly argmax over the incoming logits.""" + + for req_obj in reqs: + sample_param = req_obj.sampling_param + shm_param = sample_param.shm_param + if shm_param.top_k != 1 or shm_param.temperature != 1.0: + return False + if ( + shm_param.presence_penalty != 0.0 + or shm_param.frequency_penalty != 0.0 + or shm_param.repetition_penalty != 1.0 + ): + return False + if shm_param.exponential_decay_length_penalty.to_tuple()[1] != 1.0: + return False + out_token_len = req_obj.get_cur_total_len() - req_obj.shm_req.input_len + if out_token_len < shm_param.min_new_tokens - 1: + return False + if sample_param.invalid_token_ids: + return False + return True + + +def can_use_vocab_parallel_greedy(reqs: List[InferReq]) -> bool: + return is_vocab_parallel_greedy_enabled() and _can_use_unmodified_greedy_logits(reqs) + + +def sample(model_output: ModelOutput, reqs: List[InferReq], eos_id: List[int] = [2]): + logits = model_output.logits + if model_output.has_vocab_parallel_logits: + if not _can_use_unmodified_greedy_logits(reqs): + raise RuntimeError("vocab-parallel logits require unmodified greedy requests") + candidate_indexes = torch.argmax(logits, dim=-1, keepdim=True) + token_ids = model_output.logits_token_ids.gather(1, candidate_indexes).view(-1).long() + selected_logits = logits.gather(1, candidate_indexes).view(-1) + return token_ids, selected_logits - model_output.logits_logsumexp + ( b_req_idx, b_temperatures, diff --git a/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py b/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py index 22731439c4..124e3d6d90 100644 --- a/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py +++ b/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py @@ -3,6 +3,9 @@ from typing import List, Tuple from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context from lightllm.common.basemodel.batch_objs import ModelInput +from lightllm.server.router.model_infer.mode_backend.generic_post_process import ( + can_use_vocab_parallel_greedy, +) INT64_MAX = torch.iinfo(torch.int64).max @@ -13,6 +16,7 @@ def prepare_prefill_inputs(req_objs: List[InferReq], is_chuncked_mode: bool) -> input_ids = [] b_req_idx = [] b_seq_len = [] + b_input_len = [] b_q_seq_len = [] batch_multimodal_params = [] b_ready_cache_len = [] @@ -38,6 +42,7 @@ def prepare_prefill_inputs(req_objs: List[InferReq], is_chuncked_mode: bool) -> input_id = input_token_ids[req.cur_kv_len :] b_seq_len.append(seq_len) + b_input_len.append(req.shm_req.input_len) b_q_seq_len.append(input_token_len) input_ids.append(input_id) total_token_num += seq_len @@ -59,6 +64,7 @@ def prepare_prefill_inputs(req_objs: List[InferReq], is_chuncked_mode: bool) -> input_ids = torch.tensor(input_ids, dtype=torch.int64, device="cpu") b_req_idx = torch.tensor(b_req_idx, dtype=torch.int32, device="cpu") b_seq_len = torch.tensor(b_seq_len, dtype=torch.int32, device="cpu") + b_input_len = torch.tensor(b_input_len, dtype=torch.int32, device="cpu") b_is_decode_req = torch.tensor(b_is_decode_req, dtype=torch.bool, device="cpu") b_mtp_index = torch.tensor(b_mtp_index, dtype=torch.int32, device="cpu") b_ready_cache_len = torch.tensor(b_ready_cache_len, dtype=torch.int32, device="cpu") @@ -81,12 +87,14 @@ def prepare_prefill_inputs(req_objs: List[InferReq], is_chuncked_mode: bool) -> b_req_idx=b_req_idx, b_mtp_index=b_mtp_index, b_seq_len=b_seq_len, + b_input_len=b_input_len, b_is_decode_req=b_is_decode_req, b_ready_cache_len=b_ready_cache_len, b_prefill_start_loc=b_prefill_start_loc, is_prefill=True, b_prefill_has_output_cpu=b_prefill_has_output, multimodal_params=batch_multimodal_params, + use_vocab_parallel_greedy=can_use_vocab_parallel_greedy(run_reqs), ) return model_input, run_reqs @@ -98,6 +106,7 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In b_req_idx = [] b_mtp_index = [] b_seq_len = [] + b_input_len = [] b_q_seq_len = [] multimodal_params = [] for req in req_objs: @@ -106,6 +115,7 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In seq_len = req.get_cur_total_len() assert req.cur_kv_len == seq_len - 1, f"{req.cur_kv_len} {seq_len}" b_seq_len.append(seq_len) + b_input_len.append(req.shm_req.input_len) b_q_seq_len.append(1) total_token_num += seq_len b_mtp_index.append(0) @@ -116,6 +126,7 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In b_req_idx.append(req.req_idx) seq_len += 1 b_seq_len.append(seq_len) + b_input_len.append(req.shm_req.input_len) total_token_num += seq_len b_mtp_index.append(step + 1) multimodal_params.append(req.multimodal_params) @@ -128,6 +139,7 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In b_req_idx = torch.tensor(b_req_idx, dtype=torch.int32, device="cpu") b_seq_len = torch.tensor(b_seq_len, dtype=torch.int32, device="cpu") + b_input_len = torch.tensor(b_input_len, dtype=torch.int32, device="cpu") b_mtp_index = torch.tensor(b_mtp_index, dtype=torch.int32, device="cpu") b_position_delta = build_b_position_delta(multimodal_params) @@ -155,11 +167,13 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In b_req_idx=b_req_idx, b_mtp_index=b_mtp_index, b_seq_len=b_seq_len, + b_input_len=b_input_len, b_position_delta=b_position_delta, b_shared_seq_len=b_shared_seq_len, b_shared_radix_node_id=b_shared_radix_node_id, is_prefill=False, multimodal_params=multimodal_params, + use_vocab_parallel_greedy=can_use_vocab_parallel_greedy(run_reqs), ) return model_input, run_reqs @@ -176,6 +190,9 @@ def overlap_prepare_decode_inputs(req_objs: List[InferReq]): model_input1, run_reqs1 = prepare_decode_inputs( req_objs=decode_reqs1, ) + use_vocab_parallel_greedy = can_use_vocab_parallel_greedy(run_reqs0 + run_reqs1) + model_input0.use_vocab_parallel_greedy = use_vocab_parallel_greedy + model_input1.use_vocab_parallel_greedy = use_vocab_parallel_greedy return model_input0, run_reqs0, decode_reqs0, model_input1, run_reqs1, decode_reqs1 @@ -211,6 +228,9 @@ def overlap_prepare_prefill_inputs(req_objs: List[InferReq]): req_objs=right_reqs, is_chuncked_mode=True, ) + use_vocab_parallel_greedy = can_use_vocab_parallel_greedy(run_reqs0 + run_reqs1) + model_input0.use_vocab_parallel_greedy = use_vocab_parallel_greedy + model_input1.use_vocab_parallel_greedy = use_vocab_parallel_greedy return model_input0, run_reqs0, model_input1, run_reqs1 diff --git a/lightllm/server/router/model_infer/mtp_speculative/dp_overlap_proposers/eagle_with_att.py b/lightllm/server/router/model_infer/mtp_speculative/dp_overlap_proposers/eagle_with_att.py index 6b8c23e8fd..137a580ae6 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/dp_overlap_proposers/eagle_with_att.py +++ b/lightllm/server/router/model_infer/mtp_speculative/dp_overlap_proposers/eagle_with_att.py @@ -141,7 +141,7 @@ def propose_next_overlap( req_num_by_batch, ) ): - accepted_tail_output = ModelOutput(logits=extend_output.logits.index_select(0, accepted_tail_rows)) + accepted_tail_output = extend_output.index_select_logits_rows(accepted_tail_rows) if self.enable_dynmaic_mtp: draft_token_ids, draft_token_probs = self._gen_argmax_token_ids_and_prob(accepted_tail_output) draft_token_probs = draft_token_probs.float() diff --git a/lightllm/server/router/model_infer/mtp_speculative/planner/lightspec.py b/lightllm/server/router/model_infer/mtp_speculative/planner/lightspec.py index b25c62b9b7..a021021c86 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/planner/lightspec.py +++ b/lightllm/server/router/model_infer/mtp_speculative/planner/lightspec.py @@ -211,7 +211,15 @@ def _get_draft_cost_ms(self, req_num: int, verify_batch_size: int, draft_step: i if self.spec_mode in ("vanilla_no_att", "eagle_no_att"): return self.draft_infer_costs.estimate(req_num) * draft_step - if self.spec_mode in ("vanilla_with_att", "eagle_with_att", "eagle3"): + if self.spec_mode == "vanilla_with_att": + assert draft_step > 0, f"{self.spec_mode} requires draft_step to be greater than 0" + # Every level in the chained Vanilla proposer forwards the full + # (possibly compacted) verify layout so that its fixed-depth KV + # state covers every committed position. Pricing only the first + # level at B and the remaining levels at N systematically makes a + # wide verify budget look cheaper than the work actually run. + return self.draft_infer_costs.estimate(verify_batch_size) * draft_step + if self.spec_mode in ("eagle_with_att", "eagle3"): assert draft_step > 0, f"{self.spec_mode} requires draft_step to be greater than 0" draft_cost_ms = self.draft_infer_costs.estimate(verify_batch_size) if draft_step > 1: diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py index 3d2c0a0e86..fc3155b87f 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py @@ -79,14 +79,25 @@ def propose_next( verify_draft_input.mtp_draft_input_hiddens = target_model_output.mtp_collector.spec_hidden extend_output = draft_model.forward(verify_draft_input) - # 只在 req_num 行 logits 上进行 argmax，避免为未接受的 verify 行执行 - # vocabulary reduction。第一列 proposal 来自每个请求的 accepted tail。 - accepted_tail_output = ModelOutput(logits=extend_output.logits.index_select(0, accepted_tail_rows)) - if self.enable_dynmaic_mtp: - draft_token_ids, draft_token_probs = self._gen_argmax_token_ids_and_prob(accepted_tail_output) - schedule_scores_by_step.append(draft_token_probs.float().unsqueeze(1)) + # 第一列 proposal 来自每个请求的 accepted tail。Specialized draft + # heads can return vocab-parallel Top-1 directly; generic heads retain + # the full-logits fallback. + head_token_ids = extend_output.mtp_collector.draft_token_ids + if head_token_ids is not None: + draft_token_ids = head_token_ids.index_select(0, accepted_tail_rows) + if self.enable_dynmaic_mtp: + head_token_probs = extend_output.mtp_collector.draft_token_probs + if head_token_probs is None: + raise RuntimeError("draft head returned token ids without token probabilities") + draft_token_probs = head_token_probs.index_select(0, accepted_tail_rows) + schedule_scores_by_step.append(draft_token_probs.float().unsqueeze(1)) else: - draft_token_ids = self._gen_argmax_token_ids(accepted_tail_output) + accepted_tail_output = extend_output.index_select_logits_rows(accepted_tail_rows) + if self.enable_dynmaic_mtp: + draft_token_ids, draft_token_probs = self._gen_argmax_token_ids_and_prob(accepted_tail_output) + schedule_scores_by_step.append(draft_token_probs.float().unsqueeze(1)) + else: + draft_token_ids = self._gen_argmax_token_ids(accepted_tail_output) proposal_token_ids_by_step.append(draft_token_ids.unsqueeze(1)) if draft_step == 1: diff --git a/lightllm/server/tokenizer.py b/lightllm/server/tokenizer.py index e1a4e421d1..bf2acfbbb2 100644 --- a/lightllm/server/tokenizer.py +++ b/lightllm/server/tokenizer.py @@ -122,6 +122,14 @@ def get_tokenizer( tokenizer = QWen3_5Tokenizer( tokenizer=tokenizer, image_processor=processor.image_processor, model_cfg=model_cfg ) + elif model_type == "glm5_next" and model_cfg.get("vision_config") is not None: + from transformers import AutoProcessor + from ..models.glm5_next.tokenizer import Glm5NextTokenizer + + processor = AutoProcessor.from_pretrained(tokenizer_name) + tokenizer = Glm5NextTokenizer( + tokenizer=tokenizer, image_processor=processor.image_processor, model_cfg=model_cfg + ) elif model_cfg.get("thinker_config") is not None: from transformers import AutoProcessor diff --git a/lightllm/server/visualserver/model_infer/model_rpc.py b/lightllm/server/visualserver/model_infer/model_rpc.py index 68e0a97ca1..e67e381ac8 100644 --- a/lightllm/server/visualserver/model_infer/model_rpc.py +++ b/lightllm/server/visualserver/model_infer/model_rpc.py @@ -21,6 +21,7 @@ from lightllm.models.qwen3_vl.qwen3_visual import Qwen3VisionTransformerPretrainedModel from lightllm.models.tarsier2.tarsier2_visual import TarsierVisionTransformerPretrainedModel from lightllm.models.qwen3_omni_moe_thinker.qwen3_omni_visual import Qwen3OmniMoeVisionTransformerPretrainedModel +from lightllm.models.glm5_next.glm5_next_visual import Glm5NextVisionModel from lightllm.utils.infer_utils import set_random_seed from lightllm.utils.dist_utils import init_vision_distributed_env from lightllm.utils.envs_utils import get_env_start_args @@ -91,6 +92,10 @@ def exposed_init_model(self, kvargs): self.model = ( Qwen3VisionTransformerPretrainedModel(kvargs, **model_cfg["vision_config"]).eval().bfloat16() ) + elif self.model_type == "glm5_next" and model_cfg.get("vision_config") is not None: + if self.vit_tp != 1: + raise ValueError("GLM-5 vision supports --visual_tp 1; use --visual_dp for parallelism") + self.model = Glm5NextVisionModel(data_type=self.data_type) elif model_cfg["architectures"][0] == "TarsierForConditionalGeneration": self.model = TarsierVisionTransformerPretrainedModel(**model_cfg).eval().bfloat16() elif self.model_type == "llava": diff --git a/lightllm/utils/config_utils.py b/lightllm/utils/config_utils.py index 6695f0ec44..8fec95780e 100644 --- a/lightllm/utils/config_utils.py +++ b/lightllm/utils/config_utils.py @@ -398,6 +398,8 @@ def has_vision_module(model_path: str) -> bool: # Qwen3VisionTransformerPretrainedModel model_cfg["vision_config"] return True + elif model_type == "glm5_next": + return model_cfg.get("vision_config") is not None elif model_cfg["architectures"][0] == "TarsierForConditionalGeneration": # TarsierVisionTransformerPretrainedModel return True diff --git a/lightllm/utils/device_utils.py b/lightllm/utils/device_utils.py index 58bff90560..1bf4ff7fee 100644 --- a/lightllm/utils/device_utils.py +++ b/lightllm/utils/device_utils.py @@ -45,6 +45,11 @@ def is_sm100_gpu(): return torch.cuda.get_device_capability()[0] == 10 +@lru_cache(maxsize=None) +def is_sm90_gpu(): + return torch.cuda.get_device_capability() == (9, 0) + + @lru_cache(maxsize=None) def get_device_sm_regs_num(): import triton diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index 4fed9509a9..2cd48338e9 100644 --- a/lightllm/utils/envs_utils.py +++ b/lightllm/utils/envs_utils.py @@ -83,8 +83,30 @@ def get_deepep_num_max_dispatch_tokens_per_rank_prefill(): @lru_cache(maxsize=None) def get_deepep_num_max_dispatch_tokens_per_rank_decode(): - # 该参数需要大于单卡最大batch size，且是8的倍数。该参数与显存占用直接相关，值越大，显存占用越大，如果出现显存不足，可以尝试调小该值 - return int(os.getenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE", 256)) + # DeepEP requires this limit to cover the physical rows passed to one + # decode dispatch and to be a multiple of eight. Speculative target + # decode widens every logical request to ``mtp_step + 1`` rows, so the old + # fixed default of 256 failed as soon as (for example) a 64-request MTP5 + # CUDA Graph reached 264 rows. + configured = os.getenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE", None) + if configured is not None: + return int(configured) + + args = get_env_start_args() + logical_batch_size = max( + int(args.get("graph_max_batch_size", 0) or 0), + int(args.get("running_max_req_size", 0) or 0), + 1, + ) + verify_width = int(args.get("mtp_step", 0) or 0) + 1 if args.get("mtp_mode", None) is not None else 1 + required_tokens = logical_batch_size * verify_width + # In TP/SP + EP mode each rank dispatches only its sequence-parallel slice + # to DeepEP. CUDA Graph batch sizes are TP-aligned, but use ceil here as a + # defensive bound for non-graph decode as well. + if args.get("enable_tpsp_mix_mode", False) and args.get("enable_ep_moe", False): + tp = max(int(args.get("tp", 1) or 1), 1) + required_tokens = (required_tokens + tp - 1) // tp + return ((max(required_tokens, 256) + 7) // 8) * 8 @lru_cache(maxsize=None) diff --git a/test/kernel/test_extract_indexer_ks_dynamic.py b/test/kernel/test_extract_indexer_ks_dynamic.py new file mode 100644 index 0000000000..f948d598e6 --- /dev/null +++ b/test/kernel/test_extract_indexer_ks_dynamic.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from lightllm.models.deepseek3_2.triton_kernel.extract_indexer_ks import ( + extract_indexer_ks, + extract_indexer_ks_dynamic, +) + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def _indexer_buffer(num_slots: int): + keys = ( + torch.arange(num_slots * 128, device="cuda", dtype=torch.float32) + .remainder_(31) + .sub_(15) + .to(torch.float8_e4m3fn) + .view(num_slots, 1, 128) + ) + scales = torch.arange(1, num_slots + 1, device="cuda", dtype=torch.float32).view(num_slots, 1, 1) + buffer = torch.empty((num_slots, 1, 132), device="cuda", dtype=torch.uint8) + buffer[:, :, :128].copy_(keys.view(torch.uint8)) + buffer[:, :, 128:132].copy_(scales.view(torch.uint8).view(num_slots, 1, 4)) + return buffer, keys.view(num_slots, 128), scales.view(num_slots) + + +def _assert_packed(output_keys, output_scales, source_keys, source_scales, expected_slots): + expected_slots = torch.tensor(expected_slots, device="cuda", dtype=torch.int64) + expected_keys = source_keys.index_select(0, expected_slots) + expected_scales = source_scales.index_select(0, expected_slots) + torch.testing.assert_close( + output_keys[: len(expected_slots)].float(), + expected_keys.float(), + rtol=0, + atol=0, + ) + torch.testing.assert_close( + output_scales[: len(expected_slots)], + expected_scales, + rtol=0, + atol=0, + ) + + +def test_dynamic_layout_matches_fixed_full_width_layout(): + buffer, source_keys, source_scales = _indexer_buffer(8) + req_to_token = torch.tensor( + [[0, 1, 2, 3], [4, 5, 6, 7]], + device="cuda", + dtype=torch.int32, + ) + b_req_idx = torch.tensor([0, 0, 0, 1, 1, 1], device="cuda", dtype=torch.int32) + b_mtp_index = torch.tensor([0, 1, 2, 0, 1, 2], device="cuda", dtype=torch.int32) + b_seq_len = torch.tensor([2, 3, 4, 1, 2, 3], device="cuda", dtype=torch.int32) + + fixed_keys, fixed_scales = extract_indexer_ks( + I_buffer=buffer, + b_seq_len=b_seq_len, + b_req_idx=b_req_idx, + req_to_token_indexs=req_to_token, + out_token_num=24, + max_kv_seq_len=4, + mtp_step=2, + ) + dynamic_keys, dynamic_scales = extract_indexer_ks_dynamic( + I_buffer=buffer, + b_seq_len=b_seq_len, + b_req_idx=b_req_idx, + b_mtp_index=b_mtp_index, + req_to_token_indexs=req_to_token, + max_kv_seq_len=4, + max_request_num=2, + ) + + torch.testing.assert_close(dynamic_keys[:7].float(), fixed_keys[:7].float(), rtol=0, atol=0) + torch.testing.assert_close(dynamic_scales[:7], fixed_scales[:7], rtol=0, atol=0) + + +def test_dynamic_layout_packs_variable_request_widths(): + buffer, source_keys, source_scales = _indexer_buffer(12) + req_to_token = torch.tensor( + [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]], + device="cuda", + dtype=torch.int32, + ) + b_req_idx = torch.tensor([0, 0, 0, 1, 2], device="cuda", dtype=torch.int32) + b_mtp_index = torch.tensor([0, 1, 2, 0, 0], device="cuda", dtype=torch.int32) + b_seq_len = torch.tensor([2, 3, 4, 3, 2], device="cuda", dtype=torch.int32) + + output_keys, output_scales = extract_indexer_ks_dynamic( + I_buffer=buffer, + b_seq_len=b_seq_len, + b_req_idx=b_req_idx, + b_mtp_index=b_mtp_index, + req_to_token_indexs=req_to_token, + max_kv_seq_len=4, + max_request_num=3, + ) + + _assert_packed( + output_keys, + output_scales, + source_keys, + source_scales, + expected_slots=[0, 1, 2, 3, 4, 5, 6, 8, 9], + ) diff --git a/test/kernel/test_glm5_grouped_topk.py b/test/kernel/test_glm5_grouped_topk.py new file mode 100644 index 0000000000..7efc318c41 --- /dev/null +++ b/test/kernel/test_glm5_grouped_topk.py @@ -0,0 +1,205 @@ +#!/usr/bin/env python3 + +import argparse + +import torch +import triton + +from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_topk import ( + single_group_sigmoid_topk_kernel, + single_group_sigmoid_topk_bitonic_kernel, + triton_grouped_topk, +) + + +def torch_reference(gating_output, correction_bias, topk): + scores = gating_output.float().sigmoid() + choice_scores = scores + correction_bias + topk_ids = torch.topk(choice_scores, k=topk, dim=-1, largest=True, sorted=True).indices + topk_weights = torch.gather(scores, 1, topk_ids) + topk_weights /= topk_weights.sum(dim=-1, keepdim=True) + return topk_weights, topk_ids + + +def run_topk(gating_output, correction_bias, *, fast): + return triton_grouped_topk( + hidden_states=None, + gating_output=gating_output, + correction_bias=correction_bias, + topk=8, + renormalize=True, + num_expert_group=1, + topk_group=1, + scoring_func="sigmoid", + group_score_used_topk_num=1, + use_single_group_fast_path=fast, + ) + + +def run_fast_with_warps(gating_output, correction_bias, num_warps): + tokens, experts = gating_output.shape + topk_weights = torch.empty((tokens, 8), dtype=torch.float32, device="cuda") + topk_ids = torch.empty((tokens, 8), dtype=torch.long, device="cuda") + single_group_sigmoid_topk_kernel[(tokens,)]( + gating_output, + gating_output.stride(0), + correction_bias, + topk_weights, + topk_weights.stride(0), + topk_ids, + topk_ids.stride(0), + experts, + HAS_CORRECTION_BIAS=True, + EXPERT_BLOCK_SIZE=triton.next_power_of_2(experts), + TOPK_BLOCK_SIZE=8, + TOPK_NUM=8, + RENORMALIZE=True, + num_warps=num_warps, + num_stages=1, + ) + return topk_weights, topk_ids + + +def run_scratch_free_bitonic(gating_output, correction_bias): + tokens, experts = gating_output.shape + topk_weights = torch.empty((tokens, 8), dtype=torch.float32, device="cuda") + topk_ids = torch.empty((tokens, 8), dtype=torch.long, device="cuda") + single_group_sigmoid_topk_bitonic_kernel[(tokens,)]( + gating_output, + gating_output.stride(0), + correction_bias, + topk_weights, + topk_weights.stride(0), + topk_ids, + topk_ids.stride(0), + experts, + HAS_CORRECTION_BIAS=True, + EXPERT_BLOCK_SIZE=triton.next_power_of_2(experts), + TOPK_NUM=8, + RENORMALIZE=True, + num_warps=4, + num_stages=1, + ) + return topk_weights, topk_ids + + +def assert_correct(tokens): + generator = torch.Generator(device="cuda").manual_seed(20260828 + tokens) + gating_output = torch.randn((tokens, 288), generator=generator, dtype=torch.float32, device="cuda") + correction_bias = torch.randn((288,), generator=generator, dtype=torch.float32, device="cuda") + + ref_weights, ref_ids = torch_reference(gating_output, correction_bias, 8) + fast_weights, fast_ids = run_topk(gating_output, correction_bias, fast=True) + generic_weights, generic_ids = run_topk(gating_output, correction_bias, fast=False) + + torch.testing.assert_close(fast_ids, ref_ids, rtol=0, atol=0) + torch.testing.assert_close(generic_ids, ref_ids, rtol=0, atol=0) + torch.testing.assert_close(fast_weights, ref_weights, rtol=1e-5, atol=1e-6) + torch.testing.assert_close(generic_weights, ref_weights, rtol=1e-5, atol=1e-6) + weight_delta = (fast_weights - generic_weights).abs() + print( + f"tokens={tokens}: exact expert ids, weights match reference; " + f"fast/generic bitwise={torch.equal(fast_weights, generic_weights)} " + f"max_delta={weight_delta.max().item():.9g}" + ) + return gating_output, correction_bias + + +def capture_graph(fn): + for _ in range(3): + outputs = fn() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outputs = fn() + return graph, outputs + + +def graph_ms(graph, iterations): + for _ in range(20): + graph.replay() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iterations): + graph.replay() + end.record() + end.synchronize() + return start.elapsed_time(end) / iterations + + +def benchmark(tokens, iterations): + gating_output, correction_bias = assert_correct(tokens) + fast_graph, fast_outputs = capture_graph(lambda: run_topk(gating_output, correction_bias, fast=True)) + generic_graph, generic_outputs = capture_graph(lambda: run_topk(gating_output, correction_bias, fast=False)) + bitonic_graph, bitonic_outputs = capture_graph(lambda: run_scratch_free_bitonic(gating_output, correction_bias)) + fast_ms = graph_ms(fast_graph, iterations) + generic_ms = graph_ms(generic_graph, iterations) + bitonic_ms = graph_ms(bitonic_graph, iterations) + torch.testing.assert_close(fast_outputs[1], generic_outputs[1], rtol=0, atol=0) + torch.testing.assert_close(bitonic_outputs[0], generic_outputs[0], rtol=0, atol=0) + torch.testing.assert_close(bitonic_outputs[1], generic_outputs[1], rtol=0, atol=0) + speedup = generic_ms / fast_ms + print( + f"tokens={tokens}: graph fast={fast_ms:.6f} ms " + f"bitonic={bitonic_ms:.6f} ms generic={generic_ms:.6f} ms " + f"speedup={speedup:.2f}x bitonic_speedup={generic_ms / bitonic_ms:.2f}x" + ) + return speedup + + +def tune_warps(tokens, iterations): + gating_output, correction_bias = assert_correct(tokens) + ref_weights, ref_ids = torch_reference(gating_output, correction_bias, 8) + generic_weights, _ = run_topk(gating_output, correction_bias, fast=False) + results = [] + for num_warps in (1, 2, 4, 8, 16): + graph, outputs = capture_graph( + lambda num_warps=num_warps: run_fast_with_warps(gating_output, correction_bias, num_warps) + ) + graph.replay() + torch.cuda.synchronize() + if not torch.equal(outputs[1], ref_ids) or not torch.allclose(outputs[0], ref_weights, rtol=1e-5, atol=1e-6): + print(f"tokens={tokens}: num_warps={num_warps} INVALID") + continue + elapsed_ms = graph_ms(graph, iterations) + results.append((elapsed_ms, num_warps)) + max_delta = (outputs[0] - generic_weights).abs().max().item() + print( + f"tokens={tokens}: num_warps={num_warps} " + f"graph={elapsed_ms:.6f} ms " + f"generic_bitwise={torch.equal(outputs[0], generic_weights)} " + f"max_delta={max_delta:.9g}" + ) + best_ms, best_warps = min(results) + print(f"tokens={tokens}: best num_warps={best_warps} graph={best_ms:.6f} ms") + + +def test_glm5_single_group_topk(): + if not torch.cuda.is_available(): + return + for tokens in (1, 8, 48, 256): + assert_correct(tokens) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--benchmark", action="store_true") + parser.add_argument("--tune-warps", action="store_true") + parser.add_argument("--iterations", type=int, default=5000) + args = parser.parse_args() + if not torch.cuda.is_available(): + raise SystemExit("CUDA is required") + if args.tune_warps: + for tokens in (1, 8, 48, 256): + tune_warps(tokens, args.iterations) + elif args.benchmark: + for tokens in (1, 8, 48, 256): + benchmark(tokens, args.iterations) + else: + test_glm5_single_group_topk() + + +if __name__ == "__main__": + main() diff --git a/test/kernel/test_glm5_kda_fusions.py b/test/kernel/test_glm5_kda_fusions.py new file mode 100644 index 0000000000..83489ecb8c --- /dev/null +++ b/test/kernel/test_glm5_kda_fusions.py @@ -0,0 +1,118 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.norm.gated_rmsnorm import ( + gated_rmsnorm_forward, +) +from lightllm.models.glm5_next.layer_weights.transformer_layer_weight import ( + Glm5NextMergedKdaProjection, +) +from lightllm.models.deepseek3_2.triton_kernel.act_quant import act_quant +from lightllm.models.deepseek3_2.triton_kernel.hadamard_transform import ( + hadamard_transform, + hadamard_transform_quant_fp8, +) +from lightllm.models.deepseek3_2.triton_kernel.indexer_weight_scale import ( + scale_indexer_weights_, +) + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +@pytest.mark.parametrize("activation", ["silu", "sigmoid"]) +def test_gated_rmsnorm_activations(activation): + torch.manual_seed(0) + x = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16) + gate = torch.randn(8, 4, 128, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(128, device="cuda", dtype=torch.bfloat16) + + actual = gated_rmsnorm_forward( + x=x, + weight=weight, + bias=None, + eps=1e-6, + z=gate, + activation=activation, + run_config={"BLOCK_N": 128, "num_warps": 4}, + ) + + x_float = x.float() + expected = x_float * torch.rsqrt(x_float.square().mean(-1, keepdim=True) + 1e-6) + expected *= weight.float() + gate_float = gate.view_as(x).float() + if activation == "silu": + expected *= gate_float * gate_float.sigmoid() + else: + expected *= gate_float.sigmoid() + + torch.testing.assert_close(actual, expected.to(actual.dtype), rtol=0, atol=0) + + +def test_merged_kda_projection_mixes_sharded_and_replicated_weights(monkeypatch): + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") + projection, head_count, head_dim, in_dim = 8, 4, 2, 3 + names = [f"weight_{index}" for index in range(6)] + merged = Glm5NextMergedKdaProjection( + in_dim=in_dim, + projection=projection, + head_count=head_count, + head_dim=head_dim, + weight_names=names, + data_type=torch.float32, + tp_rank=1, + tp_world_size=2, + ) + full_weights = { + names[0]: torch.arange(0, 24, dtype=torch.float32).view(8, 3), + names[1]: torch.arange(24, 48, dtype=torch.float32).view(8, 3), + names[2]: torch.arange(48, 72, dtype=torch.float32).view(8, 3), + names[3]: torch.arange(72, 84, dtype=torch.float32).view(4, 3), + names[4]: torch.arange(84, 90, dtype=torch.float32).view(2, 3), + names[5]: torch.arange(90, 96, dtype=torch.float32).view(2, 3), + } + merged.load_hf_weights(full_weights) + + expected_weight = torch.cat( + [ + full_weights[names[0]][4:], + full_weights[names[1]][4:], + full_weights[names[2]][4:], + full_weights[names[3]][2:], + full_weights[names[4]], + full_weights[names[5]], + ] + ).cuda() + torch.testing.assert_close(merged.mm_param.weight, expected_weight) + + x = torch.arange(6, dtype=torch.float32, device="cuda").view(2, 3) + torch.testing.assert_close(merged.mm(x, use_custom_tensor_mananger=False), x @ expected_weight.T) + assert merged.verify_load() + + +@pytest.mark.parametrize("shape", [(17, 128), (3, 32, 128)]) +def test_fused_hadamard_fp8_quant_matches_two_kernel_chain(shape): + torch.manual_seed(1) + value = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + expected_value, expected_scale = act_quant( + hadamard_transform(value, scale=128 ** -0.5), + block_size=128, + scale_fmt="ue8m0", + ) + + actual_value, actual_scale = hadamard_transform_quant_fp8(value, scale=128 ** -0.5) + + assert torch.equal(actual_value, expected_value) + torch.testing.assert_close(actual_scale, expected_scale, rtol=0, atol=0) + + +def test_fused_indexer_weight_scale_matches_torch_chain(): + torch.manual_seed(2) + weights = torch.randn(257, 32, device="cuda", dtype=torch.float32) + q_scale = torch.rand(257, 32, 1, device="cuda", dtype=torch.float32) + scale = 128 ** -0.5 * 32 ** -0.5 + expected = (weights * scale).unsqueeze(-1).mul(q_scale).squeeze(-1) + + actual = scale_indexer_weights_(weights.clone(), q_scale, scale) + + torch.testing.assert_close(actual, expected, rtol=1e-6, atol=1e-7) diff --git a/test/kernel/test_glm5_mhc.py b/test/kernel/test_glm5_mhc.py new file mode 100644 index 0000000000..4cb7a1f504 --- /dev/null +++ b/test/kernel/test_glm5_mhc.py @@ -0,0 +1,250 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Numerical and launch-overhead check for the fused GLM-5 mHC kernels.""" + +import argparse + +import torch + +from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward +from lightllm.models.glm5_next.triton_kernel.mhc import ( + hc_post, + hc_post_reference, + hc_pre_norm, + hc_pre_reference, +) + + +def _max_error(actual: torch.Tensor, expected: torch.Tensor) -> float: + return (actual.float() - expected.float()).abs().max().item() + + +def _time_ms(function, iterations: int = 100) -> float: + for _ in range(5): + function() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iterations): + function() + end.record() + end.synchronize() + return start.elapsed_time(end) / iterations + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--large-post", action="store_true") + parser.add_argument("--iterations", type=int, default=100) + args = parser.parse_args() + + torch.manual_seed(1234) + device = torch.device("cuda") + streams = 4 + hidden = 4096 + for tokens in (1, 8, 48): + x = torch.randn((tokens, streams * hidden), device=device, dtype=torch.bfloat16) + fn = 0.005 * torch.randn( + ((2 + streams) * streams, streams * hidden), + device=device, + dtype=torch.float32, + ) + scale = torch.randn((3,), device=device, dtype=torch.float32) + base = torch.randn(((2 + streams) * streams,), device=device, dtype=torch.float32) + layer_output = torch.randn((tokens, hidden), device=device, dtype=torch.bfloat16) + norm_weight = torch.randn((hidden,), device=device, dtype=torch.bfloat16) + arguments = (x, fn, scale, base, streams, 1e-6, 1e-6, 20) + + expected_pre = hc_pre_reference(*arguments) + expected_pre = ( + rmsnorm_forward(expected_pre[0], weight=norm_weight, eps=1e-6), + expected_pre[1], + expected_pre[2], + ) + actual_pre = hc_pre_norm( + x, + fn, + scale, + base, + norm_weight, + streams, + 1e-6, + 1e-6, + 1e-6, + 20, + ) + torch.cuda.synchronize() + pre_errors = tuple(_max_error(actual, expected) for actual, expected in zip(actual_pre, expected_pre)) + assert pre_errors[0] <= 0.03125, pre_errors + # DeepGEMM intentionally uses TF32 and a split-K reduction, matching + # the optimized serving path rather than torch.mm's accumulation order. + assert pre_errors[1] <= 5e-4, pre_errors + assert pre_errors[2] <= 5e-4, pre_errors + + expected_post = hc_post_reference(layer_output, x, expected_pre[1], expected_pre[2], streams) + actual_post = hc_post(layer_output, x, actual_pre[1], actual_pre[2], streams) + torch.cuda.synchronize() + post_error = _max_error(actual_post, expected_post) + assert post_error <= 0.03125, post_error + + def reference_path(): + layer_input, residual_mix, post_mix = hc_pre_reference(*arguments) + rmsnorm_forward(layer_input, weight=norm_weight, eps=1e-6) + return hc_post_reference(layer_output, x, residual_mix, post_mix, streams) + + def fused_path(): + _, residual_mix, post_mix = hc_pre_norm( + x, + fn, + scale, + base, + norm_weight, + streams, + 1e-6, + 1e-6, + 1e-6, + 20, + ) + return hc_post(layer_output, x, residual_mix, post_mix, streams) + + def fused_pre_path(): + return hc_pre_norm( + x, + fn, + scale, + base, + norm_weight, + streams, + 1e-6, + 1e-6, + 1e-6, + 20, + ) + + def fused_post_path(): + return hc_post(layer_output, x, actual_pre[1], actual_pre[2], streams) + + reference_ms = _time_ms(reference_path) + fused_ms = _time_ms(fused_path) + fused_pre_ms = _time_ms(fused_pre_path) + fused_post_ms = _time_ms(fused_post_path) + print( + f"PASS tokens={tokens} pre_errors={pre_errors} " + f"post_error={post_error:.8f} reference_ms={reference_ms:.4f} " + f"fused_ms={fused_ms:.4f} pre_ms={fused_pre_ms:.4f} " + f"post_ms={fused_post_ms:.4f} " + f"speedup={reference_ms / fused_ms:.2f}x" + ) + + if args.large_post: + from sglang.kernels.ops.layernorm.mhc import mhc_post_tilelang + + tokens = 17152 + residual = torch.randn((tokens, streams * hidden), device=device, dtype=torch.bfloat16) + layer_output = torch.randn((tokens, hidden), device=device, dtype=torch.bfloat16) + residual_mix = torch.randn((tokens, streams, streams), device=device, dtype=torch.float32) + post_mix = torch.randn((tokens, streams), device=device, dtype=torch.float32) + actual = hc_post(layer_output, residual, residual_mix, post_mix, streams) + + def sgl_mhc_post(): + output = torch.empty_like(residual.view(tokens, streams, hidden)) + mhc_post_tilelang( + residual_mix, + residual.view(tokens, streams, hidden), + post_mix, + layer_output, + output, + streams, + hidden, + ) + return output + + sglang_output = sgl_mhc_post().view(tokens, -1) + torch.cuda.synchronize() + cross_error = _max_error(actual, sglang_output) + assert cross_error <= 0.0625, cross_error + triton_ms = _time_ms( + lambda: hc_post(layer_output, residual, residual_mix, post_mix, streams), + args.iterations, + ) + tilelang_ms = _time_ms( + sgl_mhc_post, + args.iterations, + ) + print( + f"PASS large_post tokens={tokens} cross_error={cross_error:.8f} " + f"triton_ms={triton_ms:.4f} tilelang_ms={tilelang_ms:.4f} " + f"triton_over_tilelang={tilelang_ms / triton_ms:.2f}x" + ) + + # Compare the full DeepGEMM prenorm + mHC-pre fusion used by both + # runtimes. The standalone kernel test has no SGLang process group, + # so disable only its symmetric-allocation context. + import contextlib + import sglang.kernels.ops.layernorm.mhc as sgl_mhc + + sgl_mhc.get_tp_group = lambda: None + sgl_mhc.is_allocation_symmetric = lambda: False + sgl_mhc.use_symmetric_memory = lambda *_args, **_kwargs: contextlib.nullcontext() + fn = 0.005 * torch.randn( + ((2 + streams) * streams, streams * hidden), + device=device, + dtype=torch.float32, + ) + scale = torch.randn((3,), device=device, dtype=torch.float32) + base = torch.randn(((2 + streams) * streams,), device=device, dtype=torch.float32) + norm_weight = torch.randn((hidden,), device=device, dtype=torch.bfloat16) + + def lightllm_pre(): + return hc_pre_norm( + residual, + fn, + scale, + base, + norm_weight, + streams, + 1e-6, + 1e-6, + 1e-6, + 20, + ) + + def sglang_pre(): + return sgl_mhc.mhc_pre( + residual.view(tokens, streams, hidden), + fn, + scale, + base, + 1e-6, + 1e-6, + 1e-6, + 2.0, + 20, + norm_weight=norm_weight, + norm_eps=1e-6, + ) + + lightllm_result = lightllm_pre() + sglang_result = sglang_pre() + torch.cuda.synchronize() + pre_cross_errors = ( + _max_error(lightllm_result[0], sglang_result[2]), + _max_error(lightllm_result[1], sglang_result[1]), + _max_error(lightllm_result[2], sglang_result[0].squeeze(-1)), + ) + assert pre_cross_errors[0] <= 0.0625, pre_cross_errors + assert pre_cross_errors[1] <= 5e-4, pre_cross_errors + assert pre_cross_errors[2] <= 5e-4, pre_cross_errors + lightllm_pre_ms = _time_ms(lightllm_pre, args.iterations) + sglang_pre_ms = _time_ms(sglang_pre, args.iterations) + print( + f"PASS large_pre tokens={tokens} errors={pre_cross_errors} " + f"lightllm_ms={lightllm_pre_ms:.4f} " + f"sglang_ms={sglang_pre_ms:.4f} " + f"sglang_speedup={lightllm_pre_ms / sglang_pre_ms:.2f}x" + ) + + +if __name__ == "__main__": + main() diff --git a/test/kernel/test_glm5_sglang_moe_compat.py b/test/kernel/test_glm5_sglang_moe_compat.py new file mode 100644 index 0000000000..7f9a7ce217 --- /dev/null +++ b/test/kernel/test_glm5_sglang_moe_compat.py @@ -0,0 +1,238 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Numerically compare LightLLM and SGLang's GLM-5 FP8 MoE paths. + +This is an integration probe for development images that provide both +packages. It intentionally uses GLM-5's TP8 decode shapes, including the +fused shared expert, block-wise FP8 scales, and clamped SwiGLU. +""" + +import argparse +import itertools +import json +from types import SimpleNamespace + +import torch + +from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe import ( + fused_experts as lightllm_fused_experts, +) +from sglang.srt.layers.moe.moe_runner.triton_utils import fused_moe as sglang_fused_moe +from sglang.srt.layers.moe.moe_runner.triton_utils import ( + fused_moe_triton_config as sglang_fused_moe_config, +) +from sglang.srt.layers.moe.moe_runner.triton_utils import override_config + + +def _fp8_randn(shape, *, scale=0.02): + return (torch.randn(shape, device="cuda", dtype=torch.bfloat16) * scale).to(torch.float8_e4m3fn) + + +def _graph_ms(fn, source, iterations): + static_input = source.clone() + for _ in range(3): + static_input.copy_(source) + fn(static_input) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + static_input.copy_(source) + fn(static_input) + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iterations): + graph.replay() + end.record() + end.synchronize() + elapsed_ms = start.elapsed_time(end) / iterations + graph.reset() + return elapsed_ms + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--benchmark", action="store_true") + parser.add_argument( + "--fuse-sum", + action="store_true", + help="Fuse the top-k sum into SGLang's down-projection kernel.", + ) + parser.add_argument( + "--tune-configs", + action="store_true", + help="Search a compact H100 config set for GLM-5's TP8 decode shape.", + ) + parser.add_argument( + "--tune-tma-configs", + action="store_true", + help="Search the compact small-M set with SGLang's up-projection TMA path.", + ) + parser.add_argument( + "--num-tokens", + type=int, + default=48, + help="Physical token count used by the MoE probe (48 main, 8 MTP draft).", + ) + parser.add_argument("--iterations", type=int, default=100) + args = parser.parse_args() + + torch.manual_seed(123) + + num_tokens = args.num_tokens + num_experts = 289 + topk = 9 + hidden_size = 4096 + tp_intermediate_size = 256 + + hidden_states = torch.randn((num_tokens, hidden_size), device="cuda", dtype=torch.bfloat16) + w13 = _fp8_randn((num_experts, 2 * tp_intermediate_size, hidden_size)) + w2 = _fp8_randn((num_experts, hidden_size, tp_intermediate_size)) + w13_scale = torch.ones( + (num_experts, 2 * tp_intermediate_size // 128, hidden_size // 128), + device="cuda", + dtype=torch.float32, + ) + w2_scale = torch.ones( + (num_experts, hidden_size // 128, tp_intermediate_size // 128), + device="cuda", + dtype=torch.float32, + ) + topk_ids = torch.randint(0, num_experts, (num_tokens, topk), device="cuda", dtype=torch.int64) + topk_weights = torch.rand((num_tokens, topk), device="cuda", dtype=torch.float32) + topk_weights.mul_(2.5 / topk_weights.sum(dim=-1, keepdim=True)) + + def run_lightllm(output): + lightllm_fused_experts( + hidden_states=output, + w1=w13, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + inplace=True, + use_fp8_w8a8=True, + w1_scale=w13_scale, + w2_scale=w2_scale, + limit=10.0, + alpha=1.0, + clamp_up_add_one=False, + ) + + lightllm_output = hidden_states.clone() + run_lightllm(lightllm_output) + + # SGLang's fused sequence consults this runtime flag. A standalone + # LightLLM process has no SGLang RuntimeContext, so provide the selected + # fused top-k sum behavior explicitly for this comparison. + standalone_exec = SimpleNamespace( + moe=SimpleNamespace(enable_fused_moe_sum_all_reduce=args.fuse_sum), + deterministic=SimpleNamespace(enable_deterministic_inference=False), + ) + sglang_fused_moe.get_exec = lambda: standalone_exec + sglang_fused_moe_config.get_exec = lambda: standalone_exec + + def run_sglang(output): + sglang_fused_moe.fused_experts_impl( + hidden_states=output, + w1=w13, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + inplace=True, + use_fp8_w8a8=True, + w1_scale=w13_scale, + w2_scale=w2_scale, + block_shape=[128, 128], + routed_scaling_factor=1.0, + filter_expert=False, + swiglu_limit=10.0, + gate_up_interleaved=False, + ) + + sglang_output = hidden_states.clone() + run_sglang(sglang_output) + torch.cuda.synchronize() + + diff = (lightllm_output.float() - sglang_output.float()).abs() + reference = lightllm_output.float().abs() + print( + "max_abs=%.6f mean_abs=%.6f max_ref=%.6f mean_ref=%.6f" + % ( + diff.max().item(), + diff.mean().item(), + reference.max().item(), + reference.mean().item(), + ) + ) + torch.testing.assert_close( + sglang_output.float(), + lightllm_output.float(), + rtol=0.08, + atol=0.08, + ) + print("GLM-5 SGLang MoE compatibility: PASS") + + if args.benchmark: + lightllm_ms = _graph_ms(run_lightllm, hidden_states, args.iterations) + sglang_ms = _graph_ms(run_sglang, hidden_states, args.iterations) + print("graph_ms lightllm=%.6f sglang=%.6f speedup=%.3fx" % (lightllm_ms, sglang_ms, lightllm_ms / sglang_ms)) + + if args.tune_configs or args.tune_tma_configs: + # GLM-5 decode has few physical tokens (48 for the main model and 8 + # for each draft) spread over 289 experts. SGLang's generic block-FP8 + # fallback uses BLOCK_SIZE_M=64, which can pad this sparse workload + # excessively. Search the small-M region instead of the generic + # 1920-config tuning space. + search_space = itertools.product( + (16, 32), + (64, 128), + (128,) if args.tune_tma_configs else (64, 128), + (1, 16), + (4, 8), + (2, 3), + ) + candidates = [ + { + "BLOCK_SIZE_M": block_m, + "BLOCK_SIZE_N": block_n, + "BLOCK_SIZE_K": block_k, + "GROUP_SIZE_M": group_m, + "num_warps": num_warps, + "num_stages": num_stages, + **({"USE_TMA": True} if args.tune_tma_configs else {}), + } + for block_m, block_n, block_k, group_m, num_warps, num_stages in search_space + ] + default_config = { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3, + } + candidates.append(default_config) + + results = [] + for config in candidates: + try: + with override_config(config): + graph_ms = _graph_ms(run_sglang, hidden_states, args.iterations) + except Exception as exc: + print("config_failed=%s error=%r" % (json.dumps(config, sort_keys=True), exc)) + continue + results.append((graph_ms, config)) + print("config_ms=%.6f config=%s" % (graph_ms, json.dumps(config, sort_keys=True))) + + results.sort(key=lambda item: item[0]) + if not results: + raise RuntimeError("all SGLang MoE tuning candidates failed") + print("top_configs=%s" % json.dumps(results[:10], sort_keys=True)) + print("best_config=%s" % json.dumps(results[0][1], sort_keys=True)) + print("best_graph_ms=%.6f" % results[0][0]) + + +if __name__ == "__main__": + main() diff --git a/test/kernel/test_glm5_short_decode_helpers.py b/test/kernel/test_glm5_short_decode_helpers.py new file mode 100644 index 0000000000..4b276a771b --- /dev/null +++ b/test/kernel/test_glm5_short_decode_helpers.py @@ -0,0 +1,23 @@ +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from lightllm.models.glm5_next_mtp.triton_kernel.zero_position_embedding import ( + zero_position_embedding_, +) + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def test_zero_position_embedding_only_changes_position_zero_rows(): + embeddings = torch.arange(15, device="cuda", dtype=torch.bfloat16).view(3, 5) + original = embeddings.clone() + position_ids = torch.tensor([0, 1, 0], device="cuda", dtype=torch.int64) + + zero_position_embedding_(embeddings, position_ids) + + torch.testing.assert_close(embeddings[0], torch.zeros_like(embeddings[0]), rtol=0, atol=0) + torch.testing.assert_close(embeddings[1], original[1], rtol=0, atol=0) + torch.testing.assert_close(embeddings[2], torch.zeros_like(embeddings[2]), rtol=0, atol=0) diff --git a/test/kernel/test_glm5_strided_causal_conv.py b/test/kernel/test_glm5_strided_causal_conv.py new file mode 100644 index 0000000000..0f1428a47d --- /dev/null +++ b/test/kernel/test_glm5_strided_causal_conv.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Check GLM-5 KDA's copy-free strided causal convolution path.""" + +import argparse + +import torch + +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( + causal_conv1d_fn, +) + + +def _inputs(seq_lens: list[int], dim: int, cache_lines: int): + total_tokens = sum(seq_lens) + token_major = torch.randn((total_tokens, dim), device="cuda", dtype=torch.bfloat16) + weight = torch.randn((dim, 4), device="cuda", dtype=torch.bfloat16) + conv_states = torch.randn((cache_lines, dim, 3), device="cuda", dtype=torch.bfloat16) + cache_indices = torch.arange(len(seq_lens), device="cuda", dtype=torch.int32) + has_initial_state = torch.tensor( + [(index % 2) == 1 for index in range(len(seq_lens))], + device="cuda", + dtype=torch.bool, + ) + query_start_loc = torch.tensor( + [0, *torch.tensor(seq_lens).cumsum(0).tolist()], + device="cuda", + dtype=torch.int32, + ) + return ( + token_major, + weight, + conv_states, + cache_indices, + has_initial_state, + query_start_loc, + ) + + +def _run(arguments, seq_lens: list[int], *, copy_free: bool): + ( + token_major, + weight, + conv_states, + cache_indices, + has_initial_state, + query_start_loc, + ) = arguments + return causal_conv1d_fn( + token_major.transpose(0, 1), + weight, + query_start_loc=query_start_loc, + cache_indices=cache_indices, + has_initial_state=has_initial_state, + conv_states=conv_states, + activation="silu", + seq_lens_cpu=seq_lens if copy_free else None, + ) + + +def _time_ms(function, iterations: int) -> float: + for _ in range(3): + function() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iterations): + function() + end.record() + end.synchronize() + return start.elapsed_time(end) / iterations + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--benchmark", action="store_true") + parser.add_argument("--iterations", type=int, default=20) + args = parser.parse_args() + + torch.manual_seed(1234) + seq_lens = [7, 1, 19, 5] + inputs = _inputs(seq_lens, dim=64, cache_lines=8) + reference_inputs = tuple(value.clone() if isinstance(value, torch.Tensor) else value for value in inputs) + actual = _run(inputs, seq_lens, copy_free=True) + expected = _run(reference_inputs, seq_lens, copy_free=False) + torch.cuda.synchronize() + output_error = (actual.float() - expected.float()).abs().max().item() + state_error = (inputs[2].float() - reference_inputs[2].float()).abs().max().item() + # The two kernels accumulate the four taps in a different order; one BF16 + # ULP at this random input scale is expected. + assert output_error <= 0.0625, output_error + assert state_error == 0.0, state_error + assert actual.transpose(0, 1).is_contiguous() + print(f"PASS correctness output_error={output_error:.8f} " f"state_error={state_error:.8f}") + + if args.benchmark: + seq_lens = [268] * 64 + strided_inputs = _inputs(seq_lens, dim=3072, cache_lines=64) + copied_inputs = tuple(value.clone() if isinstance(value, torch.Tensor) else value for value in strided_inputs) + strided_ms = _time_ms( + lambda: _run(strided_inputs, seq_lens, copy_free=True), + args.iterations, + ) + copied_ms = _time_ms( + lambda: _run(copied_inputs, seq_lens, copy_free=False), + args.iterations, + ) + print( + f"PASS benchmark tokens={sum(seq_lens)} dim=3072 " + f"strided_ms={strided_ms:.4f} copied_ms={copied_ms:.4f} " + f"speedup={copied_ms / strided_ms:.2f}x" + ) + + +if __name__ == "__main__": + main() diff --git a/test/kernel/test_glm5_vocab_parallel_top1.py b/test/kernel/test_glm5_vocab_parallel_top1.py new file mode 100644 index 0000000000..b4522f5ebc --- /dev/null +++ b/test/kernel/test_glm5_vocab_parallel_top1.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: Apache-2.0 + +import os + +import pytest +import torch +import torch.distributed as dist + +from lightllm.models.glm5_next_mtp.layer_infer.post_layer_infer import ( + vocab_parallel_top1, + vocab_parallel_top1_and_prob, +) + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def test_vocab_parallel_top1_matches_full_vocab_softmax(): + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + torch.cuda.set_device(local_rank) + initialized_here = not dist.is_initialized() + if initialized_here: + dist.init_process_group(backend="nccl", device_id=torch.device("cuda", local_rank)) + + try: + rank = dist.get_rank() + world_size = dist.get_world_size() + assert world_size >= 2 + + torch.manual_seed(20260828) + token_num = 8 + vocab_size = 154880 + assert vocab_size % world_size == 0 + full_logits = torch.randn( + (vocab_size, token_num), + dtype=torch.bfloat16, + device="cuda", + ) + # Exercise rank-boundary winners and deterministic ties. Global + # argmax must choose the lower vocabulary id just like torch.argmax. + shard_size = vocab_size // world_size + full_logits[shard_size - 1, 0] = 20 + full_logits[shard_size, 0] = 20 + full_logits[-1, 1] = 21 + + local_start = rank * shard_size + local_logits = full_logits[local_start : local_start + shard_size].contiguous() + + def alloc_func(shape, dtype, device): + return torch.empty(shape, dtype=dtype, device=device) + + token_ids, token_probs = vocab_parallel_top1_and_prob( + local_logits=local_logits, + local_vocab_start_id=local_start, + tp_world_size=world_size, + dist_group=dist.group.WORLD, + alloc_func=alloc_func, + ) + + expected_probs, expected_ids = torch.softmax(full_logits.float().t(), dim=-1).max(dim=-1) + torch.testing.assert_close(token_ids, expected_ids, rtol=0, atol=0) + torch.testing.assert_close(token_probs, expected_probs, rtol=2e-5, atol=2e-7) + torch.testing.assert_close( + vocab_parallel_top1( + local_logits, + local_start, + world_size, + dist.group.WORLD, + alloc_func, + ), + expected_ids, + rtol=0, + atol=0, + ) + + if os.environ.get("LIGHTLLM_TEST_SKIP_CUDA_GRAPH") == "1": + return + + # The draft head normally runs inside a decode CUDA Graph. Replay + # with changed logits to ensure the tiny collective and returned + # tensors are captured as live outputs rather than stale warmup data. + graph_logits = local_logits.clone() + for _ in range(2): + vocab_parallel_top1_and_prob( + graph_logits, + local_start, + world_size, + dist.group.WORLD, + alloc_func, + ) + dist.barrier() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_ids, graph_probs = vocab_parallel_top1_and_prob( + graph_logits, + local_start, + world_size, + dist.group.WORLD, + alloc_func, + ) + + torch.manual_seed(20260829) + replay_full_logits = torch.randn_like(full_logits) + replay_full_logits[shard_size * 2 - 1, 2] = 22 + graph_logits.copy_(replay_full_logits[local_start : local_start + shard_size]) + graph.replay() + torch.cuda.synchronize() + expected_probs, expected_ids = torch.softmax(replay_full_logits.float().t(), dim=-1).max(dim=-1) + torch.testing.assert_close(graph_ids, expected_ids, rtol=0, atol=0) + torch.testing.assert_close(graph_probs, expected_probs, rtol=2e-5, atol=2e-7) + finally: + if initialized_here: + dist.destroy_process_group() diff --git a/test/test_glm5_next_multimodal.py b/test/test_glm5_next_multimodal.py new file mode 100644 index 0000000000..7521b2d9d0 --- /dev/null +++ b/test/test_glm5_next_multimodal.py @@ -0,0 +1,96 @@ +from transformers.configuration_utils import PretrainedConfig + +from lightllm.models.glm5_next.model import ( + Glm5NextMultimodalTpPartModel, + Glm5NextTpPartModel, +) +from lightllm.models.glm5_next.tokenizer import Glm5NextTokenizer +from lightllm.models.registry import get_model_class +from lightllm.server.multimodal_params import MultimodalParams +from lightllm.utils.config_utils import has_vision_module + + +class _FakeImageProcessor: + patch_size = 14 + merge_size = 2 + min_image_tokens = 16 + max_image_tokens = 8000 + + @staticmethod + def get_number_of_image_patches(height, width): + assert (height, width) == (448, 448) + return 1024 + + +class _FakeTokenizer: + def __init__(self): + self.last_conversation = None + + def apply_chat_template(self, conversation, **kwargs): + self.last_conversation = conversation + return conversation[0]["content"][0]["text"] + + @staticmethod + def encode(prompt): + assert prompt == Glm5NextTokenizer.image_placeholder + return [7, 10, 12, 11, 8] + + +def _make_tokenizer(): + model_cfg = { + "image_start_token_id": 10, + "image_end_token_id": 11, + "image_token_id": 12, + } + return Glm5NextTokenizer( + tokenizer=_FakeTokenizer(), + image_processor=_FakeImageProcessor(), + model_cfg=model_cfg, + ) + + +def test_glm5_image_prompt_and_virtual_tokens(): + tokenizer = _make_tokenizer() + conversation = [ + { + "role": "user", + "content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,AA=="}}], + } + ] + + prompt = tokenizer.apply_chat_template(conversation=conversation) + assert prompt == tokenizer.image_placeholder + assert conversation[0]["content"][0]["type"] == "image_url" + + multimodal_params = MultimodalParams(images=[{"type": "base64", "data": ""}]) + image = multimodal_params.images[0] + image.image_w = image.image_h = 448 + image.token_num = tokenizer.get_image_token_length(image) + image.token_id = 1000 + + input_ids = tokenizer.encode(prompt, multimodal_params=multimodal_params) + assert image.token_num == 256 + assert image.start_idx == 2 + assert input_ids == [7, 10, *range(1000, 1256), 11, 8] + + +def test_glm5_vision_config_selects_multimodal_model(monkeypatch): + assert get_model_class({"model_type": "glm5_next"}) is Glm5NextTpPartModel + assert ( + get_model_class({"model_type": "glm5_next", "vision_config": {"hidden_size": 1024}}) + is Glm5NextMultimodalTpPartModel + ) + + monkeypatch.setattr( + PretrainedConfig, + "get_config_dict", + staticmethod( + lambda _: ( + {"model_type": "glm5_next", "vision_config": {"hidden_size": 1024}}, + {}, + ) + ), + ) + has_vision_module.cache_clear() + assert has_vision_module("unused") is True + has_vision_module.cache_clear() diff --git a/test/test_moe_prefill_dispatch.py b/test/test_moe_prefill_dispatch.py new file mode 100644 index 0000000000..eabac2d6c0 --- /dev/null +++ b/test/test_moe_prefill_dispatch.py @@ -0,0 +1,45 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.deepseek2.layer_infer.transformer_layer_infer import ( + Deepseek2TransformerLayerInfer, +) + + +@pytest.mark.parametrize("is_prefill", [False, True]) +def test_tp_moe_propagates_prefill_stage(is_prefill): + captured = {} + + class Gate: + data_type_ = torch.float32 + + @staticmethod + def mm(hidden_states): + return torch.zeros((hidden_states.shape[0], 4)) + + class Experts: + @staticmethod + def experts(*args, **kwargs): + captured.update(kwargs) + + layer = Deepseek2TransformerLayerInfer.__new__(Deepseek2TransformerLayerInfer) + layer.embed_dim_ = 4 + layer.n_shared_experts = None + layer.num_experts_per_tok = 2 + layer.norm_topk_prob = True + layer.n_group = 1 + layer.topk_group = 1 + layer_weight = SimpleNamespace( + moe_gate=Gate(), + experts=Experts(), + num_fused_shared_experts=0, + ) + infer_state = SimpleNamespace(is_prefill=is_prefill) + + output = layer._moe_ffn_tp(torch.zeros((3, 4)), infer_state, layer_weight) + + assert output.shape == (3, 4) + assert captured["is_prefill"] is is_prefill + assert captured["infer_state"] is infer_state diff --git a/test/test_pd_selector/test_pd_master_metrics.py b/test/test_pd_selector/test_pd_master_metrics.py new file mode 100644 index 0000000000..0243565cd5 --- /dev/null +++ b/test/test_pd_selector/test_pd_master_metrics.py @@ -0,0 +1,124 @@ +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, call + +import pytest +from prometheus_client import generate_latest + +from lightllm.server.httpserver_for_pd_master.manager import HttpServerManagerForPDMaster, PDManager +from lightllm.server.metrics.manager import MetricServer +from lightllm.server.metrics.metrics import Monitor +from lightllm.server.pd_io_struct import PD_Client_Obj +from lightllm.utils.error_utils import ClientDisconnected, ServerBusyError + + +class RPyCDictProxyLike: + def __init__(self, values): + self.values = values + + def __iter__(self): + return iter(self.values) + + def __getitem__(self, key): + return self.values[key] + + +def test_pd_master_exports_node_load_with_role_and_endpoint_labels(): + args = SimpleNamespace( + metric_gateway=None, + job_name="lightllm", + grouping_key=None, + enable_monitor_auth=False, + model_name="test-model", + max_req_total_len=1024, + mtp_step=0, + select_p_d_node_strategy="random", + ) + monitor = Monitor(args) + metric_client = MagicMock() + metric_client.gauge_set.side_effect = monitor.gauge_set + manager = PDManager(args, metric_client) + manager.url_to_pd_nodes["10.0.0.1:28761"] = PD_Client_Obj( + node_id=1, + client_ip_port="10.0.0.1:28761", + mode="prefill", + start_args={}, + ) + manager.url_to_pd_nodes["10.0.0.2:28764"] = PD_Client_Obj( + node_id=2, + client_ip_port="10.0.0.2:28764", + mode="decode", + start_args={}, + ) + + manager.update_node_load_info({"client_ip_port": "10.0.0.1:28761", "total_token_usage_rate": 0.25}) + manager.update_node_load_info({"client_ip_port": "10.0.0.2:28764", "total_token_usage_rate": 0.75}) + monitor.gauge_set("lightllm_pd_master_stage_waiting_requests", 2, labels={"stage": "decode"}) + + metrics = generate_latest(monitor.registry).decode() + assert ( + 'lightllm_pd_node_token_usage_ratio{endpoint="10.0.0.1:28761",model_name="test-model",role="prefill"} 0.25' + in metrics + ) + assert ( + 'lightllm_pd_node_token_usage_ratio{endpoint="10.0.0.2:28764",model_name="test-model",role="decode"} 0.75' + in metrics + ) + assert 'lightllm_pd_master_stage_waiting_requests{model_name="test-model",stage="decode"} 2.0' in metrics + + +def test_metric_server_copies_rpyc_label_proxy_before_updating_gauge(): + args = SimpleNamespace( + metric_gateway=None, + job_name="lightllm", + grouping_key=None, + enable_monitor_auth=False, + model_name="test-model", + max_req_total_len=1024, + mtp_step=0, + push_interval=10, + ) + server = MetricServer(args) + + server.exposed_gauge_set( + "lightllm_pd_master_stage_waiting_requests", + 3, + RPyCDictProxyLike({"stage": "decode"}), + ) + + metrics = generate_latest(server.monitor.registry).decode() + assert 'lightllm_pd_master_stage_waiting_requests{model_name="test-model",stage="decode"} 3.0' in metrics + + +@pytest.mark.parametrize("stage", ["prefill", "decode"]) +@pytest.mark.parametrize( + "outcome", + [ + None, + ServerBusyError(), + asyncio.CancelledError(), + ClientDisconnected(group_request_id=123), + RuntimeError("worker failed"), + ], + ids=["success", "timeout", "cancellation", "disconnect", "exception"], +) +def test_pd_master_stage_waiting_gauge_is_balanced_for_every_exit(stage, outcome): + async def run(): + manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) + manager.metric_client = MagicMock() + manager.pd_stage_waiting_request_counts = {"prefill": 0, "decode": 0} + manager._wait_for_event_or_disconnect = AsyncMock(side_effect=outcome) + + if outcome is None: + await manager._wait_for_pd_stage(AsyncMock(), AsyncMock(), 1, 123, stage) + else: + with pytest.raises(type(outcome)): + await manager._wait_for_pd_stage(AsyncMock(), AsyncMock(), 1, 123, stage) + + assert manager.pd_stage_waiting_request_counts[stage] == 0 + assert manager.metric_client.gauge_set.call_args_list == [ + call("lightllm_pd_master_stage_waiting_requests", 1, labels={"stage": stage}), + call("lightllm_pd_master_stage_waiting_requests", 0, labels={"stage": stage}), + ] + + asyncio.run(run()) diff --git a/tools/analyze_torch_trace.py b/tools/analyze_torch_trace.py new file mode 100644 index 0000000000..0da113e23b --- /dev/null +++ b/tools/analyze_torch_trace.py @@ -0,0 +1,65 @@ +#!/usr/bin/env python3 + +"""Aggregate CUDA kernel time from one or more Torch profiler traces.""" + +from __future__ import annotations + +import argparse +from collections import defaultdict +import gzip +import json +from pathlib import Path + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("traces", nargs="+", type=Path) + parser.add_argument("--graph-only", action="store_true") + parser.add_argument("--graph-id", type=int) + parser.add_argument("--top", type=int, default=40) + args = parser.parse_args() + + totals: dict[str, list[float]] = defaultdict(lambda: [0.0, 0.0, 0.0]) + graph_totals: dict[int, list[float]] = defaultdict(lambda: [0.0, 0.0]) + for path in args.traces: + opener = gzip.open if path.suffix == ".gz" else open + with opener(path, "rt", encoding="utf-8") as trace_file: + events = json.load(trace_file)["traceEvents"] + for event in events: + if event.get("cat") != "kernel" or "dur" not in event: + continue + graph_id = int(event.get("args", {}).get("graph id", 0)) + if args.graph_only and graph_id == 0: + continue + if args.graph_id is not None and graph_id != args.graph_id: + continue + duration = float(event["dur"]) + name = event.get("name", "") + totals[name][0] += duration + totals[name][1] += 1 + totals[name][2] = max(totals[name][2], duration) + graph_totals[graph_id][0] += duration + graph_totals[graph_id][1] += 1 + + total_us = sum(values[0] for values in totals.values()) + total_calls = int(sum(values[1] for values in totals.values())) + print(f"CUDA total: {total_us / 1000:.3f} ms across {total_calls} kernels") + print("Graph totals:") + for graph_id, (duration, count) in sorted(graph_totals.items(), key=lambda item: item[1][0], reverse=True): + print( + f" graph={graph_id:<5} total_ms={duration / 1000:10.3f} " + f"share={duration / total_us:7.2%} calls={int(count)}" + ) + print("Kernel totals:") + for name, (duration, count, maximum) in sorted(totals.items(), key=lambda item: item[1][0], reverse=True)[ + : args.top + ]: + print( + f" total_ms={duration / 1000:10.3f} share={duration / total_us:7.2%} " + f"calls={int(count):7d} avg_us={duration / count:9.3f} " + f"max_us={maximum:9.3f} {name}" + ) + + +if __name__ == "__main__": + main() diff --git a/tools/bench_glm53_allreduce.py b/tools/bench_glm53_allreduce.py new file mode 100644 index 0000000000..de377f2319 --- /dev/null +++ b/tools/bench_glm53_allreduce.py @@ -0,0 +1,171 @@ +#!/usr/bin/env python3 +"""Compare NCCL and symmetric-memory all-reduce for GLM-5.3 TP8 tensors.""" + +import argparse +import os + +import torch +import torch.distributed as dist +import torch.distributed._symmetric_memory as symm_mem +import flashinfer.comm as flashinfer_comm +from flashinfer.comm.mnnvl import TorchDistBackend + + +def elapsed_ms(fn, warmup: int, iterations: int) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iterations): + fn() + end.record() + end.synchronize() + return start.elapsed_time(end) / iterations + + +def graph_elapsed_ms(fn, warmup: int, iterations: int) -> float: + """Measure replay cost, matching LightLLM's decode CUDA-graph path.""" + for _ in range(warmup): + fn() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + fn() + for _ in range(warmup): + graph.replay() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iterations): + graph.replay() + end.record() + end.synchronize() + return start.elapsed_time(end) / iterations + + +def max_rank(value: float) -> float: + tensor = torch.tensor(value, device="cuda", dtype=torch.float64) + dist.all_reduce(tensor, op=dist.ReduceOp.MAX) + return tensor.item() + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--tokens", type=int, default=17152) + parser.add_argument("--hidden-size", type=int, default=4096) + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--iterations", type=int, default=10) + args = parser.parse_args() + + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group("nccl", device_id=torch.device("cuda", local_rank)) + group_name = dist.group.WORLD.group_name + rank = dist.get_rank() + + shape = (args.tokens, args.hidden_size) + # Keep the repeated in-place NCCL input finite without adding a reset copy to + # the timed region. Correctness of the transport setup is checked separately. + source = torch.zeros(shape, device="cuda", dtype=torch.bfloat16) + + nccl_input = source.clone() + dist.all_reduce(nccl_input) + torch.cuda.synchronize() + torch.testing.assert_close(nccl_input, torch.zeros_like(nccl_input), rtol=0, atol=0) + nccl_input.copy_(source) + torch.cuda.synchronize() + nccl_ms = elapsed_ms(lambda: dist.all_reduce(nccl_input), args.warmup, args.iterations) + nccl_graph_ms = graph_elapsed_ms(lambda: dist.all_reduce(nccl_input), args.warmup, args.iterations) + + buffer = symm_mem.empty(source.numel(), device="cuda", dtype=source.dtype) + handle = symm_mem.rendezvous(buffer, group_name) + if getattr(handle, "multicast_ptr", 0) == 0: + raise RuntimeError("symmetric-memory multicast pointer is unavailable") + symm_output = torch.empty_like(source) + + def symm_all_reduce() -> None: + buffer.copy_(source.view(-1)) + torch.ops.symm_mem.multimem_all_reduce_(buffer, "sum", group_name) + symm_output.view(-1).copy_(buffer) + + def symm_all_reduce_out_of_place() -> torch.Tensor: + buffer.copy_(source.view(-1)) + torch.ops.symm_mem.multimem_all_reduce_(buffer, "sum", group_name) + return buffer.view_as(source) + + symm_all_reduce() + torch.cuda.synchronize() + torch.testing.assert_close(symm_output, torch.zeros_like(symm_output), rtol=0, atol=0) + torch.cuda.synchronize() + symm_ms = elapsed_ms(symm_all_reduce, args.warmup, args.iterations) + symm_graph_ms = graph_elapsed_ms(symm_all_reduce, args.warmup, args.iterations) + symm_out_ms = elapsed_ms(symm_all_reduce_out_of_place, args.warmup, args.iterations) + symm_out_graph_ms = graph_elapsed_ms(symm_all_reduce_out_of_place, args.warmup, args.iterations) + + nccl_max_ms = max_rank(nccl_ms) + symm_max_ms = max_rank(symm_ms) + symm_out_max_ms = max_rank(symm_out_ms) + nccl_graph_max_ms = max_rank(nccl_graph_ms) + symm_graph_max_ms = max_rank(symm_graph_ms) + symm_out_graph_max_ms = max_rank(symm_out_graph_ms) + + cpu_group = dist.new_group(list(range(dist.get_world_size())), backend="gloo") + workspace = flashinfer_comm.create_allreduce_fusion_workspace( + backend="trtllm", + world_size=dist.get_world_size(), + rank=rank, + max_token_num=args.tokens, + hidden_dim=args.hidden_size, + dtype=source.dtype, + comm_backend=TorchDistBackend(group=cpu_group), + ) + + def flashinfer_all_reduce() -> torch.Tensor: + return flashinfer_comm.allreduce_fusion( + input=source, + workspace=workspace, + pattern=flashinfer_comm.AllReduceFusionPattern.kAllReduce, + ) + + fi_output = flashinfer_all_reduce() + torch.cuda.synchronize() + torch.testing.assert_close(fi_output, torch.zeros_like(fi_output), rtol=0, atol=0) + fi_ms = elapsed_ms(flashinfer_all_reduce, args.warmup, args.iterations) + fi_graph_ms = graph_elapsed_ms(flashinfer_all_reduce, args.warmup, args.iterations) + fi_max_ms = max_rank(fi_ms) + fi_graph_max_ms = max_rank(fi_graph_ms) + + if rank == 0: + nbytes = source.numel() * source.element_size() + eager_best = min( + (nccl_max_ms, "nccl"), + (symm_max_ms, "symm"), + (symm_out_max_ms, "symm_out"), + (fi_max_ms, "flashinfer"), + )[1] + graph_best = min( + (nccl_graph_max_ms, "nccl"), + (symm_graph_max_ms, "symm"), + (symm_out_graph_max_ms, "symm_out"), + (fi_graph_max_ms, "flashinfer"), + )[1] + print( + f"shape={shape} bytes={nbytes} nccl_ms={nccl_max_ms:.6f} " + f"symm_multimem_ms={symm_max_ms:.6f} symm_out_ms={symm_out_max_ms:.6f} " + f"flashinfer_ms={fi_max_ms:.6f} " + f"best={eager_best} " + f"graph_nccl_ms={nccl_graph_max_ms:.6f} " + f"graph_symm_ms={symm_graph_max_ms:.6f} " + f"graph_symm_out_ms={symm_out_graph_max_ms:.6f} " + f"graph_flashinfer_ms={fi_graph_max_ms:.6f} " + f"graph_best={graph_best}" + ) + workspace.destroy() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tools/bench_glm53_kda_chunk_h.py b/tools/bench_glm53_kda_chunk_h.py new file mode 100644 index 0000000000..0c5fed8578 --- /dev/null +++ b/tools/bench_glm53_kda_chunk_h.py @@ -0,0 +1,121 @@ +#!/usr/bin/env python3 +"""Benchmark GLM-5.3 KDA chunk-state kernel configs on a real packed shape.""" + +from __future__ import annotations + +import argparse +import itertools +import json +import statistics + +import torch + +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.chunk_delta_h import ( + chunk_gated_delta_rule_fwd_h, +) +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.index import prepare_chunk_indices + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--sequences", type=int, default=64) + parser.add_argument("--sequence-length", type=int, default=268) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--repeats", type=int, default=7) + return parser.parse_args() + + +@torch.inference_mode() +def main() -> None: + args = parse_args() + device = torch.device("cuda") + dtype = torch.bfloat16 + batch = 1 + heads = 8 + key_dim = 128 + value_dim = 128 + total_tokens = args.sequences * args.sequence_length + + shape = (batch, total_tokens, heads, key_dim) + k = torch.zeros(shape, device=device, dtype=dtype) + w = torch.zeros(shape, device=device, dtype=dtype) + u = torch.zeros((batch, total_tokens, heads, value_dim), device=device, dtype=dtype) + # The fused safe-gate+cumsum stage keeps cumulative decay in fp32; the + # exp2 path in the state kernel expects that exact dtype. + gk = torch.zeros(shape, device=device, dtype=torch.float32) + initial_state = torch.zeros((args.sequences, heads, key_dim, value_dim), device=device, dtype=dtype) + cu_seqlens = torch.arange( + 0, + total_tokens + 1, + args.sequence_length, + device=device, + dtype=torch.int32, + ) + chunk_indices = prepare_chunk_indices(cu_seqlens, 64) + + def run(config: dict[str, int]) -> None: + output = chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + gk=gk, + initial_state=initial_state, + output_final_state=True, + save_new_value=True, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=64, + use_exp2=True, + run_config=config, + ) + del output + + results = [] + for value_tile, num_warps, num_stages in itertools.product((32, 64), (2, 4), (2, 3, 4)): + config = { + "BV": value_tile, + "num_warps": num_warps, + "num_stages": num_stages, + } + for _ in range(args.warmup): + run(config) + torch.cuda.synchronize() + + samples_ms = [] + for _ in range(args.repeats): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + run(config) + end.record() + end.synchronize() + samples_ms.append(start.elapsed_time(end)) + results.append( + { + **config, + "median_ms": statistics.median(samples_ms), + "min_ms": min(samples_ms), + } + ) + + results.sort(key=lambda item: item["median_ms"]) + print( + json.dumps( + { + "shape": { + "sequences": args.sequences, + "sequence_length": args.sequence_length, + "total_tokens": total_tokens, + "heads": heads, + "key_dim": key_dim, + "value_dim": value_dim, + }, + "results": results, + }, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/tools/bench_glm53_sglang_moe.py b/tools/bench_glm53_sglang_moe.py new file mode 100644 index 0000000000..6b2a378b30 --- /dev/null +++ b/tools/bench_glm53_sglang_moe.py @@ -0,0 +1,254 @@ +#!/usr/bin/env python3 +"""Tune SGLang TritonMoE tiles for GLM-5.3-Flash TP8 on H100.""" + +import argparse +import json +import traceback +from contextlib import contextmanager +from types import SimpleNamespace + +import torch + + +SMALL_M_CONFIGS = [ + # Decode runs with 24 physical rows for c8 + MTP2. The original sweep + # barely sampled this regime, so cover the useful tile/scheduling axes. + (8, 32, 8, 4, 3), + (8, 64, 8, 4, 3), + (8, 128, 8, 4, 3), + (8, 64, 4, 4, 3), + (8, 64, 16, 4, 3), + (8, 64, 32, 4, 3), + (8, 128, 4, 4, 3), + (8, 128, 16, 4, 3), + (8, 128, 32, 4, 3), + (8, 64, 8, 4, 2), + (8, 64, 8, 4, 4), + (8, 128, 8, 4, 2), + (8, 128, 8, 4, 4), + (8, 64, 8, 8, 3), + (8, 128, 8, 8, 3), + (16, 32, 16, 4, 3), + (16, 64, 4, 4, 3), + (16, 64, 8, 4, 3), + (16, 64, 16, 4, 3), + (16, 64, 32, 4, 3), + (16, 128, 4, 4, 3), + (16, 128, 8, 4, 3), + (16, 128, 16, 4, 3), + (16, 128, 32, 4, 3), + (16, 64, 16, 4, 2), + (16, 64, 16, 4, 4), + (16, 128, 16, 4, 2), + (16, 128, 16, 4, 4), + (16, 64, 16, 8, 3), + (16, 128, 16, 8, 3), + (32, 64, 16, 4, 3), + (32, 128, 16, 4, 3), +] + + +LARGE_M_CONFIGS = [ + (64, 128, 32, 4, 3), + (64, 128, 8, 4, 3), + (64, 128, 16, 4, 3), + (64, 128, 64, 4, 3), + (32, 128, 16, 4, 3), + (128, 128, 32, 4, 3), + (128, 128, 32, 8, 3), + (64, 64, 32, 4, 3), + (64, 256, 32, 8, 3), + (128, 256, 32, 8, 3), + (64, 128, 32, 4, 2), + (64, 128, 32, 4, 4), +] + + +CONFIGS = SMALL_M_CONFIGS + LARGE_M_CONFIGS + + +def make_config(values): + block_m, block_n, group_m, num_warps, num_stages = values + return { + "BLOCK_SIZE_M": block_m, + "BLOCK_SIZE_N": block_n, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": group_m, + "num_warps": num_warps, + "num_stages": num_stages, + } + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--tokens", type=int, default=17152) + parser.add_argument( + "--tp-size", + type=int, + default=8, + choices=(4, 8), + help="Tensor-parallel size; GLM-5 has 2048 total expert intermediate rows.", + ) + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--iters", type=int, default=10) + parser.add_argument("--max-configs", type=int, default=len(CONFIGS)) + parser.add_argument( + "--config-set", + choices=("all", "small", "large"), + default="all", + help="Restrict the sweep to decode-sized or prefill-sized tile configs.", + ) + parser.add_argument( + "--fuse-sum", + action="store_true", + help="Fuse the top-k expert sum into SGLang's down-projection kernel.", + ) + parser.add_argument( + "--tune-down", + action="store_true", + help="Tune a separate down-projection config with the best measured up config.", + ) + parser.add_argument( + "--fixed-up-config", + choices=("decode", "prefill"), + default="decode", + help="Fixed up-projection config used by --tune-down.", + ) + args = parser.parse_args() + + from sglang.srt.layers.moe.moe_runner.triton_utils import override_config + from sglang.srt.layers.moe.moe_runner.triton_utils import fused_moe + from sglang.srt.layers.moe.moe_runner.triton_utils import ( + fused_moe_triton_config, + ) + + standalone_exec = SimpleNamespace( + moe=SimpleNamespace(enable_fused_moe_sum_all_reduce=args.fuse_sum), + deterministic=SimpleNamespace(enable_deterministic_inference=False), + ) + fused_moe.get_exec = lambda: standalone_exec + fused_moe_triton_config.get_exec = lambda: standalone_exec + + device = torch.device("cuda:0") + experts, hidden, intermediate, topk = 289, 4096, 2048 // args.tp_size, 9 + x = torch.zeros((args.tokens, hidden), dtype=torch.bfloat16, device=device) + w1 = torch.zeros((experts, intermediate * 2, hidden), dtype=torch.float8_e4m3fn, device=device) + w2 = torch.zeros((experts, hidden, intermediate), dtype=torch.float8_e4m3fn, device=device) + w1_scale = torch.ones((experts, 4, 32), dtype=torch.float32, device=device) + w2_scale = torch.ones((experts, 32, 2), dtype=torch.float32, device=device) + rows = torch.arange(args.tokens, dtype=torch.int64, device=device)[:, None] + cols = torch.arange(topk, dtype=torch.int64, device=device)[None, :] + topk_ids = (rows * topk + cols) % experts + topk_weights = torch.full((args.tokens, topk), 1.0 / topk, dtype=torch.float32, device=device) + + fixed_up_config = make_config( + { + "decode": (16, 64, 16, 4, 3), + "prefill": (64, 128, 64, 4, 3), + }[args.fixed_up_config] + ) + + @contextmanager + def config_context(config): + if not args.tune_down: + with override_config(config): + yield + return + original = fused_moe.try_get_optimal_moe_config + + def resolve_config(*resolve_args, return_down_config=False, **resolve_kwargs): + if return_down_config: + return fixed_up_config, (config, None) + return fixed_up_config + + fused_moe.try_get_optimal_moe_config = resolve_config + try: + yield + finally: + fused_moe.try_get_optimal_moe_config = original + + def run(config): + with config_context(config): + fused_moe.fused_experts_impl( + hidden_states=x, + w1=w1, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + inplace=True, + use_fp8_w8a8=True, + w1_scale=w1_scale, + w2_scale=w2_scale, + block_shape=[128, 128], + routed_scaling_factor=1.0, + filter_expert=False, + swiglu_limit=10.0, + gate_up_interleaved=False, + ) + + def component_times(config): + original = fused_moe.invoke_fused_moe_kernel + events = [] + + def timed_kernel(*kernel_args, **kernel_kwargs): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + result = original(*kernel_args, **kernel_kwargs) + end.record() + events.append((start, end)) + return result + + fused_moe.invoke_fused_moe_kernel = timed_kernel + try: + run(config) + torch.cuda.synchronize() + finally: + fused_moe.invoke_fused_moe_kernel = original + if len(events) != 2: + raise RuntimeError(f"expected two MoE GEMM calls, got {len(events)}") + return events[0][0].elapsed_time(events[0][1]), events[1][0].elapsed_time(events[1][1]) + + results = [] + selected_configs = { + "all": CONFIGS, + "small": SMALL_M_CONFIGS, + "large": LARGE_M_CONFIGS, + }[args.config_set] + if args.tune_down: + # Both GEMMs share the same token alignment, so the down projection's + # BLOCK_SIZE_M must match the fixed up projection. + selected_configs = [values for values in CONFIGS if values[0] == fixed_up_config["BLOCK_SIZE_M"]] + for values in selected_configs[: args.max_configs]: + config = make_config(values) + try: + for _ in range(args.warmup): + run(config) + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(args.iters): + run(config) + end.record() + end.synchronize() + elapsed_ms = start.elapsed_time(end) / args.iters + up_ms, down_ms = component_times(config) + result = {"ms": elapsed_ms, "up_ms": up_ms, "down_ms": down_ms, **config} + if args.tune_down: + result["fixed_up_config"] = fixed_up_config + except Exception as exc: + result = {"error": f"{type(exc).__name__}: {exc}", **config} + if not results: + traceback.print_exc() + torch.cuda.synchronize() + results.append(result) + print(json.dumps(result, sort_keys=True), flush=True) + + valid = [result for result in results if "ms" in result] + if valid: + print("BEST", json.dumps(min(valid, key=lambda result: result["ms"]), sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/tools/bench_glm53_sparse_prefill.py b/tools/bench_glm53_sparse_prefill.py new file mode 100644 index 0000000000..df39ebbc3b --- /dev/null +++ b/tools/bench_glm53_sparse_prefill.py @@ -0,0 +1,160 @@ +"""Compare GLM-5.3 sparse prefill kernels at serving-sized shapes. + +This benchmark deliberately includes FlashMLA's TP head-padding allocation and +copy, because that work is part of every LightLLM prefill attention layer. +TileLang receives the native eight TP-local query heads. +""" + +import argparse +import gc +import json +import statistics +from collections.abc import Callable + +import torch + + +def make_causal_indices(tokens: int, sequence_length: int, topk: int) -> torch.Tensor: + """Build the packed-request causal index layout used by the serving test.""" + + token_ids = torch.arange(tokens, dtype=torch.int32, device="cuda") + positions = token_ids.remainder(sequence_length) + sequence_starts = token_ids - positions + columns = torch.arange(topk, dtype=torch.int32, device="cuda").view(1, -1) + indices = sequence_starts.view(-1, 1) + columns + indices.masked_fill_(columns > positions.view(-1, 1), -1) + return indices.unsqueeze(1) + + +def measure_ms( + name: str, + operation: Callable[[], torch.Tensor], + warmup: int, + iterations: int, +) -> dict[str, float | str]: + result = None + for _ in range(warmup): + result = operation() + torch.cuda.synchronize() + del result + + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + baseline_bytes = torch.cuda.memory_allocated() + samples = [] + for _ in range(iterations): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + result = operation() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end)) + del result + + peak_delta_bytes = torch.cuda.max_memory_allocated() - baseline_bytes + return { + "name": name, + "median_ms": statistics.median(samples), + "min_ms": min(samples), + "max_ms": max(samples), + "peak_delta_gib": peak_delta_bytes / 2 ** 30, + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--tokens", type=int, required=True) + parser.add_argument("--sequence-length", type=int, default=1024) + parser.add_argument("--topk", type=int, default=2048) + parser.add_argument("--local-heads", type=int, default=8) + parser.add_argument("--required-heads", type=int, default=64) + parser.add_argument("--head-dim", type=int, default=512) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--iterations", type=int, default=5) + args = parser.parse_args() + + if args.topk % 64: + raise ValueError("TileLang requires --topk to be divisible by 64") + if args.required_heads % args.local_heads: + raise ValueError("--required-heads must be divisible by --local-heads") + + from sgl_kernel.flash_mla import flash_mla_sparse_fwd + from sglang.kernels.ops.attention.dsa.tilelang_kernel import tilelang_sparse_fwd + + torch.manual_seed(0) + q = torch.randn( + (args.tokens, args.local_heads, args.head_dim), + dtype=torch.bfloat16, + device="cuda", + ) + kv = torch.randn( + (args.tokens, 1, args.head_dim), + dtype=torch.bfloat16, + device="cuda", + ) + indices = make_causal_indices(args.tokens, args.sequence_length, args.topk) + scale = args.head_dim ** -0.5 + + padded_q = q.new_zeros((args.tokens, args.required_heads, args.head_dim)) + padded_q[:, : args.local_heads].copy_(q) + + def flashmla_kernel() -> torch.Tensor: + return flash_mla_sparse_fwd( + padded_q, + kv, + indices, + scale, + d_v=args.head_dim, + )[0] + + def flashmla_lightllm_path() -> torch.Tensor: + q_input = q.new_zeros((args.tokens, args.required_heads, args.head_dim)) + q_input[:, : args.local_heads].copy_(q) + return flash_mla_sparse_fwd(q_input, kv, indices, scale, d_v=args.head_dim,)[ + 0 + ][:, : args.local_heads] + + def tilelang_lightllm_path() -> torch.Tensor: + output = tilelang_sparse_fwd( + q, + kv, + indices, + scale, + d_v=args.head_dim, + ) + return output.squeeze(0) if output.ndim == 4 else output + + results = [] + for name, operation in ( + ("flashmla_padded_kernel", flashmla_kernel), + ("flashmla_lightllm_path", flashmla_lightllm_path), + ("tilelang_lightllm_path", tilelang_lightllm_path), + ): + results.append(measure_ms(name, operation, args.warmup, args.iterations)) + gc.collect() + torch.cuda.empty_cache() + + flash_ms = next(r["median_ms"] for r in results if r["name"] == "flashmla_lightllm_path") + tile_ms = next(r["median_ms"] for r in results if r["name"] == "tilelang_lightllm_path") + print( + json.dumps( + { + "shape": { + "tokens": args.tokens, + "sequence_length": args.sequence_length, + "topk": args.topk, + "local_heads": args.local_heads, + "required_heads": args.required_heads, + "head_dim": args.head_dim, + }, + "results": results, + "tilelang_speedup_over_lightllm_flashmla": flash_ms / tile_ms, + }, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/tools/bench_glm53_sparse_prefill_tp.py b/tools/bench_glm53_sparse_prefill_tp.py new file mode 100644 index 0000000000..9f886f85da --- /dev/null +++ b/tools/bench_glm53_sparse_prefill_tp.py @@ -0,0 +1,255 @@ +"""Prototype TP head/token transpose for GLM-5.3 sparse prefill attention.""" + +import argparse +import gc +import json +import statistics +from collections.abc import Callable + +import torch +import torch.distributed as dist + + +def make_causal_indices(tokens: int, sequence_length: int, topk: int) -> torch.Tensor: + token_ids = torch.arange(tokens, dtype=torch.int32, device="cuda") + positions = token_ids.remainder(sequence_length) + sequence_starts = token_ids - positions + columns = torch.arange(topk, dtype=torch.int32, device="cuda").view(1, -1) + indices = sequence_starts.view(-1, 1) + columns + indices.masked_fill_(columns > positions.view(-1, 1), -1) + return indices.unsqueeze(1) + + +def head_shards_to_token_shards(q: torch.Tensor, world_size: int) -> torch.Tensor: + tokens, local_heads, head_dim = q.shape + if tokens % world_size: + raise ValueError(f"tokens={tokens} must be divisible by world_size={world_size}") + tokens_per_rank = tokens // world_size + received = torch.empty_like(q) + dist.all_to_all_single(received, q) + return ( + received.view(world_size, tokens_per_rank, local_heads, head_dim) + .permute(1, 0, 2, 3) + .contiguous() + .view(tokens_per_rank, world_size * local_heads, head_dim) + ) + + +def token_shards_to_head_shards(output: torch.Tensor, world_size: int) -> torch.Tensor: + tokens_per_rank, global_heads, head_dim = output.shape + if global_heads % world_size: + raise ValueError(f"global_heads={global_heads} must be divisible by world_size={world_size}") + local_heads = global_heads // world_size + send = ( + output.view(tokens_per_rank, world_size, local_heads, head_dim) + .permute(1, 0, 2, 3) + .contiguous() + .view(world_size * tokens_per_rank, local_heads, head_dim) + ) + received = torch.empty_like(send) + dist.all_to_all_single(received, send) + return received + + +def global_max(value: float, device: torch.device) -> float: + tensor = torch.tensor(value, dtype=torch.float64, device=device) + dist.all_reduce(tensor, op=dist.ReduceOp.MAX) + return tensor.item() + + +def measure_ms( + name: str, + operation: Callable[[], torch.Tensor], + warmup: int, + iterations: int, + device: torch.device, +) -> dict[str, float | str]: + result = None + for _ in range(warmup): + result = operation() + torch.cuda.synchronize() + del result + + gc.collect() + torch.cuda.empty_cache() + dist.barrier() + torch.cuda.reset_peak_memory_stats() + baseline_bytes = torch.cuda.memory_allocated() + samples = [] + for _ in range(iterations): + dist.barrier() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + result = operation() + end.record() + end.synchronize() + samples.append(global_max(start.elapsed_time(end), device)) + del result + + peak_delta_gib = global_max( + (torch.cuda.max_memory_allocated() - baseline_bytes) / 2 ** 30, + device, + ) + return { + "name": name, + "median_ms": statistics.median(samples), + "min_ms": min(samples), + "max_ms": max(samples), + "peak_delta_gib": peak_delta_gib, + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--tokens", type=int, required=True) + parser.add_argument("--sequence-length", type=int, default=1024) + parser.add_argument("--topk", type=int, default=2048) + parser.add_argument("--local-heads", type=int, default=8) + parser.add_argument("--head-dim", type=int, default=512) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--iterations", type=int, default=5) + parser.add_argument("--validate", action="store_true") + args = parser.parse_args() + + local_rank = int(__import__("os").environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group("nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + device = torch.device("cuda", local_rank) + global_heads = args.local_heads * world_size + if args.tokens % world_size: + raise ValueError("--tokens must be divisible by the TP world size") + if global_heads != 64: + raise ValueError(f"FlashMLA on Hopper requires 64 global heads, got {global_heads}") + + from sgl_kernel.flash_mla import flash_mla_sparse_fwd + + torch.manual_seed(1000 + rank) + q = torch.randn( + (args.tokens, args.local_heads, args.head_dim), + dtype=torch.bfloat16, + device=device, + ) + torch.manual_seed(0) + kv = torch.randn( + (args.tokens, 1, args.head_dim), + dtype=torch.bfloat16, + device=device, + ) + indices = make_causal_indices(args.tokens, args.sequence_length, args.topk) + scale = args.head_dim ** -0.5 + token_start = rank * (args.tokens // world_size) + token_end = token_start + args.tokens // world_size + local_indices = indices[token_start:token_end] + comm_workspace = torch.empty_like(q) + transpose_workspace = torch.empty( + (args.tokens // world_size, global_heads, args.head_dim), + dtype=q.dtype, + device=q.device, + ) + + def padded_flashmla() -> torch.Tensor: + q_input = q.new_zeros((args.tokens, global_heads, args.head_dim)) + q_input[:, : args.local_heads].copy_(q) + return flash_mla_sparse_fwd(q_input, kv, indices, scale, d_v=args.head_dim,)[ + 0 + ][:, : args.local_heads] + + def transposed_flashmla() -> torch.Tensor: + transposed_q = head_shards_to_token_shards(q, world_size) + transposed_output = flash_mla_sparse_fwd( + transposed_q, + kv, + local_indices, + scale, + d_v=args.head_dim, + )[0] + return token_shards_to_head_shards(transposed_output, world_size) + + def transposed_flashmla_workspace() -> torch.Tensor: + dist.all_to_all_single(comm_workspace, q) + comm_rank_major = comm_workspace.view( + world_size, + args.tokens // world_size, + args.local_heads, + args.head_dim, + ) + transpose_workspace.view( + args.tokens // world_size, + world_size, + args.local_heads, + args.head_dim, + ).copy_(comm_rank_major.permute(1, 0, 2, 3)) + transposed_output = flash_mla_sparse_fwd( + transpose_workspace, + kv, + local_indices, + scale, + d_v=args.head_dim, + )[0] + comm_workspace.view(world_size, args.tokens // world_size, args.local_heads, args.head_dim,).copy_( + transposed_output.view( + args.tokens // world_size, + world_size, + args.local_heads, + args.head_dim, + ).permute(1, 0, 2, 3) + ) + dist.all_to_all_single(transpose_workspace.view_as(q), comm_workspace) + return transpose_workspace.view_as(q) + + validation = None + if args.validate: + baseline = padded_flashmla().contiguous() + candidate = transposed_flashmla_workspace().clone() + torch.cuda.synchronize() + difference = (baseline - candidate).abs().float() + validation = { + "max_abs_diff": global_max(difference.max().item(), device), + "mean_abs_diff_max_rank": global_max(difference.mean().item(), device), + "allclose": bool(torch.allclose(baseline, candidate, rtol=0.05, atol=0.05)), + } + validation_tensor = torch.tensor(int(validation["allclose"]), device=device) + dist.all_reduce(validation_tensor, op=dist.ReduceOp.MIN) + validation["allclose"] = bool(validation_tensor.item()) + del baseline, candidate, difference + torch.cuda.empty_cache() + + results = [] + for name, operation in ( + ("flashmla_head_padding", padded_flashmla), + ("flashmla_tp_head_token_transpose", transposed_flashmla), + ("flashmla_tp_transpose_workspace", transposed_flashmla_workspace), + ): + results.append(measure_ms(name, operation, args.warmup, args.iterations, device)) + + if rank == 0: + baseline_ms = results[0]["median_ms"] + candidate_ms = results[-1]["median_ms"] + print( + json.dumps( + { + "shape": { + "tokens_per_rank_before_transpose": args.tokens, + "tokens_per_rank_after_transpose": args.tokens // world_size, + "sequence_length": args.sequence_length, + "topk": args.topk, + "local_heads": args.local_heads, + "global_heads": global_heads, + "head_dim": args.head_dim, + "world_size": world_size, + }, + "validation": validation, + "results": results, + "transpose_speedup_over_padding": baseline_ms / candidate_ms, + }, + indent=2, + ) + ) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tools/check_glm53_symm_out_of_place.py b/tools/check_glm53_symm_out_of_place.py new file mode 100644 index 0000000000..be8f928160 --- /dev/null +++ b/tools/check_glm53_symm_out_of_place.py @@ -0,0 +1,58 @@ +#!/usr/bin/env python3 +"""Exercise the production SymmMem alias path in normal and inference modes.""" + +import os +from contextlib import nullcontext + +import torch +import torch.distributed as dist + +from lightllm.distributed.symm_mem_all_reduce import SymmMemAllreduce + + +def _check(reducer: SymmMemAllreduce, *, inference: bool) -> None: + rank = dist.get_rank() + expected = sum(range(1, dist.get_world_size() + 1)) + context = torch.inference_mode() if inference else nullcontext() + with context: + for rows in (64, 32, 64): + value = torch.full( + (rows, 4096), + rank + 1, + dtype=torch.bfloat16, + device="cuda", + ) + value.data = reducer.all_reduce_out_of_place(value) + # Match the model's immediate post-reduction view operation. + viewed = value.view(rows, 4, 1024) + torch.testing.assert_close( + viewed, + torch.full_like(viewed, expected), + rtol=0, + atol=0, + ) + + +def main() -> None: + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group("nccl", device_id=torch.device("cuda", local_rank)) + reducer = SymmMemAllreduce( + dist.group.WORLD, + torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + if reducer.disabled: + raise RuntimeError("SymmMemAllreduce unexpectedly disabled") + if not reducer.buffer.is_inference(): + raise RuntimeError("SymmMem workspace is not an inference tensor") + _check(reducer, inference=False) + _check(reducer, inference=True) + dist.barrier() + if dist.get_rank() == 0: + print("SymmMem out-of-place normal+inference alias checks passed") + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tools/run_glm53_h100_container.sh b/tools/run_glm53_h100_container.sh new file mode 100755 index 0000000000..b563085fb6 --- /dev/null +++ b/tools/run_glm53_h100_container.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +set -euo pipefail + +image="${LIGHTLLM_GLM53_IMAGE:-lightllm-glm53:16k256-kpool}" +name="${LIGHTLLM_GLM53_CONTAINER:-glm53-lightllm-16k256-kpool}" +model_dir="${LIGHTLLM_GLM53_MODEL_DIR:-/home/devsft/models/GLM-5.3-Flash}" +cache_dir="${LIGHTLLM_GLM53_CACHE_DIR:-/home/devsft/cache-glm53-lightllm}" +triton_cache_dir="${LIGHTLLM_GLM53_TRITON_CACHE_DIR:-/home/devsft/cache-glm53-triton}" +deep_gemm_cache_dir="${LIGHTLLM_GLM53_DEEP_GEMM_CACHE_DIR:-/home/devsft/cache-glm53-deep-gemm}" + +if [[ ! -d "${model_dir}" ]]; then + echo "model directory does not exist: ${model_dir}" >&2 + exit 1 +fi + +mkdir -p "${cache_dir}" "${triton_cache_dir}" "${deep_gemm_cache_dir}" + +exec sudo docker run --rm --name "${name}" \ + --gpus all \ + --ipc host \ + --network host \ + -e LIGHTLLM_ENABLE_KPOOL_DECODE_FASTPATH=1 \ + -e PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \ + --ulimit memlock=-1:-1 \ + --ulimit nofile=1048576:1048576 \ + -v "${model_dir}:/model:ro" \ + -v "${cache_dir}:/root/.cache" \ + -v "${triton_cache_dir}:/root/.triton" \ + -v "${deep_gemm_cache_dir}:/root/.deep_gemm" \ + "${image}" diff --git a/tools/run_glm53_long_prompt_bench.py b/tools/run_glm53_long_prompt_bench.py new file mode 100644 index 0000000000..a31795d498 --- /dev/null +++ b/tools/run_glm53_long_prompt_bench.py @@ -0,0 +1,65 @@ +#!/usr/bin/env python3 +"""Run SGLang bench_serving with one cached long prompt reused per request.""" + +import copy +import os +import pickle +import runpy +import sys +import types +from pathlib import Path + +from sglang.benchmark.datasets.random import RandomDataset + + +def _install_prompt_cache() -> None: + cache_path = Path(os.environ["GLM53_LONG_PROMPT_CACHE"]) + original_load = RandomDataset.load + + def cached_load(self, tokenizer, model_id=None): + expected = { + "input_len": self.input_len, + "output_len": self.output_len, + "range_ratio": self.range_ratio, + } + if cache_path.exists(): + with cache_path.open("rb") as cache_file: + cached = pickle.load(cache_file) + if cached["metadata"] != expected: + raise ValueError(f"Long-prompt cache metadata mismatch: {cached['metadata']} != {expected}") + row = cached["row"] + print(f"Loaded long prompt from {cache_path}", flush=True) + else: + num_requests = self.num_requests + self.num_requests = 1 + try: + row = original_load(self, tokenizer, model_id)[0] + finally: + self.num_requests = num_requests + cache_path.parent.mkdir(parents=True, exist_ok=True) + with cache_path.open("wb") as cache_file: + pickle.dump({"metadata": expected, "row": row}, cache_file) + print(f"Cached long prompt at {cache_path}", flush=True) + + return [copy.copy(row) for _ in range(self.num_requests)] + + RandomDataset.load = cached_load + + +def main() -> None: + _install_prompt_cache() + + disaggregation_utils = types.ModuleType("sglang.srt.disaggregation.utils") + disaggregation_utils.FAKE_BOOTSTRAP_HOST = "fake" + sys.modules[disaggregation_utils.__name__] = disaggregation_utils + + network_utils = types.ModuleType("sglang.srt.utils.network") + network_utils.NetworkAddress = object + sys.modules[network_utils.__name__] = network_utils + + sys.argv = ["sglang.bench_serving", *sys.argv[1:]] + runpy.run_module("sglang.bench_serving", run_name="__main__") + + +if __name__ == "__main__": + main() diff --git a/unit_tests/common/basemodel/attention/test_flashmla_sparse_tp.py b/unit_tests/common/basemodel/attention/test_flashmla_sparse_tp.py new file mode 100644 index 0000000000..cfb5eabf92 --- /dev/null +++ b/unit_tests/common/basemodel/attention/test_flashmla_sparse_tp.py @@ -0,0 +1,86 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.attention.nsa import flashmla_sparse + + +def _disabled_modes(): + return SimpleNamespace( + enable_tpsp_mix_mode=False, + enable_prefill_cudagraph=False, + enable_prefill_microbatch_overlap=False, + enable_prefill_decode_mixed=False, + ) + + +def _infer_state(world_size=8): + return SimpleNamespace( + dist_group=SimpleNamespace(dp_world_size=world_size), + max_cache_len=0, + need_dp_prefill_balance=False, + use_replicated_attention_ep=False, + ) + + +def test_copy_received_head_shards_makes_token_major_global_heads(): + # Two source ranks sent the same destination's two-token chunk. + received = torch.tensor([[[10]], [[11]], [[20]], [[21]]]) + output = torch.empty((2, 2, 1), dtype=received.dtype) + + flashmla_sparse._copy_received_head_shards(received, output, world_size=2) + + torch.testing.assert_close(output, torch.tensor([[[10], [20]], [[11], [21]]])) + + +def test_copy_token_shard_for_head_scatter_makes_destination_blocks(): + output = torch.tensor([[[10], [20]], [[11], [21]]]) + send = torch.empty((4, 1, 1), dtype=output.dtype) + + flashmla_sparse._copy_token_shard_for_head_scatter(output, send, world_size=2) + + torch.testing.assert_close(send, torch.tensor([[[10]], [[11]], [[20]], [[21]]])) + + +def test_tp_head_token_transpose_accepts_validated_glm_tp8_layout(monkeypatch): + monkeypatch.setattr(flashmla_sparse, "get_env_start_args", _disabled_modes) + q = torch.empty((4096, 8, 2)) + + assert flashmla_sparse._should_use_tp_head_token_transpose(q, _infer_state(), required_heads=64) + + +@pytest.mark.parametrize( + ("tokens", "heads", "world_size", "state_change", "mode_change"), + [ + (2048, 8, 8, {}, {}), + (4097, 8, 8, {}, {}), + (4096, 16, 8, {}, {}), + (4096, 8, 1, {}, {}), + (4096, 8, 8, {"max_cache_len": 1}, {}), + (4096, 8, 8, {"need_dp_prefill_balance": True}, {}), + (4096, 8, 8, {"use_replicated_attention_ep": True}, {}), + (4096, 8, 8, {}, {"enable_tpsp_mix_mode": True}), + (4096, 8, 8, {}, {"enable_prefill_cudagraph": True}), + (4096, 8, 8, {}, {"enable_prefill_microbatch_overlap": True}), + (4096, 8, 8, {}, {"enable_prefill_decode_mixed": True}), + ], +) +def test_tp_head_token_transpose_rejects_unvalidated_layouts( + monkeypatch, + tokens, + heads, + world_size, + state_change, + mode_change, +): + modes = _disabled_modes() + for name, value in mode_change.items(): + setattr(modes, name, value) + monkeypatch.setattr(flashmla_sparse, "get_env_start_args", lambda: modes) + state = _infer_state(world_size) + for name, value in state_change.items(): + setattr(state, name, value) + q = torch.empty((tokens, heads, 2)) + + assert not flashmla_sparse._should_use_tp_head_token_transpose(q, state, required_heads=64) diff --git a/unit_tests/common/basemodel/attention/test_tilelang_sparse.py b/unit_tests/common/basemodel/attention/test_tilelang_sparse.py new file mode 100644 index 0000000000..448276703c --- /dev/null +++ b/unit_tests/common/basemodel/attention/test_tilelang_sparse.py @@ -0,0 +1,27 @@ +import pytest +import torch + +from lightllm.common.basemodel.attention.nsa.tilelang_sparse import pad_sparse_indices + + +def test_pad_sparse_indices_adds_masked_block_tail(): + indices = torch.arange(65, dtype=torch.int32).view(1, 65) + + padded = pad_sparse_indices(indices) + + assert padded.shape == (1, 1, 128) + torch.testing.assert_close(padded[0, 0, :65], indices[0]) + assert torch.all(padded[0, 0, 65:] == -1) + + +def test_pad_sparse_indices_preserves_aligned_storage(): + indices = torch.zeros((4, 1, 128), dtype=torch.int32) + + assert pad_sparse_indices(indices) is indices + + +def test_pad_sparse_indices_rejects_invalid_shape_or_block(): + with pytest.raises(ValueError, match="2D or 3D"): + pad_sparse_indices(torch.zeros((2, 3, 4, 5), dtype=torch.int32)) + with pytest.raises(ValueError, match="positive"): + pad_sparse_indices(torch.zeros((2, 64), dtype=torch.int32), block_size=0) diff --git a/unit_tests/common/basemodel/test_cuda_graph_layout.py b/unit_tests/common/basemodel/test_cuda_graph_layout.py index 17ddcd98da..ce84bf5585 100644 --- a/unit_tests/common/basemodel/test_cuda_graph_layout.py +++ b/unit_tests/common/basemodel/test_cuda_graph_layout.py @@ -76,3 +76,31 @@ def test_batch_step_size_after_split_controls_capture_range(_graph_args): 42, 56, ] + + +def test_extra_batch_sizes_are_merged_and_bounded(_graph_args): + graph = CudaGraph( + batch_step_size_before_split=6, + split_batch_size=24, + batch_step_size_after_split=96, + max_batch_size=384, + extra_batch_sizes=[1, 2, 4, 20, 36, 52, 64, 999], + ) + + assert graph.cuda_graph_batch_sizes == [ + 1, + 2, + 4, + 6, + 12, + 18, + 20, + 24, + 36, + 52, + 64, + 120, + 216, + 312, + 384, + ] diff --git a/unit_tests/common/basemodel/test_hidden_collector.py b/unit_tests/common/basemodel/test_hidden_collector.py index 22748b7294..6d9937d941 100644 --- a/unit_tests/common/basemodel/test_hidden_collector.py +++ b/unit_tests/common/basemodel/test_hidden_collector.py @@ -103,18 +103,22 @@ def test_noop_collector_keeps_normal_forward_output_minimal(): def test_mtp_head_output_collector_returns_and_clears_outputs(): collector = MtpHeadOutputCollector() draft_token_ids = torch.arange(6) + draft_token_probs = torch.linspace(0.1, 0.6, 6) confidence_logits = torch.arange(6).view(2, 3) collector.add_mtp_outputs( draft_token_ids=draft_token_ids, + draft_token_probs=draft_token_probs, confidence_logits=confidence_logits, ) output = collector.finish_output(infer_state=None) assert output.spec_hidden is None assert output.draft_token_ids is draft_token_ids + assert output.draft_token_probs is draft_token_probs assert output.confidence_logits is confidence_logits assert collector.draft_token_ids is None + assert collector.draft_token_probs is None assert collector.confidence_logits is None diff --git a/unit_tests/common/basemodel/test_model_output.py b/unit_tests/common/basemodel/test_model_output.py index 6f9477e294..54553ce589 100644 --- a/unit_tests/common/basemodel/test_model_output.py +++ b/unit_tests/common/basemodel/test_model_output.py @@ -7,20 +7,46 @@ from lightllm.common.basemodel.batch_objs import ModelInput, ModelMtpOutputCollector, ModelOutput +def test_vocab_parallel_metadata_follows_row_selection(): + output = ModelOutput( + logits=torch.tensor([[8.0], [7.0], [6.0]]), + logits_token_ids=torch.tensor([[4], [9], [2]]), + logits_logsumexp=torch.tensor([8.5, 7.25, 6.75]), + ) + + selected = output.index_select_logits_rows(torch.tensor([2, 0])) + + torch.testing.assert_close(selected.logits.view(-1), torch.tensor([6.0, 8.0])) + torch.testing.assert_close(selected.logits_token_ids.view(-1), torch.tensor([2, 4])) + torch.testing.assert_close(selected.logits_logsumexp, torch.tensor([6.75, 8.5])) + + def test_decode_unpad_slices_spec_output_with_logits(): model = TpPartBaseModel.__new__(TpPartBaseModel) output = ModelOutput( logits=torch.arange(24).view(6, 4), - mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.arange(18).view(6, 3)), + logits_token_ids=torch.arange(100, 124).view(6, 4), + logits_logsumexp=torch.arange(6, dtype=torch.float32), + mtp_collector=ModelMtpOutputCollector( + spec_hidden=torch.arange(18).view(6, 3), + draft_token_ids=torch.arange(6), + draft_token_probs=torch.linspace(0.1, 0.6, 6), + ), ) unpadded = model._create_unpad_decode_model_output(output, origin_batch_size=4) assert unpadded.logits.shape == (4, 4) + assert unpadded.logits_token_ids.shape == (4, 4) + assert unpadded.logits_logsumexp.shape == (4,) assert unpadded.mtp_collector.spec_hidden.shape == (4, 3) + assert unpadded.mtp_collector.draft_token_ids.shape == (4,) + assert unpadded.mtp_collector.draft_token_probs.shape == (4,) # Unpadding returns a shallow output copy and leaves the graph-owned # tensors on the original ModelOutput intact. assert output.logits.shape == (6, 4) + assert output.logits_token_ids.shape == (6, 4) + assert output.logits_logsumexp.shape == (6,) assert output.mtp_collector.spec_hidden.shape == (6, 3) @@ -28,6 +54,8 @@ def test_prefill_unpad_uses_token_rows_for_spec_hidden(): model = TpPartBaseModel.__new__(TpPartBaseModel) output = ModelOutput( logits=torch.arange(20).view(5, 4), + logits_token_ids=torch.arange(100, 120).view(5, 4), + logits_logsumexp=torch.arange(5, dtype=torch.float32), mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.arange(24).view(8, 3)), prompt_logics=torch.arange(32).view(8, 4), ) @@ -39,6 +67,8 @@ def test_prefill_unpad_uses_token_rows_for_spec_hidden(): ) assert unpadded.logits.shape == (3, 4) + assert unpadded.logits_token_ids.shape == (3, 4) + assert unpadded.logits_logsumexp.shape == (3,) assert unpadded.mtp_collector.spec_hidden.shape == (6, 3) assert unpadded.prompt_logics.shape == (6, 4) @@ -118,6 +148,7 @@ def test_decode_pads_only_once_after_selecting_execution_path(monkeypatch): model = TpPartBaseModel.__new__(TpPartBaseModel) model.args = SimpleNamespace(enable_tpsp_mix_mode=enable_tpsp_mix_mode) model.tp_world_size_ = tp_world_size + model.decode_batch_multiplier = 1 model.mem_manager = SimpleNamespace(HOLD_TOKEN_MEMINDEX=99) model.req_manager = SimpleNamespace(HOLD_REQUEST_ID=88, req_to_token_indexs=object()) diff --git a/unit_tests/common/basemodel/test_mtp_manager.py b/unit_tests/common/basemodel/test_mtp_manager.py index c7ad04f56b..c1aa7a8870 100644 --- a/unit_tests/common/basemodel/test_mtp_manager.py +++ b/unit_tests/common/basemodel/test_mtp_manager.py @@ -53,8 +53,8 @@ def test_decode_batch_multiplier(monkeypatch, spec_mode, is_draft_model, expecte [ (False, False, 8), (True, False, 1), - (False, True, 1), - (True, True, 1), + (False, True, 8), + (True, True, 8), ], ) def test_decode_cuda_graph_grow_step_size(monkeypatch, dynamic_verify, is_draft_model, expected): @@ -68,6 +68,40 @@ def test_decode_cuda_graph_grow_step_size(monkeypatch, dynamic_verify, is_draft_ assert MtpManager.get_instance().get_decode_cuda_graph_grow_step_size(is_draft_model) == expected +@pytest.mark.parametrize( + "spec_mode,is_draft_model,expected", + [ + ("eagle_with_att", True, 8), + ("eagle3", True, 8), + ("eagle_no_att", True, 1), + ("vanilla_with_att", True, 8), + ("eagle_with_att", False, 8), + ], +) +def test_decode_cuda_graph_batch_multiplier(monkeypatch, spec_mode, is_draft_model, expected): + args = SimpleNamespace(mtp_mode=spec_mode, mtp_step=7, mtp_dynamic_verify=False) + monkeypatch.setattr(mtp_manager_module, "get_env_start_args", lambda: args) + + manager = MtpManager.get_instance() + assert manager.get_decode_cuda_graph_batch_multiplier(is_draft_model) == expected + + +@pytest.mark.parametrize( + "spec_mode,expected", + [ + ("eagle_with_att", True), + ("eagle3", True), + ("eagle_no_att", False), + ("vanilla_with_att", False), + ], +) +def test_recurrent_attention_draft_keeps_logical_graph_schedule(monkeypatch, spec_mode, expected): + args = SimpleNamespace(mtp_mode=spec_mode, mtp_step=7, mtp_dynamic_verify=False) + monkeypatch.setattr(mtp_manager_module, "get_env_start_args", lambda: args) + + assert MtpManager.get_instance().draft_model_needs_logical_batch_graphs(True) is expected + + @pytest.mark.parametrize( "spec_mode,is_draft_model,expected", [ diff --git a/unit_tests/common/basemodel/test_prefill_cuda_graph_selection.py b/unit_tests/common/basemodel/test_prefill_cuda_graph_selection.py new file mode 100644 index 0000000000..efd1f384fe --- /dev/null +++ b/unit_tests/common/basemodel/test_prefill_cuda_graph_selection.py @@ -0,0 +1,121 @@ +from types import SimpleNamespace + +import pytest + +import lightllm.common.basemodel.prefill_cuda_graph as prefill_cuda_graph + + +class _DecodeGraph: + mempool = object() + + +def _args(**overrides): + values = { + "enable_prefill_microbatch_overlap": False, + "prefill_cudagraph_max_handle_token": 32768, + "prefill_cudagraph_token_nums": None, + "prefill_cudagraph_batch_sizes": None, + "prefill_cudagraph_capture_attention": False, + "batch_max_tokens": 32768, + "enable_tpsp_mix_mode": False, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def test_explicit_token_nums_only_run_exact_shapes(monkeypatch): + monkeypatch.setattr( + prefill_cuda_graph, + "get_env_start_args", + lambda: _args( + prefill_cudagraph_token_nums=[17152, 352, 17152], + prefill_cudagraph_batch_sizes=[64, 1, 64], + ), + ) + + graph = prefill_cuda_graph.PrefillCudaGraph(_DecodeGraph(), tp_world_size=8) + + assert graph.graph_handle_token_nums == [352, 17152] + assert graph.can_run(352, batch_size=1, max_q_seq_len=352, max_kv_seq_len=352, max_cache_len=0) + assert graph.can_run(17152, batch_size=64, max_q_seq_len=268, max_kv_seq_len=268, max_cache_len=0) + assert not graph.can_run(17152, batch_size=1, max_q_seq_len=17152, max_kv_seq_len=17152, max_cache_len=0) + assert not graph.can_run(17152, batch_size=64, max_q_seq_len=269, max_kv_seq_len=269, max_cache_len=0) + assert not graph.can_run(17152, batch_size=64, max_q_seq_len=268, max_kv_seq_len=300, max_cache_len=32) + assert not graph.can_run(351, batch_size=1, max_q_seq_len=351, max_kv_seq_len=351, max_cache_len=0) + assert not graph.can_run(16000, batch_size=64, max_q_seq_len=250, max_kv_seq_len=250, max_cache_len=0) + + +def test_explicit_token_nums_respect_limits(monkeypatch): + monkeypatch.setattr( + prefill_cuda_graph, + "get_env_start_args", + lambda: _args( + prefill_cudagraph_max_handle_token=20000, + batch_max_tokens=18000, + prefill_cudagraph_token_nums=[0, 17152, 20000], + prefill_cudagraph_batch_sizes=[1, 64, 64], + ), + ) + + graph = prefill_cuda_graph.PrefillCudaGraph(_DecodeGraph(), tp_world_size=8) + + assert graph.graph_handle_token_nums == [17152] + + +def test_explicit_token_nums_reject_empty_valid_set(monkeypatch): + monkeypatch.setattr( + prefill_cuda_graph, + "get_env_start_args", + lambda: _args( + prefill_cudagraph_token_nums=[0, 40000], + prefill_cudagraph_batch_sizes=[1, 64], + ), + ) + + with pytest.raises(ValueError, match="prefill_cudagraph_token_nums"): + prefill_cuda_graph.PrefillCudaGraph(_DecodeGraph(), tp_world_size=8) + + +def test_explicit_layout_rejects_missing_or_mismatched_batch_sizes(monkeypatch): + monkeypatch.setattr( + prefill_cuda_graph, + "get_env_start_args", + lambda: _args(prefill_cudagraph_token_nums=[17152]), + ) + with pytest.raises(ValueError, match="prefill_cudagraph_batch_sizes is required"): + prefill_cuda_graph.PrefillCudaGraph(_DecodeGraph(), tp_world_size=8) + + monkeypatch.setattr( + prefill_cuda_graph, + "get_env_start_args", + lambda: _args( + prefill_cudagraph_token_nums=[17152, 352], + prefill_cudagraph_batch_sizes=[64], + ), + ) + with pytest.raises(ValueError, match="same number of entries"): + prefill_cuda_graph.PrefillCudaGraph(_DecodeGraph(), tp_world_size=8) + + +def test_explicit_layout_rejects_nonuniform_sequence_shape(monkeypatch): + monkeypatch.setattr( + prefill_cuda_graph, + "get_env_start_args", + lambda: _args( + prefill_cudagraph_token_nums=[17153], + prefill_cudagraph_batch_sizes=[64], + ), + ) + with pytest.raises(ValueError, match="must be divisible"): + prefill_cuda_graph.PrefillCudaGraph(_DecodeGraph(), tp_world_size=8) + + +def test_attention_capture_requires_exact_layout(monkeypatch): + monkeypatch.setattr( + prefill_cuda_graph, + "get_env_start_args", + lambda: _args(prefill_cudagraph_capture_attention=True), + ) + + with pytest.raises(ValueError, match="requires --prefill_cudagraph_token_nums"): + prefill_cuda_graph.PrefillCudaGraph(_DecodeGraph(), tp_world_size=8) diff --git a/unit_tests/common/basemodel/test_sglang_triton_moe_config.py b/unit_tests/common/basemodel/test_sglang_triton_moe_config.py new file mode 100644 index 0000000000..64898b4156 --- /dev/null +++ b/unit_tests/common/basemodel/test_sglang_triton_moe_config.py @@ -0,0 +1,100 @@ +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl import ( + triton_impl, +) + + +def test_glm5_decode_uses_measured_h100_moe_config(monkeypatch): + triton_impl._get_sglang_triton_moe_configs.cache_clear() + monkeypatch.setattr(triton_impl.torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(triton_impl.torch.cuda, "get_device_name", lambda _: "NVIDIA H100 80GB HBM3") + + up_config, down_config = triton_impl._get_sglang_triton_moe_configs( + (289, 512, 4096), + (289, 4096, 256), + 9, + False, + 24, + ) + + assert up_config["BLOCK_SIZE_M"] == 16 + assert up_config["BLOCK_SIZE_N"] == 64 + assert up_config["num_stages"] == 3 + assert down_config["BLOCK_SIZE_M"] == up_config["BLOCK_SIZE_M"] + assert down_config["BLOCK_SIZE_N"] == 64 + assert down_config["num_stages"] == 2 + + +def test_glm5_draft_uses_measured_h100_down_config(monkeypatch): + triton_impl._get_sglang_triton_moe_configs.cache_clear() + monkeypatch.setattr(triton_impl.torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(triton_impl.torch.cuda, "get_device_name", lambda _: "NVIDIA H100 80GB HBM3") + + up_config, down_config = triton_impl._get_sglang_triton_moe_configs( + (289, 512, 4096), + (289, 4096, 256), + 9, + False, + 8, + ) + + assert up_config["BLOCK_SIZE_M"] == 16 + assert down_config["GROUP_SIZE_M"] == 8 + assert down_config["num_stages"] == 3 + + +def test_glm5_prefill_uses_measured_h100_moe_config(monkeypatch): + triton_impl._get_sglang_triton_moe_configs.cache_clear() + monkeypatch.setattr(triton_impl.torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(triton_impl.torch.cuda, "get_device_name", lambda _: "NVIDIA H100 80GB HBM3") + + (up_config, down_config) = triton_impl._get_sglang_triton_moe_configs( + (289, 512, 4096), + (289, 4096, 256), + 9, + True, + 8192, + ) + + assert up_config["BLOCK_SIZE_M"] == 64 + assert up_config["BLOCK_SIZE_N"] == 128 + assert up_config["GROUP_SIZE_M"] == 64 + assert down_config["BLOCK_SIZE_M"] == up_config["BLOCK_SIZE_M"] + assert down_config["BLOCK_SIZE_N"] == 128 + assert down_config["GROUP_SIZE_M"] == 8 + + +def test_glm5_large_prefill_uses_independent_down_tile(monkeypatch): + triton_impl._get_sglang_triton_moe_configs.cache_clear() + monkeypatch.setattr(triton_impl.torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(triton_impl.torch.cuda, "get_device_name", lambda _: "NVIDIA H100 80GB HBM3") + + up_config, down_config = triton_impl._get_sglang_triton_moe_configs( + (289, 512, 4096), + (289, 4096, 256), + 9, + True, + 65536, + ) + + assert up_config["BLOCK_SIZE_N"] == 128 + assert up_config["GROUP_SIZE_M"] == 64 + assert down_config["BLOCK_SIZE_M"] == up_config["BLOCK_SIZE_M"] + assert down_config["BLOCK_SIZE_N"] == 64 + assert down_config["GROUP_SIZE_M"] == 32 + + +def test_non_h100_has_no_sglang_override(monkeypatch): + triton_impl._get_sglang_triton_moe_configs.cache_clear() + monkeypatch.setattr(triton_impl.torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(triton_impl.torch.cuda, "get_device_name", lambda _: "NVIDIA H20") + + assert ( + triton_impl._get_sglang_triton_moe_configs( + (289, 512, 4096), + (289, 4096, 256), + 9, + False, + 24, + ) + is None + ) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_fused_gate.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_fused_gate.py new file mode 100644 index 0000000000..44a5ac8c44 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_fused_gate.py @@ -0,0 +1,76 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.fused_recurrent import ( + fused_recurrent_gated_delta_rule, +) +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda import kda_safe_gate + + +if not torch.cuda.is_available(): + pytest.skip("CUDA required", allow_module_level=True) + + +@pytest.mark.parametrize("sequence_length", [3, 6]) +def test_fused_kda_gate_matches_materialized_gate(sequence_length): + torch.manual_seed(2026 + sequence_length) + request_count, query_heads, value_heads, key_dim, value_dim = 3, 2, 4, 64, 64 + token_count = request_count * sequence_length + slot_count = token_count + 8 + + q = torch.randn(1, token_count, query_heads, key_dim, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(1, token_count, value_heads, value_dim, device="cuda", dtype=torch.bfloat16) + raw_gate = torch.randn(1, token_count, value_heads * key_dim, device="cuda", dtype=torch.bfloat16) + raw_beta = torch.randn(1, token_count, value_heads, device="cuda", dtype=torch.bfloat16) + a_log = torch.randn(value_heads, device="cuda", dtype=torch.float32) * 0.1 + gate_bias = torch.randn(value_heads * key_dim, device="cuda", dtype=torch.float32) * 0.1 + state = torch.randn( + slot_count, + value_heads, + key_dim, + value_dim, + device="cuda", + dtype=torch.bfloat16, + ) + state_indices = torch.arange(token_count, device="cuda", dtype=torch.int32).view(request_count, sequence_length) + cu_seqlens = torch.arange(request_count + 1, device="cuda", dtype=torch.int64) * sequence_length + accepted = torch.full((request_count,), sequence_length, device="cuda", dtype=torch.int32) + + reference_state = state.clone() + reference, _ = fused_recurrent_gated_delta_rule( + q=q, + k=k, + v=v, + g=kda_safe_gate(raw_gate, a_log, gate_bias), + beta=raw_beta.float().sigmoid(), + initial_state=reference_state, + cu_seqlens=cu_seqlens, + ssm_state_indices=state_indices, + ssm_state_write_indices=state_indices, + num_accepted_tokens=accepted, + use_qk_l2norm_in_kernel=True, + is_kda=True, + ) + + fused_state = state.clone() + fused, _ = fused_recurrent_gated_delta_rule( + q=q, + k=k, + v=v, + initial_state=fused_state, + cu_seqlens=cu_seqlens, + ssm_state_indices=state_indices, + ssm_state_write_indices=state_indices, + num_accepted_tokens=accepted, + use_qk_l2norm_in_kernel=True, + A_log=a_log, + dt_bias=gate_bias, + a_raw=raw_gate.reshape(token_count, value_heads * key_dim), + b_raw=raw_beta.reshape(token_count, value_heads), + is_kda=True, + kda_lower_bound=-5.0, + ) + + torch.testing.assert_close(fused, reference, rtol=2e-3, atol=2e-3) + torch.testing.assert_close(fused_state, reference_state, rtol=2e-3, atol=2e-3) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py new file mode 100644 index 0000000000..c70452aa08 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py @@ -0,0 +1,201 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.attention.linear import gdn as gdn_module +from lightllm.common.basemodel.triton_kernel.linear_att import causal_conv1d_mtp +from lightllm.common.basemodel.triton_kernel.linear_att import mtp_fused_recurrent as mtp_recurrent +from lightllm.common.triton_utils.autotuner import Autotuner + + +def _layout_tensors(total_tokens=8, K=64, V=128): + H, HV = 1, 2 + q = torch.empty((1, total_tokens, H, K), dtype=torch.bfloat16) + k = torch.empty_like(q) + v = torch.empty((1, total_tokens, HV, V), dtype=torch.bfloat16) + return q, k, v + + +def test_v1_config_space_and_fallback(): + configs = mtp_recurrent._get_mtp_fused_recurrent_configs() + assert len(configs) == 60 + assert len({tuple(sorted(config.items())) for config in configs}) == len(configs) + assert {config["BV"] for config in configs} == {4, 8, 16, 32, 64} + assert {config["num_warps"] for config in configs} == {1, 2, 4, 8} + assert {config["num_stages"] for config in configs} == {1, 2, 3} + assert mtp_recurrent._default_mtp_fused_recurrent_config(128) == { + "BV": 8, + "num_warps": 1, + "num_stages": 3, + } + + +def test_key_separates_fixed_and_dynamic_specializations(): + q, _, v = _layout_tensors() + initial_state = torch.empty((16, 2, 64, 128), dtype=torch.bfloat16) + cu_seqlens = torch.empty(3, dtype=torch.int64) + + dynamic_key = mtp_recurrent._get_mtp_fused_recurrent_static_key(q, v, initial_state, fixed_seq_len=0) + fixed_key = mtp_recurrent._get_mtp_fused_recurrent_static_key(q, v, initial_state, fixed_seq_len=4) + + assert dynamic_key["fixed_seq_len"] == 0 + assert fixed_key["fixed_seq_len"] == 4 + assert dynamic_key != fixed_key + assert mtp_recurrent._get_mtp_fused_recurrent_run_key(q, cu_seqlens) == 2 * 10 ** 9 + 8 + + +def test_tuner_declares_mutated_state(): + assert mtp_recurrent._mtp_fused_recurrent_gated_delta_rule_autotuned.mutates_args == ["initial_state"] + + +def test_gdn_cuda_graph_warmup_scopes_v1_autotune_to_first_layer_microbatch(monkeypatch): + Autotuner.end_autotune_warmup() + calls = [] + q, k, v = _layout_tensors() + backend = SimpleNamespace( + activation="silu", + conv_kernel_dim=4, + model=SimpleNamespace( + is_mtp_draft_model=False, + mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 3), + ), + uses_dynamic_spec_verify_layout=lambda: False, + _rearrange_mixed_qkv=lambda _, decode: (q, k, v), + ) + state = gdn_module.LinearAttDecodeAttState(backend=backend) + state.b1_mtp_cu_q_seq_len = torch.tensor([0, 4, 8], dtype=torch.int32) + state.b_conv_buffer_idx = torch.tensor([0, 1], dtype=torch.int32) + state.b_ssm_buffer_idx = torch.zeros((2, 4), dtype=torch.int32) + state.b_num_accepted_tokens = torch.ones(2, dtype=torch.int32) + infer_state = SimpleNamespace(is_cuda_graph=True, microbatch_index=0) + layer_weight = SimpleNamespace( + layer_num_=0, + linear_conv1d=SimpleNamespace(mm_param=SimpleNamespace(weight=torch.empty(1)), bias=None), + linear_A_log=SimpleNamespace(weight=torch.empty(1)), + linear_dt_bias=SimpleNamespace(weight=torch.empty(1)), + ) + + monkeypatch.setattr(gdn_module, "get_triton_autotune_level", lambda: 1) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + monkeypatch.setattr(causal_conv1d_mtp, "causal_conv1d_update", lambda mixed_qkv, *args, **kwargs: mixed_qkv) + + def fake_recurrent(**kwargs): + calls.append(Autotuner.is_autotune_warmup()) + return kwargs["q"], kwargs["initial_state"] + + monkeypatch.setattr(gdn_module, "mtp_fused_recurrent_gated_delta_rule", fake_recurrent) + + args = ( + torch.empty((8, 1)), + torch.empty(1), + torch.empty(1), + torch.empty((8, 2)), + torch.empty((8, 2)), + infer_state, + layer_weight, + ) + state._gdn_mtp_kernel(*args) + layer_weight.layer_num_ = 1 + state._gdn_mtp_kernel(*args) + + assert calls == [True, False] + assert not Autotuner.is_autotune_warmup() + + +def _make_fixed_case(): + torch.manual_seed(11) + batch, seq_len, H, HV, K, V = 2, 4, 1, 2, 64, 128 + total_tokens = batch * seq_len + q = torch.randn((1, total_tokens, H, K), device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn((1, total_tokens, HV, V), device="cuda", dtype=torch.bfloat16) + initial_state = torch.randn((10, HV, K, V), device="cuda", dtype=torch.bfloat16) + cu_seqlens = torch.arange(batch + 1, device="cuda", dtype=torch.int64) * seq_len + read_indices = torch.tensor([[0, 0, 0, 0], [1, 1, 1, 1]], device="cuda", dtype=torch.int32) + write_indices = torch.tensor([[2, 3, 4, 5], [6, 7, 8, 9]], device="cuda", dtype=torch.int32) + accepted = torch.full((batch,), seq_len, device="cuda", dtype=torch.int32) + A_log = torch.randn(HV, device="cuda", dtype=torch.float32) * 0.1 + dt_bias = torch.randn(HV, device="cuda", dtype=torch.float32) * 0.1 + a_raw = torch.randn((total_tokens, HV), device="cuda", dtype=torch.bfloat16) + b_raw = torch.randn_like(a_raw) + return ( + q, + k, + v, + initial_state, + cu_seqlens, + read_indices, + write_indices, + accepted, + A_log, + dt_bias, + a_raw, + b_raw, + ) + + +def _launch(case, run_config, *, fixed_seq_len): + ( + q, + k, + v, + initial_state, + cu_seqlens, + read_indices, + write_indices, + accepted, + A_log, + dt_bias, + a_raw, + b_raw, + ) = case + state = initial_state.clone() + kwargs = { + "q": q, + "k": k, + "v": v, + "initial_state": state, + "cu_seqlens": cu_seqlens, + "ssm_state_indices": read_indices, + "ssm_state_write_indices": write_indices, + "num_accepted_tokens": accepted, + "A_log": A_log, + "dt_bias": dt_bias, + "a_raw": a_raw, + "b_raw": b_raw, + "fixed_seq_len": fixed_seq_len, + "run_config": run_config, + } + output, _ = mtp_recurrent._mtp_fused_recurrent_gated_delta_rule_autotuned(**kwargs) + return output, state, kwargs + + +_RETAINED_CONFIGS = mtp_recurrent._get_mtp_fused_recurrent_configs() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("fixed_seq_len", [0, 4]) +@pytest.mark.parametrize("run_config", _RETAINED_CONFIGS) +def test_retained_configs_match_fallback(run_config, fixed_seq_len): + case = _make_fixed_case() + fallback = mtp_recurrent._default_mtp_fused_recurrent_config(128) + expected_output, expected_state, _ = _launch(case, fallback, fixed_seq_len=fixed_seq_len) + output, state, _ = _launch(case, run_config, fixed_seq_len=fixed_seq_len) + + torch.testing.assert_close(output.float(), expected_output.float(), atol=2e-2, rtol=2e-2) + torch.testing.assert_close(state.float(), expected_state.float(), atol=5.0, rtol=1e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_tuner_warmup_does_not_mutate_caller_state(): + case = _make_fixed_case() + fallback = mtp_recurrent._default_mtp_fused_recurrent_config(128) + _, state, kwargs = _launch(case, fallback, fixed_seq_len=4) + state.copy_(case[3]) + state_before = state.clone() + tuner = mtp_recurrent._mtp_fused_recurrent_gated_delta_rule_autotuned + + tuner.kernel_warmup(tuner._static_key(**kwargs), **kwargs) + + torch.testing.assert_close(state, state_before, atol=0, rtol=0) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py index 11035c83a3..a54ab21ecb 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py @@ -60,6 +60,7 @@ def _run_both( dt_bias=dt_bias, a_raw=a_raw, b_raw=b_raw, + fixed_seq_len=0, ) return o_old, o_new, fs_old, fs_new diff --git a/unit_tests/common/basemodel/triton_kernel/test_gated_rmsnorm.py b/unit_tests/common/basemodel/triton_kernel/test_gated_rmsnorm.py new file mode 100644 index 0000000000..f03fef4629 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/test_gated_rmsnorm.py @@ -0,0 +1,43 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.norm.gated_rmsnorm import ( + gated_rmsnorm_forward, +) + + +@pytest.mark.parametrize("norm_before_gate", [True, False]) +def test_gated_rmsnorm_accepts_strided_3d_gate(norm_before_gate): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for gated RMSNorm test") + + torch.manual_seed(123) + tokens, heads, head_dim = 5, 4, 128 + packed_dim = heads * head_dim + 256 + x = torch.randn((tokens * heads, head_dim), device="cuda", dtype=torch.bfloat16) + weight = torch.randn((head_dim,), device="cuda", dtype=torch.bfloat16) + packed = torch.randn((tokens, packed_dim), device="cuda", dtype=torch.bfloat16) + z = packed[:, 128 : 128 + heads * head_dim].view(tokens, heads, head_dim) + assert not z.is_contiguous() + + run_config = {"BLOCK_N": head_dim, "num_warps": 1} + expected = gated_rmsnorm_forward( + x=x, + weight=weight, + bias=None, + eps=1e-6, + z=z.contiguous().view(-1, head_dim), + norm_before_gate=norm_before_gate, + run_config=run_config, + ) + actual = gated_rmsnorm_forward( + x=x, + weight=weight, + bias=None, + eps=1e-6, + z=z, + norm_before_gate=norm_before_gate, + run_config=run_config, + ) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) diff --git a/unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_greedy.py b/unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_greedy.py new file mode 100644 index 0000000000..279268a85c --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_greedy.py @@ -0,0 +1,65 @@ +import importlib + +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.post_process.greedy_sample import ( + greedy_sample_local_stats, +) + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for Triton kernels") + + +@pytest.mark.parametrize("token_num", [1, 7, 64]) +def test_vocab_parallel_greedy_matches_full_logits(monkeypatch, token_num): + module = importlib.import_module("lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_greedy") + tp_world_size = 4 + local_vocab_size = 8192 + vocab_size = tp_world_size * local_vocab_size + generator = torch.Generator(device="cuda").manual_seed(20260826 + token_num) + local_logits_by_rank = [ + torch.randn( + (local_vocab_size, token_num), + dtype=torch.bfloat16, + device="cuda", + generator=generator, + ) + for _ in range(tp_world_size) + ] + + # Exercise deterministic tie-breaking both between local reduction blocks + # and across tensor-parallel ranks. The smallest global token id must win. + local_logits_by_rank[0][4097, 0] = 20.0 + local_logits_by_rank[0][3, 0] = 20.0 + local_logits_by_rank[3][2, 0] = 20.0 + + local_stats_by_rank = [ + greedy_sample_local_stats(local_logits.transpose(0, 1).contiguous()) for local_logits in local_logits_by_rank + ] + + def fake_all_gather_into_tensor(output_, input_, **_kwargs): + for output, local_stats in zip(output_, local_stats_by_rank): + output.copy_(local_stats) + + monkeypatch.setattr(module, "all_gather_into_tensor", fake_all_gather_into_tensor) + actual_logits, actual_ids, actual_logsumexp = module.vocab_parallel_greedy( + local_logits_by_rank[0], + vocab_size=vocab_size, + tp_world_size=tp_world_size, + group=None, + alloc_func=torch.empty, + ) + actual_ids = actual_ids.view(-1) + actual_logprobs = actual_logits.view(-1) - actual_logsumexp + + full_logits = torch.cat(local_logits_by_rank, dim=0).transpose(0, 1).float() + expected_ids = full_logits.argmax(dim=1) + expected_logits = full_logits.gather(1, expected_ids[:, None]).view(-1) + expected_logsumexp = torch.logsumexp(full_logits, dim=1) + expected_logprobs = torch.log_softmax(full_logits, dim=1).gather(1, expected_ids[:, None]).squeeze(1) + + torch.testing.assert_close(actual_ids, expected_ids, rtol=0, atol=0) + torch.testing.assert_close(actual_logits.view(-1), expected_logits, rtol=0, atol=0) + torch.testing.assert_close(actual_logsumexp, expected_logsumexp, rtol=2e-4, atol=2e-4) + torch.testing.assert_close(actual_logprobs, expected_logprobs, rtol=2e-4, atol=2e-4) diff --git a/unit_tests/models/deepseek3_2/triton_kernel/test_topk_index_to_mem_index.py b/unit_tests/models/deepseek3_2/triton_kernel/test_topk_index_to_mem_index.py index 5e3aaf38d0..a666e48646 100644 --- a/unit_tests/models/deepseek3_2/triton_kernel/test_topk_index_to_mem_index.py +++ b/unit_tests/models/deepseek3_2/triton_kernel/test_topk_index_to_mem_index.py @@ -1,23 +1,35 @@ import torch import pytest -from lightllm.models.deepseek3_2.triton_kernel.topk_index_to_mem_index import trans_topk_index_to_mem_index +from lightllm.models.deepseek3_2.triton_kernel.topk_index_to_mem_index import ( + trans_topk_index_to_mem_index, +) -def test_trans_topk_index_to_mem_index(): +@pytest.mark.parametrize("topk", [2048, 2176]) +def test_trans_topk_index_to_mem_index(topk): """Test trans_topk_index_to_mem_index converts topk indices to memory indices correctly.""" batch_size = 1 - topk = 2048 # Create topk_index tensor with some valid indices and some -1 (padding) topk_index = torch.zeros((batch_size, topk), dtype=torch.int32, device="cuda") - topk_index[:, 0:2048] = torch.arange(0, 2048, dtype=torch.int32, device="cuda") + topk_index[:, 0 : topk - 1] = torch.arange(0, topk - 1, dtype=torch.int32, device="cuda") + topk_index[:, -1] = -1 + ragged_start_index = torch.tensor([2], dtype=torch.int32, device="cuda") # Create ragged_mem_index lookup table - ragged_mem_index = torch.arange(0, 2048, dtype=torch.int32, device="cuda") + 10 - - topk_mem_index = trans_topk_index_to_mem_index(topk_index, ragged_mem_index) - - assert torch.equal(topk_mem_index, (torch.arange(0, 2048, dtype=torch.int32, device="cuda") + 10).view(1, -1)) + ragged_mem_index = torch.arange(0, topk + 2, dtype=torch.int32, device="cuda") + 10 + + topk_mem_index = trans_topk_index_to_mem_index(topk_index, ragged_start_index, ragged_mem_index) + + expected_index = torch.cat( + ( + torch.arange(2, topk + 1, dtype=torch.int32, device="cuda"), + torch.tensor([-1], dtype=torch.int32, device="cuda"), + ) + ).view(1, -1) + expected_mem_index = torch.where(expected_index != -1, expected_index + 10, -1) + assert torch.equal(topk_index, expected_index) + assert torch.equal(topk_mem_index, expected_mem_index) if __name__ == "__main__": diff --git a/unit_tests/models/gemma4/test_post_layer_infer.py b/unit_tests/models/gemma4/test_post_layer_infer.py new file mode 100644 index 0000000000..1d6db329d3 --- /dev/null +++ b/unit_tests/models/gemma4/test_post_layer_infer.py @@ -0,0 +1,99 @@ +from types import SimpleNamespace + +import torch + +import lightllm.models.llama.layer_infer.post_layer_infer as llama_post_layer +from lightllm.models.gemma4.layer_infer.post_layer_infer import Gemma4PostLayerInfer + + +def test_vocab_parallel_greedy_softcaps_local_logits_before_reduction(monkeypatch): + post_layer = Gemma4PostLayerInfer.__new__(Gemma4PostLayerInfer) + post_layer.final_logit_softcapping = 2.0 + post_layer.tp_world_size_ = 2 + post_layer.alloc_tensor = torch.empty + post_layer._norm = lambda hidden, infer_state, layer_weight: hidden + + local_logits = torch.tensor( + [ + [4.0, -4.0], + [2.0, -2.0], + [1.0, -1.0], + ], + dtype=torch.bfloat16, + ) + + class LMHead: + vocab_size = 6 + + def __call__(self, input, alloc_func): + return local_logits + + sparse_logits = torch.tensor([[5.0], [1.0]]) + token_ids = torch.tensor([[5], [1]]) + logsumexp = torch.tensor([5.25, 1.5]) + captured = {} + + def fake_vocab_parallel_greedy(logits, **kwargs): + captured["logits"] = logits + return sparse_logits, token_ids, logsumexp + + monkeypatch.setattr(llama_post_layer, "vocab_parallel_greedy", fake_vocab_parallel_greedy) + + infer_state = SimpleNamespace( + dist_group=None, + use_vocab_parallel_greedy=True, + logits_token_ids=None, + logits_logsumexp=None, + ) + + result = post_layer._lm_head_and_gather( + hidden=torch.empty((2, 3)), + token_num=2, + layer_weight=SimpleNamespace(lm_head_weight_=LMHead()), + infer_state=infer_state, + ) + + expected_logits = ( + torch.tanh(local_logits.float() / post_layer.final_logit_softcapping) * post_layer.final_logit_softcapping + ) + assert captured["logits"].dtype == torch.float32 + torch.testing.assert_close(captured["logits"], expected_logits) + assert result is sparse_logits + assert infer_state.logits_token_ids is token_ids + assert infer_state.logits_logsumexp is logsumexp + + +def test_full_logits_softcap_after_float32_conversion(monkeypatch): + post_layer = Gemma4PostLayerInfer.__new__(Gemma4PostLayerInfer) + post_layer.final_logit_softcapping = 2.0 + post_layer.tp_world_size_ = 1 + post_layer.alloc_tensor = torch.empty + post_layer._norm = lambda hidden, infer_state, layer_weight: hidden + + local_logits = torch.tensor( + [ + [1.234375, -3.140625], + [2.71875, -0.333984375], + [0.10009765625, -1.609375], + ], + dtype=torch.bfloat16, + ) + + class LMHead: + vocab_size = 3 + + def __call__(self, input, alloc_func): + return local_logits + + result = post_layer._lm_head_and_gather( + hidden=torch.empty((2, 3), dtype=torch.bfloat16), + token_num=2, + layer_weight=SimpleNamespace(lm_head_weight_=LMHead()), + infer_state=SimpleNamespace(dist_group=None, use_vocab_parallel_greedy=True), + force_full_logits=True, + ) + + expected = torch.tanh(local_logits.T.float() / post_layer.final_logit_softcapping) + expected *= post_layer.final_logit_softcapping + assert result.dtype == torch.float32 + torch.testing.assert_close(result, expected) diff --git a/unit_tests/models/qwen3next/test_transformer_layer_infer.py b/unit_tests/models/qwen3next/test_transformer_layer_infer.py new file mode 100644 index 0000000000..08334539d1 --- /dev/null +++ b/unit_tests/models/qwen3next/test_transformer_layer_infer.py @@ -0,0 +1,35 @@ +from types import SimpleNamespace + +import torch + +from lightllm.models.qwen3next.layer_infer.transformer_layer_infer import ( + Qwen3NextTransformerLayerInfer, +) + + +def test_linear_post_passes_strided_gate_without_materializing(): + tokens, heads, head_dim = 3, 4, 8 + packed_dim = heads * head_dim + 16 + packed = torch.randn((tokens, packed_dim)) + z = packed[:, 8 : 8 + heads * head_dim].view(tokens, heads, head_dim) + core_attn_out = torch.randn((1, tokens, heads, head_dim)) + calls = [] + + def linear_norm(input, gate_value, eps): + calls.append((input, gate_value, eps)) + return input + + layer_weight = SimpleNamespace( + linear_norm=linear_norm, + linear_out_proj=SimpleNamespace(mm=lambda input: input), + ) + layer_infer = object.__new__(Qwen3NextTransformerLayerInfer) + layer_infer.eps_ = 1e-6 + + output = layer_infer._linear_post(core_attn_out, z, layer_weight) + + assert len(calls) == 1 + assert calls[0][0].shape == (tokens * heads, head_dim) + assert calls[0][1] is z + assert not calls[0][1].is_contiguous() + assert output.shape == (tokens, heads * head_dim) diff --git a/unit_tests/models/test_vocab_parallel_greedy_output.py b/unit_tests/models/test_vocab_parallel_greedy_output.py new file mode 100644 index 0000000000..88e5a7c282 --- /dev/null +++ b/unit_tests/models/test_vocab_parallel_greedy_output.py @@ -0,0 +1,74 @@ +from types import SimpleNamespace + +import torch + +from lightllm.common.basemodel.batch_objs import ModelOutput +from lightllm.models.qwen3_dspark.layer_infer.post_layer_infer import Qwen3DSparkPostLayerInfer +from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend + + +def test_argmax_restores_global_token_ids_and_exact_probabilities(): + backend = ModeBackend.__new__(ModeBackend) + output = ModelOutput( + logits=torch.tensor([[3.0, 1.0], [0.0, 5.0]]), + logits_token_ids=torch.tensor([[30, 10], [100, 500]]), + logits_logsumexp=torch.tensor([4.0, 5.25]), + ) + + token_ids = backend._gen_argmax_token_ids(output) + token_ids_with_prob, probs = backend._gen_argmax_token_ids_and_prob(output) + + torch.testing.assert_close(token_ids, torch.tensor([30, 500])) + torch.testing.assert_close(token_ids_with_prob, token_ids) + torch.testing.assert_close(probs, torch.exp(torch.tensor([-1.0, -0.25]))) + + +def test_dense_argmax_keeps_column_index_semantics(): + backend = ModeBackend.__new__(ModeBackend) + output = ModelOutput(logits=torch.tensor([[1.0, 4.0, 2.0]])) + + torch.testing.assert_close(backend._gen_argmax_token_ids(output), torch.tensor([1])) + + +def test_dspark_confidence_path_receives_global_token_ids(): + post = Qwen3DSparkPostLayerInfer.__new__(Qwen3DSparkPostLayerInfer) + post.block_size_ = 2 + post.markov_rank_ = 0 + post._slice_get_last_input = lambda input_embeddings, infer_state: (input_embeddings, 4) + sparse_logits = torch.tensor([[4.0], [5.0], [7.0], [9.0]]) + sparse_token_ids = torch.tensor([[40], [50], [70], [90]]) + + def gather_vocab_parallel(*args, **kwargs): + infer_state = args[3] + infer_state.logits_token_ids = sparse_token_ids + return sparse_logits + + post._lm_head_and_gather = gather_vocab_parallel + observed = {} + + def predict_confidence(block_hidden, anchor_token_ids, sampled_tokens, layer_weight): + observed["sampled_tokens"] = sampled_tokens + return None + + post.predict_confidence_logits = predict_confidence + + class Collector: + def add_mtp_outputs(self, **kwargs): + self.outputs = kwargs + + collector = Collector() + infer_state = SimpleNamespace( + is_prefill=False, + input_ids=torch.tensor([1, 0, 2, 0]), + logits_token_ids=None, + hidden_collector=collector, + ) + + returned_logits = post.token_forward( + input_embdings=torch.ones((4, 3)), + infer_state=infer_state, + layer_weight=object(), + ) + + torch.testing.assert_close(returned_logits, sparse_logits) + torch.testing.assert_close(observed["sampled_tokens"], torch.tensor([[40, 50], [70, 90]])) diff --git a/unit_tests/server/router/model_infer/mode_backend/test_chunked_prefill_mega_moe_overlap.py b/unit_tests/server/router/model_infer/mode_backend/test_chunked_prefill_mega_moe_overlap.py new file mode 100644 index 0000000000..7d85ad21c5 --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_chunked_prefill_mega_moe_overlap.py @@ -0,0 +1,64 @@ +from lightllm.server.router.model_infer.mode_backend.chunked_prefill import ( + impl as chunked_prefill_impl, +) +from lightllm.server.router.model_infer.mode_backend.chunked_prefill.impl import ( + ChunkedPrefillBackend, +) + + +class _FakeEvent: + def __init__(self, calls=None): + self.recorded = False + self.synchronized = False + self.calls = calls + + def record(self): + self.recorded = True + + def synchronize(self): + self.synchronized = True + if self.calls is not None: + self.calls.append("synchronize") + + +class _FakeEventPack: + def __init__(self, calls): + self.calls = calls + + def notify_forward_and_wait_post_handle(self): + self.calls.append("notify_forward") + + +def test_record_forward_completion_keeps_cpu_bookkeeping_async(monkeypatch): + event = _FakeEvent() + monkeypatch.setattr(chunked_prefill_impl.torch.cuda, "Event", lambda: event) + backend = ChunkedPrefillBackend.__new__(ChunkedPrefillBackend) + backend._serialize_sm90_mega_moe_forwards = True + + assert backend._record_forward_completion() is event + assert event.recorded + assert not event.synchronized + + +def test_mega_moe_waits_before_notifying_next_forward(): + calls = [] + event = _FakeEvent(calls) + event_pack = _FakeEventPack(calls) + backend = ChunkedPrefillBackend.__new__(ChunkedPrefillBackend) + backend._serialize_sm90_mega_moe_forwards = True + + backend._notify_next_forward_when_safe(event_pack, event) + + assert calls == ["synchronize", "notify_forward"] + + +def test_regular_path_notifies_next_forward_before_waiting(): + calls = [] + event = _FakeEvent(calls) + event_pack = _FakeEventPack(calls) + backend = ChunkedPrefillBackend.__new__(ChunkedPrefillBackend) + backend._serialize_sm90_mega_moe_forwards = False + + backend._notify_next_forward_when_safe(event_pack, event) + + assert calls == ["notify_forward", "synchronize"] diff --git a/unit_tests/server/router/model_infer/mode_backend/test_vocab_parallel_greedy_sampling.py b/unit_tests/server/router/model_infer/mode_backend/test_vocab_parallel_greedy_sampling.py new file mode 100644 index 0000000000..fddcda9380 --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_vocab_parallel_greedy_sampling.py @@ -0,0 +1,91 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.batch_objs import ModelOutput +from lightllm.utils.envs_utils import enable_env_vars +from lightllm.server.router.model_infer.mode_backend.generic_post_process import ( + _can_use_unmodified_greedy_logits, + can_use_vocab_parallel_greedy, + sample, +) + + +def make_req(**overrides): + values = { + "top_k": 1, + "temperature": 1.0, + "presence_penalty": 0.0, + "frequency_penalty": 0.0, + "repetition_penalty": 1.0, + "min_new_tokens": 1, + "decay_factor": 1.0, + "invalid_token_ids": [], + "output_len": 0, + } + values.update(overrides) + shm_param = SimpleNamespace( + top_k=values["top_k"], + temperature=values["temperature"], + presence_penalty=values["presence_penalty"], + frequency_penalty=values["frequency_penalty"], + repetition_penalty=values["repetition_penalty"], + min_new_tokens=values["min_new_tokens"], + exponential_decay_length_penalty=SimpleNamespace(to_tuple=lambda: (1, values["decay_factor"])), + ) + input_len = 10 + return SimpleNamespace( + sampling_param=SimpleNamespace( + shm_param=shm_param, + invalid_token_ids=values["invalid_token_ids"], + ), + shm_req=SimpleNamespace(input_len=input_len), + get_cur_total_len=lambda: input_len + values["output_len"], + ) + + +def test_accepts_unmodified_greedy_requests(): + assert _can_use_unmodified_greedy_logits([make_req(), make_req(output_len=5)]) + + +def test_feature_gate_requires_environment_and_eligible_batch(monkeypatch): + monkeypatch.delenv("LIGHTLLM_VOCAB_PARALLEL_GREEDY", raising=False) + enable_env_vars.cache_clear() + assert not can_use_vocab_parallel_greedy([make_req()]) + + monkeypatch.setenv("LIGHTLLM_VOCAB_PARALLEL_GREEDY", "1") + enable_env_vars.cache_clear() + assert can_use_vocab_parallel_greedy([make_req()]) + assert not can_use_vocab_parallel_greedy([make_req(top_k=2)]) + enable_env_vars.cache_clear() + + +def test_samples_explicit_vocab_parallel_output_exactly(): + model_output = ModelOutput( + logits=torch.tensor([[9.0], [4.0]], dtype=torch.float32), + logits_token_ids=torch.tensor([[17], [3]], dtype=torch.int64), + logits_logsumexp=torch.tensor([9.25, 5.5], dtype=torch.float32), + ) + + token_ids, token_logprobs = sample(model_output, [make_req(), make_req()]) + + torch.testing.assert_close(token_ids, torch.tensor([17, 3])) + torch.testing.assert_close(token_logprobs, torch.tensor([-0.25, -1.5])) + + +@pytest.mark.parametrize( + "override", + [ + {"top_k": 2}, + {"temperature": 0.5}, + {"presence_penalty": 0.1}, + {"frequency_penalty": 0.1}, + {"repetition_penalty": 1.1}, + {"decay_factor": 1.1}, + {"invalid_token_ids": [7]}, + {"min_new_tokens": 2}, + ], +) +def test_rejects_logits_modifiers(override): + assert not _can_use_unmodified_greedy_logits([make_req(**override)]) diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py b/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py index 7baf34061b..9b6995a6b6 100644 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py +++ b/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py @@ -739,7 +739,7 @@ def test_lightspec_eagle_draft_always_keeps_the_extend_candidate(): assert plan.draft_step >= 1 -def test_vanilla_with_attention_planner_prices_extend_then_normal_batches(): +def test_vanilla_with_attention_planner_prices_every_chained_level_at_verify_batch(): planner = build_lightspec_planner(spec_mode="vanilla_with_att") planner.draft_infer_costs.update(batch_size=2, infer_cost_ms=0.1) @@ -749,7 +749,7 @@ def test_vanilla_with_attention_planner_prices_extend_then_normal_batches(): draft_step=3, ) - assert np.isclose(draft_cost_ms, 0.8) + assert np.isclose(draft_cost_ms, 1.2) with pytest.raises(AssertionError, match="requires draft_step to be greater than 0"): planner._get_draft_cost_ms(req_num=4, verify_batch_size=8, draft_step=0) diff --git a/unit_tests/server/router/test_dp_model_capacity.py b/unit_tests/server/router/test_dp_model_capacity.py new file mode 100644 index 0000000000..236e368372 --- /dev/null +++ b/unit_tests/server/router/test_dp_model_capacity.py @@ -0,0 +1,39 @@ +from types import SimpleNamespace + +import pytest + +from lightllm.server.router.manager import resolve_model_max_req_num + + +def test_model_request_capacity_defaults_to_global_running_limit(): + args = SimpleNamespace( + running_max_req_size=256, + per_dp_max_req_size=None, + dp=8, + nnodes=1, + ) + + assert resolve_model_max_req_num(args) == 256 + + +def test_model_request_capacity_accepts_balanced_per_dp_limit(): + args = SimpleNamespace( + running_max_req_size=256, + per_dp_max_req_size=40, + dp=8, + nnodes=1, + ) + + assert resolve_model_max_req_num(args) == 40 + + +def test_model_request_capacity_rejects_less_than_balanced_share(): + args = SimpleNamespace( + running_max_req_size=256, + per_dp_max_req_size=31, + dp=8, + nnodes=1, + ) + + with pytest.raises(ValueError, match="minimum 32"): + resolve_model_max_req_num(args) diff --git a/unit_tests/server/router/test_prefill_coalescing.py b/unit_tests/server/router/test_prefill_coalescing.py new file mode 100644 index 0000000000..0c0c1bec4a --- /dev/null +++ b/unit_tests/server/router/test_prefill_coalescing.py @@ -0,0 +1,69 @@ +from types import SimpleNamespace + +from lightllm.server.router import manager as router_manager +from lightllm.server.router.manager import RouterManager + + +class _ReqQueue: + def __init__(self, waiting_req_num): + self.waiting_req_num = waiting_req_num + + def get_wait_req_num(self): + return self.waiting_req_num + + +def _make_router(waiting_req_num, interval=0.5, running_req_num=0, pending_req_num=0): + router = RouterManager.__new__(RouterManager) + router.args = SimpleNamespace(running_max_req_size=64) + router.prefill_coalesce_interval = interval + router._prefill_coalesce_deadline = None + router.req_queue = _ReqQueue(waiting_req_num) + router.running_batch = None if running_req_num == 0 else SimpleNamespace(reqs=[None] * running_req_num) + router.schedule_new_batch = None if pending_req_num == 0 else SimpleNamespace(reqs=[None] * pending_req_num) + return router + + +def test_prefill_coalescing_is_disabled_by_default(): + router = _make_router(waiting_req_num=1, interval=0.0) + + assert not router._should_defer_prefill_batch() + + +def test_partial_burst_waits_only_until_original_deadline(monkeypatch): + now = [10.0] + monkeypatch.setattr(router_manager.time, "monotonic", lambda: now[0]) + router = _make_router(waiting_req_num=12) + + assert router._should_defer_prefill_batch() + assert router._prefill_coalesce_deadline == 10.5 + + now[0] = 10.49 + router.req_queue.waiting_req_num = 48 + assert router._should_defer_prefill_batch() + assert router._prefill_coalesce_deadline == 10.5 + + now[0] = 10.5 + assert not router._should_defer_prefill_batch() + assert router._prefill_coalesce_deadline is None + + +def test_full_runnable_batch_skips_coalescing(monkeypatch): + monkeypatch.setattr(router_manager.time, "monotonic", lambda: 10.0) + router = _make_router(waiting_req_num=48, running_req_num=16) + + assert not router._should_defer_prefill_batch() + assert router._prefill_coalesce_deadline is None + + +def test_expired_burst_launches_as_soon_as_running_batch_clears(monkeypatch): + now = [10.0] + monkeypatch.setattr(router_manager.time, "monotonic", lambda: now[0]) + router = _make_router(waiting_req_num=40, running_req_num=64) + + assert router._should_defer_prefill_batch() + now[0] = 11.0 + assert router._should_defer_prefill_batch() + + router.running_batch = None + assert not router._should_defer_prefill_batch() + assert router._prefill_coalesce_deadline is None diff --git a/unit_tests/server/test_mtp_start_args.py b/unit_tests/server/test_mtp_start_args.py index 87e12479d4..0e5504bc98 100644 --- a/unit_tests/server/test_mtp_start_args.py +++ b/unit_tests/server/test_mtp_start_args.py @@ -4,11 +4,11 @@ from lightllm.server.core.objs.start_args_type import StartArgs -def test_mtp_requires_cuda_graph(monkeypatch): +def test_dynamic_mtp_requires_cuda_graph(monkeypatch): monkeypatch.setattr("lightllm.server.api_start._set_envs_and_config", lambda args: None) - args = StartArgs(mtp_mode="vanilla_no_att", disable_cudagraph=True) + args = StartArgs(mtp_mode="vanilla_no_att", mtp_dynamic_verify=True, disable_cudagraph=True) - with pytest.raises(AssertionError, match="only supported on Prefill nodes"): + with pytest.raises(AssertionError, match="--mtp_dynamic_verify"): _launch_subprocesses(args) diff --git a/unit_tests/utils/test_envs_utils.py b/unit_tests/utils/test_envs_utils.py new file mode 100644 index 0000000000..55e95bc3d5 --- /dev/null +++ b/unit_tests/utils/test_envs_utils.py @@ -0,0 +1,74 @@ +from easydict import EasyDict + +from lightllm.utils import envs_utils + + +def _set_start_args(monkeypatch, **kwargs): + args = EasyDict( + graph_max_batch_size=kwargs.get("graph_max_batch_size", 256), + running_max_req_size=kwargs.get("running_max_req_size", 256), + mtp_mode=kwargs.get("mtp_mode"), + mtp_step=kwargs.get("mtp_step", 0), + enable_tpsp_mix_mode=kwargs.get("enable_tpsp_mix_mode", False), + enable_ep_moe=kwargs.get("enable_ep_moe", False), + tp=kwargs.get("tp", 1), + ) + monkeypatch.setattr(envs_utils, "get_env_start_args", lambda: args) + envs_utils.get_deepep_num_max_dispatch_tokens_per_rank_decode.cache_clear() + + +def test_deepep_decode_limit_covers_speculative_physical_batch(monkeypatch): + monkeypatch.delenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE", raising=False) + _set_start_args( + monkeypatch, + graph_max_batch_size=64, + running_max_req_size=64, + mtp_mode="eagle_with_att", + mtp_step=5, + ) + + assert envs_utils.get_deepep_num_max_dispatch_tokens_per_rank_decode() == 384 + + +def test_deepep_decode_limit_uses_sequence_parallel_local_batch(monkeypatch): + monkeypatch.delenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE", raising=False) + _set_start_args( + monkeypatch, + graph_max_batch_size=64, + running_max_req_size=64, + mtp_mode="eagle_with_att", + mtp_step=5, + enable_tpsp_mix_mode=True, + enable_ep_moe=True, + tp=8, + ) + + assert envs_utils.get_deepep_num_max_dispatch_tokens_per_rank_decode() == 256 + + +def test_deepep_decode_limit_is_aligned_and_keeps_minimum(monkeypatch): + monkeypatch.delenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE", raising=False) + _set_start_args( + monkeypatch, + graph_max_batch_size=63, + running_max_req_size=32, + mtp_mode="eagle_with_att", + mtp_step=4, + ) + assert envs_utils.get_deepep_num_max_dispatch_tokens_per_rank_decode() == 320 + + _set_start_args(monkeypatch, graph_max_batch_size=8, running_max_req_size=8) + assert envs_utils.get_deepep_num_max_dispatch_tokens_per_rank_decode() == 256 + + +def test_deepep_decode_limit_honors_explicit_override(monkeypatch): + monkeypatch.setenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE", "512") + _set_start_args( + monkeypatch, + graph_max_batch_size=64, + running_max_req_size=64, + mtp_mode="eagle_with_att", + mtp_step=5, + ) + + assert envs_utils.get_deepep_num_max_dispatch_tokens_per_rank_decode() == 512