[ROCm][Kimi-K3] Add moonmath kernels for gfx942 MXFP4 SiTU MoE and MLA decode, with decode context parallelism - #57412
tarik-sarac wants to merge 5 commits into
Conversation
…#56723) Net diff of vllm-project#56723 at 903e47e, applied unchanged as the base for the ROCm moonmath DCP change; it is not part of that PR. Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
…oE and MLA decode Signed-off-by: Tarik Sarac <tariksarac120@gmail.com>
…rallelism Builds on vllm-project#56723, which runs the dense DSpark drafter at DCP=1 next to a DCP target. On gfx942 the moonmath kernel now serves decode and verify for a DCP target: each rank attends its own KV shard, with the causal limit taken from the global lengths, and returns its LSE, converted from base 2 to the natural log the merge expects. - moonmath_mla: accept the fused query DCP gathers, pass the global lengths and CP rank, and return the LSE. - rocm_aiter_mla: segmented_dcp_verify lets a subclass keep varlen DCP verify on the flat per-token view instead of AITER's segmented metadata. - Layout: the AITER backends have no block-outer KV layout, so a replicated draft next to a sharded target falls back to padded pages rather than failing. Signed-off-by: Tarik Sarac <tariksarac120@gmail.com>
…in fp32 - MLA: the AITER MLA builder gets _decode_reads_aiter_schedule. The moonmath builder skips get_mla_metadata_v1 for batches its kernel serves, using the same routing predicate as the impl (5 fewer launches per decode step). - MoE: the down GEMM no longer applies a bf16 copy of the routing weight; one Triton kernel sums each token's top-k rows weighted in fp32 (2 fewer kernels per MoE layer, and the weight keeps the router's precision). Signed-off-by: Tarik Sarac <tariksarac120@gmail.com>
AITER MLA decode, and moonmath's decode built on it, flattens each layer's cache with .view(), which fails when the layout does not keep the layer's pages contiguous. The backend did not declare that, so under DCP with a replicated drafter whose backend has a block-outer layout (e.g. TRITON_ATTN), vllm-project#56723's layout preference picked BLNHC and decode crashed at startup. Declare the layer-contiguous layouts, as the sparse AITER MLA backend does. Signed-off-by: Tarik Sarac <tariksarac120@gmail.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
|
This pull request has merge conflicts that must be resolved before it can be |
2 participants
Stacked on #56723: the first commit is its net diff at
903e47e4fd, unchanged. This PR isthe four commits after it (17 files, +862/-33). E2E results also use #56434 (ROCm drafter
hidden-state fix) and, for the
TRITON_ATTNdrafter, #57409.Purpose
Kimi-K3 on AMD MI325X (gfx942) has no kernel for its MXFP4 SiTU MoE, no bf16-query / fp8-KV MLA
decode for its 12 heads per rank, and no MLA decode that works under decode context parallelism
(DCP). This PR adds them from the optional
moonmath_amdpackage; without it nothing changes.moonmath_mxfp4_moe.py):MoonmathW4A16SituExpertssubclasses the AITER W4A16experts and claims only SiTU. The routing weight is applied in an fp32 top-k reduce.
AITER_TRITON_MXFP4_BF16aligns SiTU to 128, since the generic 256 round-up adds 33% MoE memory.moonmath_mla.py,ROCM_MOONMATH_MLA): subclasses the AITER MLA backend and replacesonly decode and verify; batches outside the kernel's domain go to AITER.
K3 with [PCP][DCP] Support DCP target model with non-DCP Dspark #56723 on ROCm:
resolve_kv_cache_layoutkeeps a layer-outer layout (padding the draft pages) when nobackend has a block-outer one, instead of failing.
AiterMLABackenddeclares layer-contiguous layouts: its decode.view()s each layer's cache,which crashed under
BLNHCwith aTRITON_ATTNdrafter.Why this is not a duplicate
(int4), [Feature] Enable AITER MXFP4 MoE on gfx942 and optimize tile configurations for MI325X Target Kimi K3 running on MI325X #50817 (AITER tile configs).
cover
ROCM_AITER_MLAwith an MLA drafter; [Spec decode] Support Kimi-K3 DCP with DSpark #52188 is NVIDIA.Test Plan
Unit tests on MI325X with ROCm nightly
0bfc7a15d0, as shipped vs with this branch: the test filesthis PR and #56723 touch, plus the AITER MLA / AITER_FA tests.
E2E: Kimi-K3, TP8 + DCP8, c=8, same-prompt decode differential (first probe discarded).
Test Result
Unit: no test goes from pass to fail, and all 99 new or DCP-specific tests pass. The 4 failures
are identical in both arms (a model not cached offline).
E2E (backends
ROCM_MOONMATH_MLAandMoonmathW4A16SituExperts; step = 8 x acceptance / tok/s):ROCM_AITER_FATRITON_ATTN(#57409)GPU KV cache at DCP8: 5.08M tokens.