File size: 4,024 Bytes
63ec7fb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | 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()
|