NOSSLETTER.techLIVE
back to archive

[ROCm][Perf][Muse Spark] FlyDSL decode monokernel for DSpark verify steps on MI355X (--monokernel)

Language: Python+8985 -3performance
concepts demonstrated
GPU kernel launchesPerformance optimizationCommand-line flagsBenchmarking and accuracy evaluation
The Issue & Context

Before this PR, Muse Spark on 8x MI355X GPUs ran decoder layers using multiple kernel launches per layer, limiting throughput and interactivity especially for multi-request workloads with DSpark verification steps.

Root Cause Analysis

The existing decoder implementation used separate kernel launches for attention output, query, output gate, and KV insert steps within each decoder layer, causing overhead and underutilization of GPU resources on MI355X hardware when running DSpark workloads with many requests.

Solution Strategy & Architectural Pattern

This PR introduces a new FlyDSL-based monokernel decoder path that fuses multiple decoder steps (layers 1-60) into persistent kernel launches covering attention output through KV insert, implemented in the vllm/models/muse_spark/amd/mono directory (notably kernels/k3.py, k5.py, decoder.py). It is opt-in via the --monokernel flag and falls back to stock for unsupported cases. The monokernel approach reduces kernel launch overhead and improves GPU utilization, validated by comprehensive benchmarks and accuracy tests.

Key Takeaway for Contributors

Fusing multiple GPU kernel steps into a single persistent launch can significantly improve throughput and latency by reducing launch overhead and better utilizing hardware, especially on complex multi-step workloads like transformer decoding.