# 粗粒度缩放（低效）
A_fp8 = quantize(A, scale=global_scale)  # 全局缩放

# 细粒度缩放（DeepGEMM采用）
A_fp8 = quantize(A, scale=per_block_scale)  # 每块独立缩放

import torch
import deep_gemm

# 输入矩阵 (M=1024, N=4096, K=4096)
M, N, K = 1024, 4096, 4096

# FP8量化
A_fp8, A_scale = fp8_quantize(A)  # [M, K]
B_fp8, B_scale = fp8_quantize(B)  # [N, K] (转置后)

# 调用GEMM
D = deep_gemm.fp8_gemm_nt(
    A_fp8,          # 非转置
    B_fp8,          # 转置后传入
    D_dtype=torch.bfloat16,
    lhs_scale=A_scale,  # FP32格式（SM90）
    rhs_scale=B_scale,
    num_sms=120,       # 使用120个SM
)

# 连续布局分组GEMM
# 每个专家处理不同数量的token
m_grouped_fp8_gemm_nt_contiguous(
    inputs,      # [total_tokens, K] 所有专家的token拼接
    weights,     # [num_experts, N, K]
    scales,     # [num_experts]
    expert_indices,  # 每个token属于哪个专家
)

# 获取对称内存缓冲区（需要PyTorch >= 2.9）
buffer = deep_gemm.get_symm_buffer_for_mega_moe(
    group, num_experts, num_max_tokens_per_rank,
    num_topk, hidden, intermediate_hidden
)

# 权重变换
transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe(
    l1_weights, l2_weights
)

# 填充缓冲区
buffer.x[:num_tokens].copy_(x_fp8)
buffer.topk_idx[:num_tokens].copy_(topk_idx)
buffer.topk_weights[:num_tokens].copy_(topk_weights)

# 调用Mega MoE内核
y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda')
deep_gemm.fp8_fp4_mega_moe(y, transformed_l1, transformed_l2, buffer)

# MQA (Multi-Query Attention) 评分
# 用于token到token的logit计算
output = deep_gemm.fp8_mqa_logits(
    q,           # [seq_len, num_heads, head_dim]
    kv,          # [seq_len_kv, head_dim]
    weights,     # [seq_len, num_heads]
    cu_seq_len_k_start,  # 每个query对应的kv起始位置
    cu_seq_len_k_end,    # 每个query对应的kv结束位置
)

# FP8 × FP4 GEMM（超低精度推理）
fp8_fp4_gemm_nt(
    A_fp8,       # FP8输入
    B_fp4,       # FP4权重（更紧凑）
    scales,      # UE8M0格式缩放因子
)

# 启用NVRTC（编译快10倍，可能有少量性能损失）
export DG_JIT_USE_NVRTC=1

# 禁用NVRTC（编译慢，但性能最优）
export DG_JIT_USE_NVRTC=0

# 1. 克隆仓库（包含子模块）
git clone --recursive git@github.com:deepseek-ai/DeepGEMM.git
cd DeepGEMM

# 2. 开发模式构建
./develop.sh

# 3. 安装
./install.sh

# 4. 验证安装
python -c "import deep_gemm; print(deep_gemm.__version__)"

import torch
import deep_gemm

# 创建随机FP8输入
M, N, K = 1024, 4096, 4096
A = torch.randn(M, K, device='cuda', dtype=torch.float8_e4m3fn)
B = torch.randn(N, K, device='cuda', dtype=torch.float8_e4m3fn)

# 量化
A_fp8, A_sf = deep_gemm.FP8_e4m3(A)
B_fp8, B_sf = deep_gemm.FP8_e4m3(B)

# GEMM计算
D = deep_gemm.fp8_gemm_nt(
    A_fp8, B_fp8,
    lhs_scale=A_sf,
    rhs_scale=B_sf,
    D_dtype=torch.bfloat16
)

print(f"Output shape: {D.shape}")  # [1024, 4096]

# 设置使用的SM数量（留一些给系统）
deep_gemm.set_num_sms(120)  # H100有144个SM

# 设置Tensor Core利用率（用于资源预留）
deep_gemm.set_tc_util(0.95)  # 保留5%给其他操作

# 设置PDL（Programmatic Dependent Launch）
deep_gemm.set_pdl(1)  # 启用依赖内核调度优化

# 获取最优M/K对齐
alignment = deep_gemm.get_theoretical_mk_alignment_for_contiguous_layout()

# 启用line info（用于nsys/nsysd分析）
export DG_JIT_WITH_LINEINFO=1

# 导出PTX（查看生成代码）
export DG_JIT_DUMP_PTX=1

# 导出SASS（查看最终汇编）
export DG_JIT_DUMP_SASS=1

