admin/matmul-analysis
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update models.py: fp4 support

Browse Source
This commit is contained in:
admin
2026-09-03 12:42:35 +00:00
parent c2028a41ed
commit db8399b62c
1 changed files with 7 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
BMM/BMM_Theory/bmm_theory/models.py
View File

@@ -28,13 +28,15 @@ DTYPE_BYTES = {
"fp8_e4m3": 1, "fp8_e4m3": 1,
"fp8_e5m2": 1, "fp8_e5m2": 1,
"int8": 1, "int8": 1,
"fp4": 0.5, # 半字节 (对齐 bmmv3, 2026-09-03)
"fp4_e2m1": 0.5,
} }
# Cube 累加器 (L0C) 中元素字节数: 16bit 输入 -> fp32 累加; fp8 输入 -> fp32 累加 # Cube 累加器 (L0C) 中元素字节数: 16bit 输入 -> fp32 累加; fp8/fp4 输入 -> fp32 累加
L0C_DTYPE_BYTES = 4 L0C_DTYPE_BYTES = 4
def dtype_bytes(dtype: str) -> int: def dtype_bytes(dtype: str) -> float:
key = dtype.strip().lower() key = dtype.strip().lower()
if key not in DTYPE_BYTES: if key not in DTYPE_BYTES:
raise ValueError(f"不支持的 dtype: {dtype!r}, 支持 {sorted(DTYPE_BYTES)}") raise ValueError(f"不支持的 dtype: {dtype!r}, 支持 {sorted(DTYPE_BYTES)}")
@@ -109,12 +111,13 @@ class BmmCase:
return max(self.batch_a, self.batch_b) return max(self.batch_a, self.batch_b)
@property @property
def dtype_in_bytes(self) -> int: def dtype_in_bytes(self) -> float:
# A/B 输入元素字节数 (要求 A/B 同 dtype, 不一致时取较大者并在校验中报 warning) # A/B 输入元素字节数 (要求 A/B 同 dtype, 不一致时取较大者并在校验中报 warning)
# 返回 float 以兼容 fp4 (0.5B)
return max(dtype_bytes(self.dtype_a), dtype_bytes(self.dtype_b)) return max(dtype_bytes(self.dtype_a), dtype_bytes(self.dtype_b))
@property @property
def dtype_out_bytes(self) -> int: def dtype_out_bytes(self) -> float:
return dtype_bytes(self.dtype_c) return dtype_bytes(self.dtype_c)
@property @property