| import torch |
|
|
| from kernels.benchmark import Benchmark |
|
|
|
|
| SILU_QUANT_SHAPES = [ |
| ("decode_r1_h4096", 1, 4096), |
| ("decode_r2_h4096", 2, 4096), |
| ("decode_r4_h4096", 4, 4096), |
| ("decode_r8_h4096", 8, 4096), |
| ("decode_r1_h8192", 1, 8192), |
| ("decode_r4_h8192", 4, 8192), |
| ("decode_r8_h8192", 8, 8192), |
| ("decode_r1_h12288", 1, 12288), |
| ("decode_r4_h12288", 4, 12288), |
| ("decode_r8_h12288", 8, 12288), |
| ("decode_r1_h16384", 1, 16384), |
| ("decode_r4_h16384", 4, 16384), |
| ("decode_r8_h16384", 8, 16384), |
| ("small_r16_h4096", 16, 4096), |
| ("small_r32_h4096", 32, 4096), |
| ("small_r16_h8192", 16, 8192), |
| ("small_r32_h8192", 32, 8192), |
| ("small_r16_h12288", 16, 12288), |
| ("small_r32_h12288", 32, 12288), |
| ("small_r16_h16384", 16, 16384), |
| ("small_r32_h16384", 32, 16384), |
| ("prefill_r64_h4096", 64, 4096), |
| ("prefill_r128_h4096", 128, 4096), |
| ("prefill_r256_h4096", 256, 4096), |
| ("video_r1024_h4096", 1024, 4096), |
| ("video_r2520_h4096", 2520, 4096), |
| ("prefill_r64_h8192", 64, 8192), |
| ("prefill_r128_h8192", 128, 8192), |
| ("prefill_r256_h8192", 256, 8192), |
| ("video_r1024_h8192", 1024, 8192), |
| ("video_r2520_h8192", 2520, 8192), |
| ("prefill_r64_h12288", 64, 12288), |
| ("prefill_r128_h12288", 128, 12288), |
| ("prefill_r256_h12288", 256, 12288), |
| ("video_r1024_h12288", 1024, 12288), |
| ("video_r2520_h12288", 2520, 12288), |
| ] |
|
|
|
|
| def _scale_bytes(rows: int, cols: int) -> int: |
| n_blocks = cols // 16 |
| return ((rows + 127) // 128) * ((n_blocks + 3) // 4) * 512 |
|
|
|
|
| class SiluMulNvfp4SplitBenchmark(Benchmark): |
| seed = 17 |
|
|
| def _setup_shape(self, rows: int, cols: int) -> None: |
| self.gate = torch.randn( |
| (rows, cols), device=self.device, dtype=torch.bfloat16 |
| ).contiguous() |
| self.up = torch.randn( |
| (rows, cols), device=self.device, dtype=torch.bfloat16 |
| ).contiguous() |
| self.packed = torch.empty( |
| (rows, cols // 2), device=self.device, dtype=torch.uint8 |
| ) |
| self.scales = torch.zeros( |
| (_scale_bytes(rows, cols),), device=self.device, dtype=torch.uint8 |
| ) |
| self.out = self.packed |
|
|
| def _benchmark(self) -> None: |
| self.kernel.silu_mul_quant_nvfp4_swizzled_bf16( |
| self.gate, |
| self.up, |
| packed=self.packed, |
| scales=self.scales, |
| ) |
|
|
|
|
| class SiluMulNvfp4MergedBenchmark(Benchmark): |
| seed = 19 |
|
|
| def _setup_shape(self, rows: int, cols: int) -> None: |
| gate = torch.randn( |
| (rows, cols), device=self.device, dtype=torch.bfloat16 |
| ).contiguous() |
| up = torch.randn((rows, cols), device=self.device, dtype=torch.bfloat16) |
| self.merged = torch.cat([gate, up], dim=1).contiguous() |
| self.packed = torch.empty( |
| (rows, cols // 2), device=self.device, dtype=torch.uint8 |
| ) |
| self.scales = torch.zeros( |
| (_scale_bytes(rows, cols),), device=self.device, dtype=torch.uint8 |
| ) |
| self.out = self.packed |
|
|
| def _benchmark(self) -> None: |
| self.kernel.silu_mul_merged_quant_nvfp4_swizzled_bf16( |
| self.merged, |
| packed=self.packed, |
| scales=self.scales, |
| ) |
|
|
|
|
| def _register_shapes() -> None: |
| for label, rows, cols in SILU_QUANT_SHAPES: |
|
|
| def setup_split(self, rows=rows, cols=cols) -> None: |
| self._setup_shape(rows, cols) |
|
|
| def benchmark_split(self) -> None: |
| self._benchmark() |
|
|
| setattr(SiluMulNvfp4SplitBenchmark, f"setup_{label}", setup_split) |
| setattr(SiluMulNvfp4SplitBenchmark, f"benchmark_{label}", benchmark_split) |
|
|
| def setup_merged(self, rows=rows, cols=cols) -> None: |
| self._setup_shape(rows, cols) |
|
|
| def benchmark_merged(self) -> None: |
| self._benchmark() |
|
|
| setattr(SiluMulNvfp4MergedBenchmark, f"setup_{label}", setup_merged) |
| setattr(SiluMulNvfp4MergedBenchmark, f"benchmark_{label}", benchmark_merged) |
|
|
|
|
| _register_shapes() |
|
|
