| import torch |
|
|
| from kernels.benchmark import Benchmark |
|
|
|
|
| LAYOUT_SHAPES = [ |
| ("rows1_d4096", 1, 4096), |
| ("rows2_d4096", 2, 4096), |
| ("rows31_d4096", 31, 4096), |
| ("rows32_d4096", 32, 4096), |
| ("rows33_d4096", 33, 4096), |
| ("rows127_d4096", 127, 4096), |
| ("rows128_d4096", 128, 4096), |
| ("rows129_d4096", 129, 4096), |
| ("rows16_d1024", 16, 1024), |
| ("rows16_d2048", 16, 2048), |
| ("rows16_d8192", 16, 8192), |
| ("rows16_d12288", 16, 12288), |
| ("rows64_d16384", 64, 16384), |
| ] |
|
|
|
|
| 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 blk in range(n_blocks): |
| cb = blk // 4 |
| ci = blk % 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, blk] |
| return out.to(scales.device) |
|
|
|
|
| class Nvfp4ScaleFactorReshapeBenchmark(Benchmark): |
| seed = 7 |
|
|
| def _setup_shape(self, rows: int, D: int) -> None: |
| self.scales = torch.randint( |
| 0, |
| 256, |
| (rows, D // 16), |
| device=self.device, |
| dtype=torch.uint8, |
| ) |
| n_col_super = ((D // 16) + 3) // 4 |
| self.out = torch.zeros( |
| ((rows + 127) // 128) * n_col_super * 512, |
| device=self.device, |
| dtype=torch.uint8, |
| ) |
|
|
| def _reference(self): |
| return _reference_swizzle(self.scales) |
|
|
| def setup_rows1_d4096(self): |
| self._setup_shape(1, 4096) |
|
|
| def benchmark_rows1_d4096(self): |
| self.kernel.nvfp4_sf_linear_to_swizzled(self.scales, out=self.out) |
|
|
| def verify_rows1_d4096(self): |
| return self._reference() |
|
|
| def setup_rows16_d12288(self): |
| self._setup_shape(16, 12288) |
|
|
| def benchmark_rows16_d12288(self): |
| self.kernel.nvfp4_sf_linear_to_swizzled(self.scales, out=self.out) |
|
|
| def verify_rows16_d12288(self): |
| return self._reference() |
|
|
| def setup_rows64_d16384(self): |
| self._setup_shape(64, 16384) |
|
|
| def benchmark_rows64_d16384(self): |
| self.kernel.nvfp4_sf_linear_to_swizzled(self.scales, out=self.out) |
|
|
| def verify_rows64_d16384(self): |
| return self._reference() |
|
|
| def setup_rows128_d4096(self): |
| self._setup_shape(128, 4096) |
|
|
| def benchmark_rows128_d4096(self): |
| self.kernel.nvfp4_sf_linear_to_swizzled(self.scales, out=self.out) |
|
|
| def verify_rows128_d4096(self): |
| return self._reference() |
|
|
| def setup_rows129_d4096(self): |
| self._setup_shape(129, 4096) |
|
|
| def benchmark_rows129_d4096(self): |
| self.kernel.nvfp4_sf_linear_to_swizzled(self.scales, out=self.out) |
|
|
| def verify_rows129_d4096(self): |
| return self._reference() |
|
|
|
|
| def _register_layout_shapes() -> None: |
| for label, rows, D in LAYOUT_SHAPES: |
|
|
| def setup(self, rows=rows, D=D) -> None: |
| self._setup_shape(rows, D) |
|
|
| def benchmark(self) -> None: |
| self.kernel.nvfp4_sf_linear_to_swizzled(self.scales, out=self.out) |
|
|
| def verify(self): |
| return self._reference() |
|
|
| setattr(Nvfp4ScaleFactorReshapeBenchmark, f"setup_{label}", setup) |
| setattr(Nvfp4ScaleFactorReshapeBenchmark, f"benchmark_{label}", benchmark) |
| setattr(Nvfp4ScaleFactorReshapeBenchmark, f"verify_{label}", verify) |
|
|
|
|
| _register_layout_shapes() |
|
|
