QuantFlashAttn HIF8 Per-Tensor 量化模式实现文档

1. 背景

1.1 HIF8 Per-Tensor 量化背景

QuantFlashAttn 算子此前仅支持 MXFP8（quant_compute_mode=1）量化模式，该模式采用 per-block group 量化（每 32 个元素一组，使用 FLOAT8_E8M0 类型的 block scale）。随着大模型推理场景对精度和性能的需求演进，HIFLOAT8（HiFloat8）数据类型被引入，其核心优势在于：

• Per-Tensor 量化：整个 Tensor 共享一个 FLOAT32 标量 descale，无需 per-block scale 数据，简化了数据准备和内存访问
• 更高精度：HIFLOAT8 采用 1 位符号位 + 4 位指数位 + 3 位尾数位格式，相较 FP8_E4M3FN 在部分场景下精度更优
• 更少内存开销：descale tensor 从多维 block scale 数组简化为 shape=(1,) 的标量，显著减少 descale 数据的内存占用

1.2 支持规格

1、 Layout：支持BNSD、BSND、TND
2、 Dtype与量化形式：
a. Q/K/V：hifloat8
b. QScale：fp32，per-tensor，shape：[1]
c. KScale：fp32，per-tensor，shape：[1]
d. VScale：fp32，per-tensor，shape：[1]
e. PScale：fp32，per-tensor静态量化，shape [1]，可选参数，不传默认值为1.0f
f. 输出：bf16
3、 Shape要求：
a. D=128
b. BNS支持泛化
4、 属性要求：
a. quant_mode：7：A8C8_QKV_HIF8_PER_TENSOR_P_FP8_E4M3_PER_TENSOR_SOFTMAX_FP32
b. mask_mode：0（仅支持不传入mask）/ 3 atten_mask [2048, 2048]
c. 不支持PA
d. 支持softmaxLSE

2. 主要方案

2.1 方案大纲

┌──────────────────────────────────────────────────────────────┐
│                      Host 侧 (Tiling)                         │
│                                                               │
│  1. GetQuantMode(): 校验 quant_compute_mode ∈ {1, 7}          │
│  2. ParseAxisInfo(): HIF8 时 early return（跳过 4D/5D 校验）   │
│  3. Checker 校验:                                             │
│     ├── Q/K/V dtype = DT_HIFLOAT8                             │
│     ├── descale dtype = DT_FLOAT, shape = (1,), 1D            │
│     ├── layout: kv=TND only, out∈{TND,BSND,BNSD}             │
│     └── headDim = 128                                         │
│  4. Tiling 配置:                                              │
│     ├── SplitPolicy: SINNER_256（非 SINNER_512）               │
│     ├── Config: 2/3 (S2Aligned256) + HIF8 fallback            │
│     └── QuantMode: QFA_HIF8_FP32_PREFILL = 7                  │
│  5. 生成 TilingKey                                            │
└──────────────────────────┬───────────────────────────────────┘
                           │
┌──────────────────────────▼───────────────────────────────────┐
│                    Kernel 侧 (Entry)                          │
│                                                               │
│  #elif (ORIG_DTYPE_Q == DT_HIFLOAT8)                          │
│    → quant_flash_attn_hif8<...>()                             │
│      ├── INPUT_T = hifloat8_t                                 │
│      ├── OUT_T = bfloat16_t                                   │
│      ├── useDn = true (仅 DN 路径)                            │
│      ├── CubeBlock = QuantFlashAttnBlockCubeHif8              │
│      ├── VecFaBlock = QuantFlashAttnBlockVecHif8              │
│      ├── VecFdBlock = QuantFlashAttnBlockVecFlashDecode       │
│      └── Kernel = QuantFlashAttnKernelHif8                    │
└──────────────────────────┬───────────────────────────────────┘
                           │
