在 AMD MI300X 上跑 DeepSeek-V4-Flash：FP8、Triton 与 AITER 的真实踩坑记录

import torch
import os

def recalibrate_fp8_scales(weight_path, target_device="cuda:0"):
    """
    加载 DeepSeek-V4-Flash 的 FP8 权重，
    将 per-row scale 转换为 AITER 所需的 per-block scale。
    实际部署中 block_size 通常为 128 或 256，
    需根据 AITER 内核版本确认。
    """
    state = torch.load(weight_path, map_location=target_device)
    recalibrated = {}

    block_size = 128  # AITER v0.4 默认块大小，请按实际版本调整

    for name, tensor in state.items():
        if tensor.dtype == torch.float8_e4m3fn:
            # 原始 per-row scale 存在相邻 key 中，命名约定为 xxx.scale
            scale_key = name.replace(".weight", ".weight_scale_inv")
            if scale_key in state:
                row_scales = state[scale_key].float()  # shape: [num_rows]
                num_rows = row_scales.shape[0]

                # 按 block_size 分块，取块内最大 scale 作为块 scale
                padded = torch.nn.functional.pad(
                    row_scales,
                    (0, block_size - num_rows % block_size),
                    value=1.0
                )
                blocked = padded.view(-1, block_size)
                block_scales = blocked.max(dim=1).values  # shape: [num_blocks]

                # 用块 scale 重新量化权重
                weight_fp32 = tensor.float()
                weight_requant = (weight_fp32 * row_scales.unsqueeze(1))
                weight_requant = weight_requant / block_scales.unsqueeze(1)
                recalibrated[name] = weight_requant.to(torch.float8_e4m3fn)
                recalibrated[scale_key.replace("_inv", "_block_inv")] = block_scales
            else:
                recalibrated[name] = tensor
        else:
            recalibrated[name] = tensor

    return recalibrated

# 使用示例（MI300X 上 ROCm 的 device 标记仍为 cuda）
weights = recalibrate_fp8_scales(
    "deepseek-v4-flash/model.fp8.safetensors",
    target_device="cuda:0"
)

# 查看当前 MI300X 的 GCD 拓扑
rocm-smi --showtopo

# 将推理进程绑定到 GCD 0 和 GCD 1 的计算队列
# 方法一：通过环境变量限制可见设备（每个 GCD 视为独立 device）
export HIP_VISIBLE_DEVICES=0  # 单 GCD 模式，牺牲一半内存但避免死锁

# 方法二：在代码中手动分配 kernel 到不同 GCD
# Triton 目前不支持直接指定 GCD，需要通过 HIP stream 间接控制

import torch

def split_gcd_streams():
    """
    在双 GCD 的 MI300X 上创建两个 HIP stream，
    将 attention kernel 分流到不同 GCD。
    注意：此方案依赖 ROCm 6.2+ 的多 stream 支持，
    且需要 Triton kernel 通过 stream 参数提交。
    """
    device = torch.device("cuda:0")
    # MI300X 双 GCD 在 ROCm 中暴露为同一 device 的不同 stream
    stream_gcd0 = torch.cuda.Stream(device=device, priority=0)
    stream_gcd1 = torch.cuda.Stream(device=device, priority=0)

    return stream_gcd0, stream_gcd1

# 在推理循环中使用
stream0, stream1 = split_gcd_streams()

# 假设你有 64 个 attention head，分两组
with torch.cuda.stream(stream0):
    # 前 32 个 head 的 Triton kernel 在此提交
    attention_group0 = triton_attention(q[:32], k[:32], v[:32])

with torch.cuda.stream(stream1):
    # 后 32 个 head 的 Triton kernel 在此提交
    attention_group1 = triton_attention(q[32:], k[32:], v[64:])

# 同步两个 stream
torch.cuda.synchronize()
output = torch.cat([attention_group0, attention_group1], dim=0)

# 在 Triton kernel 定义中调整参数
@triton.jit
def flash_attn_kernel(
    Q_ptr, K_ptr, V_ptr, O_ptr,
    stride_qm, stride_kn, stride_vn, stride_om,
    BLOCK_M: tl.constexpr = 64,
    BLOCK_N: tl.constexpr = 64,
    # 关键调整：MI300X 上增加 NUM_STAGES 提升寄存器复用
    NUM_STAGES: tl.constexpr = 4,  # H100 默认 2，MI300X 建议 3-4
    NUM_WARPS: tl.constexpr = 8,   # H100 默认 4，MI300X 建议 8
):
    # kernel body 略
    pass

"""DeepSeek-V4-Flash 推理框架的 AITER shim 层
将 NVIDIA 内核接口映射到 AITER 内核接口。
假设 AITER 已通过 pip install aiter 安装，
且 ROCm 版本 >= 6.2。
"""
import torch
import aiter

def flash_attn_varlen_aiter(
    q, k, v,
    cu_seqlens_q, cu_seqlens_k,
    max_seqlen_q, max_seqlen_k,
    causal=True,
    sm_scale=None,
):
    """
    将 DeepSeek 的 flash_attn_varlen 调用
    转换为 AITER 的 aiter_flash_attn 调用。
    """
    batch_size = cu_seqlens_q.shape[0] - 1

    # 从 cu_seqlens 反推每个序列的实际长度
    seq_lens = cu_seqlens_q[1:] - cu_seqlens_q[:-1]

    # AITER 要求的 block_table：每个序列的 KV block 映射
    # 对于非 PagedAttention 场景，block_table 是简单的连续映射
    block_size = 64  # AITER 默认 KV block 大小
    num_blocks_per_seq = (max_seqlen_k + block_size - 1) // block_size
    total_blocks = num_blocks_per_seq * batch_size

    block_table = torch.arange(
        total_blocks, device=q.device, dtype=torch.int32
    ).view(batch_size, num_blocks_per_seq)

    # 调用 AITER
    output = aiter.flash_attn(
        q, k, v,
        seq_lens=seq_lens,
        block_table=block_table,
        causal=causal,
        sm_scale=sm_scale or 1.0 / (q.shape[-1] ** 0.5),
    )
    return output


def rmsnorm_add_aiter(x, residual, weight, eps=1e-6):
    """
    将 NVIDIA 的 rmsnorm_with_residual_add
    映射到 AITER 的 aiter_rmsnorm。
    关键：关闭 inplace 模式以保留 residual 原始值。
    """
    # AITER 的 rmsnorm 返回归一化结果，
    # 残差加法需要手动完成（因为关闭了 inplace）
    normed = aiter.rmsnorm(x, weight, eps, residual_inplace=False)
    return normed + residual

# 用 rocprof 统计 MoE 层的 GEMM 调用模式
# 输出中每个 "KernelName: aiter_moe_gemm" 就是一次独立 GEMM
rocprof --stats ./run_deepseek_v4_flash.py \
    --model deepseek-v4-flash \
    --batch-size 8 \
    --seq-length 4096 \
    --topk 6

# 查看输出 CSV 中的 KernelExecution 时间分布
# 如果 moe_gemm 占总时间 > 40%，优先优化 MoE 路由