Skip to content

Kimi-K3 FP8 MI355X vLLM agentic recipe (base + DSpark MTP) - #2585

Draft
xiaohuguo2023 wants to merge 1 commit into
mainfrom
xiaohuguo/kimik3-fp8-mi355x-vllm-recipe
Draft

Kimi-K3 FP8 MI355X vLLM agentic recipe (base + DSpark MTP)#2585
xiaohuguo2023 wants to merge 1 commit into
mainfrom
xiaohuguo/kimik3-fp8-mi355x-vllm-recipe

Conversation

@xiaohuguo2023

Copy link
Copy Markdown
Collaborator

Description

ROCm/MI355X (gfx950, TP8) recipe for the Kimi-K3 agentic trace-replay
benchmark via vLLM + AITER — the AMD counterpart to
benchmarks/single_node/agentic/kimik3_fp4_b300_vllm*.sh. Same AIPerf cc-traces
harness and dataset as the B300 baseline, so the MI355X points drop onto the same
Pareto apples-to-apples.

Two recipes (benchmarks/single_node/agentic/):

  • kimik3_fp4_mi355x_vllm.sh — base (non-spec).
  • kimik3_fp4_mi355x_vllm_mtp.sh — DSpark speculative decoding (K=2, synthetic
    golden AL 2.51); the arm comparable to kimik3_fp4_b300_vllm_mtp.sh.

Precision. Mixed AITER path — MoE experts A8W4 (AITER_SITUV2_A8W4, 4-bit
weights / 8-bit activations), dense MLA attention in fp8 ASM (ROCM_AITER_MLA
asm-persistent), fp8 KV cache. Not a uniform NVFP4 checkpoint. Recipe filenames
keep the fp4_* stem for parity with the B300 sibling; the branch/title say FP8 to
reflect the dominant activation/attention path. (Happy to make these consistent
either way — see open question below.)

Reproduce-from-base-image. No prebuilt image and no copy of the NV waiver
layout. The image pinned in configs/amd-master.yaml is the stock ROCm vLLM
nightly
(cb8104839c...); apply_k3_container_patches.sh idempotently patches it
in-container at startup:

  • aiter rebuild @ pin (asserts #4579 / #4575), tuned GEMM CSV merge, triton 3.7.0;
  • five vLLM ASM-MLA patches — 12→16 head pad for decode + prefill, 16-head PS
    metadata, fp8-PS skip, wvsplitk;
  • the DSpark draft layer + FlyDSL→torch decode reroute.

No-op once markers are present; SKIP_K3_BOOTSTRAP=1 for a pre-baked image.

Mandated serving config (unchanged from the validated sweep): gpu-mem 0.95,
max-num-seqs 64, MNBT 16384, FULL_AND_PIECEWISE cudagraphs with {12,36} capture
sizes (lands conc-4/12 DSpark decode on FULL graphs, removing the
get_mla_metadata_v1 host bubble), ROCM_AITER_MLA asm padding, KV-offload via
native / vllm-simple backends.

Config. configs/amd-master.yaml adds kimik3-fp4-mi355x-vllm-agentic (base)
and kimik3-fp4-mi355x-vllm-agentic-dspark (spec-decoding: mtp) blocks; both
validate via generate_sweep_configs test-config.

Signed-off-by: Xiaohu Guo <Xiaohu.Guo@amd.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Status: No status

Development

Successfully merging this pull request may close these issues.

1 participant