┌──────────────────────────▼───────────────────────────────────┐
│                 Kernel.Init() 执行流程                        │
│                                                               │
│  1. 从 GM 读取 per-tensor descale 标量:                       │
│     ├── qScale = *((__gm__ float *)dequantScaleQuery)         │
│     ├── kScale = *((__gm__ float *)dequantScaleKey)           │
│     └── vScale = *((__gm__ float *)dequantScaleValue)         │
│                                                               │
│  2. 传递给 Vec block:                                         │
│     ├── vecFaBlock.SetDescaleQK(qScale, kScale)               │
│     │   └── descaleQKValue_ = qScale * kScale                 │
│     └── vecFaBlock.SetVScale(vScale)                          │
│         └── vScaleValue_ = vScale                             │
│                                                               │
│  3. Cube block 不处理 descale:                                │
│     └── InitCubeBlock() 接受 descale 参数但不初始化 scale GM  │
└──────────────────────────┬───────────────────────────────────┘
                           │
┌──────────────────────────▼───────────────────────────────────┐
│                 Kernel.Process() 执行流程                     │
│                                                               │
│  ┌─ Cube 侧 (AIC) ────────────────────────────────────────┐  │
│  │                                                        │  │
│  │  IterateBmm1Dn:                                        │  │
│  │    ├── CopyQueryTile: L1 加载 Q (hifloat8_t)           │  │
│  │    ├── CopyKeyTile: L1 加载 K (hifloat8_t)             │  │
│  │    ├── s2BaseSizeCur = s2BaseSize (不切半)              │  │
│  │    ├── MatmulFullMX<Q_T, KV_T, T, 256, 128, dBaseSize, │  │
│  │    │   ..., float, float, INPUT_T, INPUT_T>(            │  │
│  │    │   ..., dummyScale, dummyScale)  ← 无 scale         │  │
│  │    └── FixpipeMm1Dn: L0C→UB                           │  │
│  │                                                        │  │
│  │  IterateBmm2Dn:                                        │  │
│  │    ├── CopyValueSlice: L1 加载 V (hifloat8_t)          │  │
│  │    ├── MatmulFullMX<..., dummyScale, dummyScale>       │  │
│  │    └── FixpipeMm2: L0C→UB                             │  │
│  └────────────────────────────────────────────────────────┘  │
│                                                               │
│  ┌─ Vec 侧 (AIV) ─────────────────────────────────────────┐  │
│  │                                                        │  │
│  │  ProcessVec1Dn:                                        │  │
│  │    ├── ComputeMaskLineDN: Arange+Compare 生成 mask     │  │
│  │    ├── descaleQK = 1.0 (局部变量，不使用 descaleQKValue_)│  │
│  │    ├── dScale = scale * deScaleQK (scale 已含 Q/K descale)│
│  │    ├── ProcessVec1VfDnHif8:                            │  │
│  │    │   ├── castTraitZero/One/Two/Three (hif8 4路 cast) │  │
│  │    │   ├── Or + Gather: 合并重排 cast 结果              │  │
│  │    │   └── maskLine: Arange+Compare GE (无需 mask tensor)│  │
│  │    └── SoftmaxDataCopyOut: sum/max 输出                │  │
│  │                                                        │  │
│  │  ProcessVec2OnUb:                                      │  │
│  │    ├── FlashUpdateNew/FlashUpdateLastNew: softmax 更新  │  │
│  │    ├── Muls(vec2ResUb, vec2ResUb, vScaleValue_): V descale│
│  │    ├── Cast(attenOut, vec2ResUb, CAST_ROUND): fp32→bf16│  │
│  │    └── Bmm2DataCopyOutTrans: 输出到 GM                  │  │
│  └────────────────────────────────────────────────────────┘  │
└──────────────────────────────────────────────────────────────┘

3. 主要工作

3.1 Host 侧 Tiling 修改

3.2 Kernel 侧修改

3.3 Metadata 修改

3.4 Torch Extension 修改

3.5 测试文件 (新增)