| import torch |
|
|
| from kernels.benchmark import Benchmark |
|
|
|
|
| _original_allclose = torch.allclose |
|
|
|
|
| def _bf16_max_ulp(input: torch.Tensor, other: torch.Tensor) -> int: |
| got_bits = input.detach().cpu().view(torch.int16).to(torch.int32) & 0xFFFF |
| exp_bits = other.detach().cpu().view(torch.int16).to(torch.int32) & 0xFFFF |
| got_ordered = torch.where((got_bits & 0x8000) != 0, 0x8000 - (got_bits & 0x7FFF), got_bits) |
| exp_ordered = torch.where((exp_bits & 0x8000) != 0, 0x8000 - (exp_bits & 0x7FFF), exp_bits) |
| return int((got_ordered - exp_ordered).abs().max().item()) |
|
|
|
|
| def _flashrt_allclose(input, other, rtol=1e-05, atol=1e-08, equal_nan=False): |
| if input.dtype == torch.bfloat16 and other.dtype == torch.bfloat16: |
| return _bf16_max_ulp(input, other) <= 5 |
| return _original_allclose(input, other, rtol=rtol, atol=atol, equal_nan=equal_nan) |
|
|
|
|
| torch.allclose = _flashrt_allclose |
|
|
|
|
| DECODE_SHAPES = [ |
| ("k4096_n1024", 4096, 1024), |
| ("k4096_n4096", 4096, 4096), |
| ("k4096_n12288", 4096, 12288), |
| ("k12288_n1024", 12288, 1024), |
| ("k12288_n4096", 12288, 4096), |
| ("k12288_n12288", 12288, 12288), |
| ] |
|
|
|
|
| def _swizzled_bytes(rows: int, cols: int) -> int: |
| n_blocks = cols // 16 |
| return ((rows + 127) // 128) * ((n_blocks + 3) // 4) * 512 |
|
|
|
|
| def _swizzle_constant_scale(rows: int, cols: int, value: int) -> torch.Tensor: |
| return torch.full((_swizzled_bytes(rows, cols),), value, dtype=torch.uint8) |
|
|
|
|
| def _reference_swizzle(scales: torch.Tensor) -> torch.Tensor: |
| rows, n_blocks = scales.shape |
| n_col_super = (n_blocks + 3) // 4 |
| src = scales.cpu() |
| out = torch.zeros( |
| ((rows + 127) // 128) * n_col_super * 512, |
| dtype=torch.uint8, |
| ) |
| for row in range(rows): |
| rb = row // 128 |
| ri = row % 128 |
| for block in range(n_blocks): |
| cb = block // 4 |
| ci = block % 4 |
| super_idx = rb * n_col_super + cb |
| inner_off = (ri % 32) * 16 + (ri // 32) * 4 + ci |
| out[super_idx * 512 + inner_off] = src[row, block] |
| return out |
|
|
|
|
| def _ue4m3_to_float(byte: int) -> float: |
| sign = -1.0 if (byte & 0x80) else 1.0 |
| exp = (byte >> 3) & 0x0F |
| mant = byte & 0x07 |
| if exp == 0: |
| return sign * (mant / 8.0) * (2.0 ** -6) |
| if exp == 15 and mant == 7: |
| return 0.0 |
| return sign * (1.0 + mant / 8.0) * (2.0 ** (exp - 7)) |
|
|
|
|
| def _ue4m3_lut() -> torch.Tensor: |
| return torch.tensor([_ue4m3_to_float(i) for i in range(256)], dtype=torch.float32) |
|
|
|
|
| def _fp4_codebook() -> torch.Tensor: |
| return torch.tensor( |
| [ |
| 0.0, |
| 0.5, |
| 1.0, |
| 1.5, |
| 2.0, |
| 3.0, |
| 4.0, |
| 6.0, |
| -0.0, |
| -0.5, |
| -1.0, |
| -1.5, |
| -2.0, |
| -3.0, |
| -4.0, |
| -6.0, |
| ], |
| dtype=torch.float32, |
| ) |
|
|
|
|
| def _unpack_fp4(packed: torch.Tensor) -> torch.Tensor: |
| codebook = _fp4_codebook().to(packed.device) |
| lo = packed & 0x0F |
| hi = packed >> 4 |
| out = torch.empty( |
| (packed.shape[0], packed.shape[1] * 2), |
| device=packed.device, |
| dtype=torch.float32, |
| ) |
| out[:, 0::2] = codebook[lo.long()] |
| out[:, 1::2] = codebook[hi.long()] |
| return out |
|
|
|
|
| def _reference_smallm( |
| a_packed: torch.Tensor, |
| b_packed: torch.Tensor, |
| sfa_linear: torch.Tensor, |
| sfb_linear: torch.Tensor, |
| K: int, |
| alpha: float, |
| chunk_rows: int = 256, |
| ) -> torch.Tensor: |
| device = b_packed.device |
| N = b_packed.shape[0] |
| lut = _ue4m3_lut().to(device) |
| a = _unpack_fp4(a_packed.reshape(1, -1)).reshape(K) |
| a_scale = lut[sfa_linear.reshape(-1).to(device).long()].repeat_interleave(16) |
| a = a * a_scale |
| sfb_linear = sfb_linear.to(device) |
| out = torch.empty((N,), device=device, dtype=torch.bfloat16) |
| for start in range(0, N, chunk_rows): |
| end = min(start + chunk_rows, N) |
| b = _unpack_fp4(b_packed[start:end]) |
| b_scale = lut[sfb_linear[start:end].long()].repeat_interleave(16, dim=1) |
| expected = (b * b_scale * a.reshape(1, K)).sum(dim=1) * alpha |
| out[start:end] = expected.to(torch.bfloat16) |
| return out |
|
|
|
|
| class Nvfp4W4A4DecodeMatvecBenchmark(Benchmark): |
| seed = 23 |
|
|
| def _setup_shape(self, K: int, N: int) -> None: |
| torch.manual_seed(600 + K + N) |
| if torch.cuda.is_available(): |
| torch.cuda.manual_seed_all(600 + K + N) |
| self.K = K |
| self.N = N |
| self.alpha = 0.5 |
| self.a_packed = torch.randint( |
| 0, 256, (K // 2,), device=self.device, dtype=torch.uint8 |
| ) |
| self.b_packed = torch.randint( |
| 0, 256, (N, K // 2), device=self.device, dtype=torch.uint8 |
| ) |
| self.sfa_linear = torch.randint(0, 0x78, (1, K // 16), dtype=torch.uint8) |
| self.sfb_linear = torch.randint(0, 0x78, (N, K // 16), dtype=torch.uint8) |
| self.sfa = _reference_swizzle(self.sfa_linear).to(self.device) |
| self.sfb = _reference_swizzle(self.sfb_linear).to(self.device) |
| self.out = torch.empty((N,), device=self.device, dtype=torch.bfloat16) |
|
|
| def _benchmark(self) -> None: |
| self.kernel.nvfp4_w4a4_decode_matvec_bf16out( |
| self.a_packed, |
| self.b_packed, |
| self.sfa, |
| self.sfb, |
| alpha=self.alpha, |
| out=self.out, |
| ) |
|
|
| def _reference(self) -> torch.Tensor: |
| return _reference_smallm( |
| self.a_packed, |
| self.b_packed, |
| self.sfa_linear, |
| self.sfb_linear, |
| self.K, |
| self.alpha, |
| ) |
|
|
|
|
| def _register_shapes() -> None: |
| for label, K, N in DECODE_SHAPES: |
|
|
| def setup(self, K=K, N=N) -> None: |
| self._setup_shape(K, N) |
|
|
| def benchmark(self) -> None: |
| self._benchmark() |
|
|
| def verify(self) -> torch.Tensor: |
| return self._reference() |
|
|
| setattr(Nvfp4W4A4DecodeMatvecBenchmark, f"setup_{label}", setup) |
| setattr(Nvfp4W4A4DecodeMatvecBenchmark, f"benchmark_{label}", benchmark) |
| setattr(Nvfp4W4A4DecodeMatvecBenchmark, f"verify_{label}", verify) |
|
|
|
|
| _register_shapes() |
|
|