# 显示编译时间
export DG_JIT_PRINT_LOAD_TIME=1

# 混合精度训练中的FP8 GEMM
class FP8Linear(torch.nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.weight = torch.nn.Parameter(torch.randn(
            out_features, in_features, device='cuda', dtype=torch.float8_e4m3fn
        ))
        self.scale = torch.nn.Parameter(torch.ones(out_features, device='cuda'))
    
    def forward(self, x):
        return deep_gemm.fp8_gemm_nt(
            x, self.weight.t(),
            rhs_scale=self.scale,
            D_dtype=torch.bfloat16
        )

# Prefill阶段的高效FP8推理
def prefill_with_fp8(model, input_ids):
    # FP8量化
    x_fp8, x_scale = quantize_fp8(hidden_states)
    
    # FP8 GEMM替代FP16/BF16
    for layer in model.layers:
        # Self-attention
        q = deep_gemm.fp8_gemm_nt(x_fp8, layer.q_weight, rhs_scale=layer.q_scale)
        k = deep_gemm.fp8_gemm_nt(x_fp8, layer.k_weight, rhs_scale=layer.k_scale)
        v = deep_gemm.fp8_gemm_nt(x_fp8, layer.v_weight, rhs_scale=layer.v_scale)
        
        # FFN
        ffn_out = deep_gemm.fp8_gemm_nt(
            deep_gemm.fp8_gemm_nt(x_fp8, layer.gate_weight, rhs_scale=layer.gate_scale),
            layer.up_weight.t(),
            rhs_scale=layer.up_scale
        ) * torch.nn.functional.silu(q)  # SwiGLU

# DeepSeek V3风格的MoE推理
def moe_forward_with_deepgemm(router_output, expert_weights, expert_biases):
    # Top-K专家选择
    topk_weights, topk_indices = torch.topk(router_output, k=8, dim=-1)
    
    # 连续布局分组GEMM
    output = deep_gemm.m_grouped_fp8_gemm_nt_contiguous(
        hidden_states,           # 所有token拼接
        expert_weights,          # [num_experts, N, K]
        expert_scales,           # [num_experts]
        topk_indices,            # token→expert映射
    )
    
    # 加权合并
    return output * topk_weights.unsqueeze(-1)

# 1. 克隆并初始化子模块
git clone --recursive git@github.com:deepseek-ai/DeepGEMM.git
cd DeepGEMM && git submodule update --init --recursive

# 2. 修改内核代码
# 编辑 src/kernels/*.cu

# 3. 重新编译
./develop.sh

# 4. 运行测试
python -m pytest tests/test_core.py -v

# 5. 性能基准测试
python -m pytest tests/bench_gemm.py -v

// src/kernels/my_new_kernel.cu

#include "kernel_utils.cuh"

// 定义内核配置
struct MyKernelConfig {
    int block_m, block_n, block_k;
    int stages;
    // ...
};

// JIT编译入口
torch::Tensor my_new_kernel(
    torch::Tensor input,
    torch::Tensor weight,
    // 其他参数
) {
    // 1. 选择最优配置
    auto config = select_config(input, weight);
    
    // 2. 分配输出张量
    auto output = torch::empty_like(input);
    
    // 3. 准备kernel参数
    LaunchParams<MyKernelConfig> launch_params(config);
    
    // 4. CUDA launch
    cudaKernel<<<launch_params.grid, launch_params.block>>>(
        output.data_ptr(),
        input.data_ptr(),
        weight.data_ptr(),
        // ...
    );
    
    return output;
}

#include <deepgemm/gemm.h>
#include <deepgemm/kernel_launch.cuh>

// 1. 创建配置
GemmConfig config;
config.m = 8192;           // M维度
config.n = 8192;           // N维度  
config.k = 8192;           // K维度
config.dtype_a = FP8;      // A矩阵精度
config.dtype_b = FP8;      // B矩阵精度
config.dtype_c = BF16;     // 输出精度

// 2. 创建算子
auto gemm = DeepGemmFP8::create(config);

// 3. 执行计算
gemm->forward(output, input_a, input_b);

// 4. 同步
cudaStreamSynchronize(stream);

import torch
from deepgemm import DeepGEMM

# 1. 初始化
gemm = DeepGEMM(device='cuda:0', dtype='fp8')

# 2. 准备数据
M, N, K = 8192, 8192, 8192
A = torch.randn(M, K, dtype=torch.float8_e4m3fn, device='cuda')
B = torch.randn(K, N, dtype=torch.float8_e4m3fn, device='cuda')

# 3. 执行GEMM
C = gemm(A, B)  # 返回BF16

# 4. 验证结果
assert C.dtype == torch.bfloat16