Skip to content

perf(rdna2): opt-in MoE skinny v2 decode, M=16 decode kernels, small-op pass (port of leapdragon) - #57

Draft
BlivionIaG wants to merge 1 commit into
rdna_extrasfrom
port/leap-moe-decode-v2
Draft

BlivionIaG wants to merge 1 commit into
rdna_extrasfrom
port/leap-moe-decode-v2

Conversation

@BlivionIaG

Copy link
Copy Markdown

Ports: leapdragon/vllm-rdna2-qwen fd8d939 (MoE skinny v2, decode kernels to M=16) and a80aad3 (decode small-op pass), author Aron Hsiao leapdragon@gmail.com. Adapted to rdna_extras. Everything is gfx10x-only and opt-in; defaults are unchanged.

Changes

Flag (default 0) What
VLLM_RDNA_MOE_V2=1 moe_skinny_int4_decode_v2: no LDS activation staging, 16-byte weight loads, exact natural-order int4 dequant into v_dot2, EP-local slot compaction in w2, optional residual epilogue. Same sequential moe_wna16 layout as v1. The _custom_ops wrapper falls back to v1 when the v2 contract does not hold (alignment, K/N % 32, K <= 4096, group % 32, topk <= 32).
VLLM_RDNA_DECODE_M16=1 rdna_dense_gemm (int8 shadow GEMV) takes batches up to 16, and the fused glue (rdna_gemv_act, rdna_hc_up_gate_mix, rdna_se_*) takes M == 16. The kernels always accept M=16 now (gemv_i8_rdna2: 16 native, 9..15 as 8 + rest); only the Python caps are gated. gemv_f16_rdna2 routing is unchanged: it already guards 9..32, and our crossover bench has rocBLAS ahead at M=16-24.
VLLM_RDNA_MOE_ALIAS=1 The modular MoE kernel writes into the caller's fresh output on gfx10x, as upstream does for aiter. One device copy less per MoE layer.
VLLM_RDNA_HC_V2=1 Split-K int8 rdna_gemv_act for N < 1152, K >= 4096 (HC down projection). Read in C++.

MoE batches above 8 still need VLLM_ROCM_MOE_SKINNY_MAX_M=16 (existing knob).

Files: csrc/rocm/{skinny_gemms_int4.cu,rdna_fused_glue.cu,ops.h,torch_bindings.cpp}, vllm/_custom_ops.py, vllm/envs.py, rdna_ops.py, rdna_dense_int8.py, fused_moe/modular_kernel.py, plus the new tests/kernels/quantization/test_rdna_decode_m16.py (CPU predicates; GPU kernel checks ported from his tools/rdna2/{skinny_m16,glue_m16,moe_v2}_test.py).

His numbers (his box, 4x V620 TP=4/EP, his bench, not ours)

  • One MoE layer at M=1: 88 -> 71 us; -36 % at M=8/16.
  • Single-stream decode: 62.4 -> 64.65 t/s (+3.6 %).
  • Aggregate at 1/2/4/8/16 streams: 64/107/170/234/265 t/s, vs 63/104/156/206/134 before the M=16 work.
  • Small-op pass: +1.3 % decode (66.9 -> 67.7 t/s).
  • Greedy output vs v1: 6/8 prompts identical, divergences at near-ties.

Checks run here

  • ruff 0.14.0 check + format --check on the changed Python; clang-format 21.1.2 on the changed C++ (only the new hunks kept).
  • Op registration reviewed: ops.h decl, torch_bindings.cpp def/impl, _custom_ops wrapper. The op mutates its outputs and returns (), the same pattern as v1, so it needs no fake.
  • Not compiled (no ROCm toolchain here). pytest was not run: importing vllm in the allowed env fails on missing regex.

GPU validation to run

  • Build, then run pytest tests/kernels/quantization/test_rdna_decode_m16.py tests/kernels/quantization/test_rocm_moe_skinny.py on gfx1030.
  • Correctness for each flag and all together: greedy probes, PPL, prefix, after-load and concurrency probes, 0 faults/dmesg.
  • Bench Flash-Next MTP0 and MTP2, all flags off vs MOE_V2=1 vs all four on (+ VLLM_ROCM_MOE_SKINNY_MAX_M=16): 1k/512 and 16k/1k at c=1 and c=8. Report prefill tok/s, decode tok/s/req, TTFT s and ITL ms for each.
  • DECODE_M16 and HC_V2 only matter with int8 shadows (perf(rdna2): int8 dense shadows on in the Flash-Next recipes (port of leapdragon T45) #56), so bench them on top of VLLM_RDNA_DENSE_INT8=1.

…op pass

Ported from leapdragon/vllm-rdna2-qwen@fd8d939ae and @a80aad336 (Aron
Hsiao), adapted to rdna_extras. Everything is gfx10x-only and opt-in;
defaults are unchanged.

- moe_skinny_int4_decode_v2 (VLLM_RDNA_MOE_V2=1): no LDS activation
  staging, 16-byte weight loads, exact natural-order int4 dequant into
  v_dot2, EP-local slot compaction in w2, optional residual epilogue.
  The Python wrapper falls back to v1 when the v2 shape/alignment
  contract does not hold. Same sequential moe_wna16 layout as v1.
- gemv_i8_rdna2 takes M up to 16 (16 native, 9..15 as 8 + rest); the
  fused glue ops (rdna_gemv_act, rdna_hc_up_gate_mix, rdna_se_*) take
  M == 16. The Python decode caps move to 16 only with
  VLLM_RDNA_DECODE_M16=1. gemv_f16_rdna2 routing is unchanged (it
  already guards 9..32, and rocBLAS wins at M=16-24 on our tuned rows).
- VLLM_RDNA_MOE_ALIAS=1: the modular MoE kernel writes into the
  caller's fresh output on gfx10x, as upstream does for aiter (one
  device copy less per MoE layer).
- VLLM_RDNA_HC_V2=1: split-K int8 rdna_gemv_act for N < 1152,
  K >= 4096 (HC down projection).

Tests: tests/kernels/quantization/test_rdna_decode_m16.py (CPU
predicates; GPU kernel checks ported from his tools/rdna2 tests).

Co-authored-by: Aron Hsiao <leapdragon@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant