[Issue]: RuntimeError with ATOM docker of rocm/atom:rocm7.2.1-ubuntu24.04-pytorch2.9.1-atom0.1.2 for ISL/OSL of 1k/1k and conc64 due to dtype mismatch of two matrices in one gemm kernel call.

Problem Description

RuntimeError with ATOM docker of rocm/atom:rocm7.2.1-ubuntu24.04-pytorch2.9.1-atom0.1.2 for ISL/OSL of 1k/1k and conc64 due to dtype mismatch of two matrices in one gemm kernel call.

Model: https://huggingface.co/amd/Llama-3.1-8B-Instruct-FP8-KV

We observed the following RuntimeError with ATOM docker of rocm/atom:rocm7.2.1-ubuntu24.04-pytorch2.9.1-atom0.1.2 for ISL/OSL of 1k/1k and conc64 due to dtype mismatch of two matrices in one gemm kernel call.

Docker: rocm/atom:rocm7.2.1-ubuntu24.04-pytorch2.9.1-atom0.1.2

RuntimeError: expected mat1 and mat2 to have the same dtype, but got: c10::BFloat16 != c10::Float8_e4m3fn

How to reproduce this issue?

The link of the CI/CD run is here: Benchmark | OOB | smci355-ccs-aus-n10-25 | rocm/atom-dev:vllm-latest | Llama-3-1-8B-Instruct-FP8-KV | atom | warmup · AMD-MLPerf/workloads-inference@1153c26

Operating System

ubuntu24.04

CPU

AMD EPYC 9575F 64-Core Processor x 2

GPU

AMD Instinct MI355x

ROCm Version

ROCm7.2.1

ROCm Component

No response

Steps to Reproduce

How to reproduce this issue?

The link of the CI/CD run is here: Benchmark | OOB | smci355-ccs-aus-n10-25 | rocm/atom-dev:vllm-latest | Llama-3-1-8B-Instruct-FP8-KV | atom | warmup · AMD-MLPerf/workloads-inference@1153c26

(Optional for Linux users) Output of /opt/rocm/bin/rocminfo --support

No response

Additional Information

No response

Start with the linked CI/CD run and the rocm/atom:rocm7.2.1-ubuntu24.04-pytorch2.9.1-atom0.1.2 environment, using the amd/Llama-3.1-8B-Instruct-FP8-KV model with ISL/OSL 1k/1k and concurrency 64. Trace the GEMM kernel call that reports c10::BFloat16 versus c10::Float8_e4m3fn, and consider the issue complete when the workload runs without the dtype mismatch.