| import torch |
|
|
| from kernels.benchmark import Benchmark |
|
|
| SHAPES = { |
| "qkv": (8704, 6656), |
| "o": (6656, 4096), |
| "mlp": (39936, 6656), |
| "down": (6656, 19968), |
| "lm_head": (202048, 6656), |
| } |
|
|
|
|
| def _dequant(pw): |
| grid = torch.tensor( |
| [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], device=pw.qweight.device |
| ) |
| gs = pw.global_scale.float() |
| out = torch.empty(pw.n, pw.k, dtype=torch.float16, device=pw.qweight.device) |
| step = 16384 |
| for i in range(0, pw.n, step): |
| qw = pw.qweight[i : i + step] |
| nib = torch.stack([qw & 0xF, qw >> 4], dim=-1) |
| nib = nib.reshape(qw.shape[0], pw.k).long() |
| sign = torch.where((nib & 0x8) != 0, -1.0, 1.0) |
| vals = grid[nib & 0x7] * sign |
| sf = pw.sf_rowmajor[i : i + step].view(torch.float8_e4m3fn).float() |
| scale = sf.repeat_interleave(16, dim=1) |
| out[i : i + step] = ((vals * scale) / gs).to(torch.float16) |
| return out |
|
|
|
|
| class NvFP4GemmBenchmark(Benchmark): |
| seed: int = 42 |
|
|
| def _setup_shape(self, name, m): |
| n, k = SHAPES[name] |
| torch.manual_seed(self.seed) |
| w = torch.randn(n, k, dtype=torch.bfloat16) * 0.02 |
| self.pw = self.kernel.pack(w.to(self.device)) |
| self.x = torch.randn(m, k, dtype=torch.bfloat16, device=self.device) * 0.125 |
| self.ref_w = _dequant(self.pw) |
| self.xs = self.x.to(torch.float16) |
|
|
| def setup_decode_qkv(self): |
| self._setup_shape("qkv", 1) |
|
|
| def benchmark_decode_qkv(self): |
| self.out = self.kernel.gemm(self.pw, self.x) |
|
|
| def verify_decode_qkv(self): |
| return (self.xs @ self.ref_w.T).to(torch.bfloat16) |
|
|
| def setup_decode_o(self): |
| self._setup_shape("o", 1) |
|
|
| def benchmark_decode_o(self): |
| self.out = self.kernel.gemm(self.pw, self.x) |
|
|
| def verify_decode_o(self): |
| return (self.xs @ self.ref_w.T).to(torch.bfloat16) |
|
|
| def setup_decode_mlp(self): |
| self._setup_shape("mlp", 1) |
|
|
| def benchmark_decode_mlp(self): |
| self.out = self.kernel.gemm(self.pw, self.x) |
|
|
| def verify_decode_mlp(self): |
| return (self.xs @ self.ref_w.T).to(torch.bfloat16) |
|
|
| def setup_decode_down(self): |
| self._setup_shape("down", 1) |
|
|
| def benchmark_decode_down(self): |
| self.out = self.kernel.gemm(self.pw, self.x) |
|
|
| def verify_decode_down(self): |
| return (self.xs @ self.ref_w.T).to(torch.bfloat16) |
|
|
| def setup_decode_lm_head(self): |
| self._setup_shape("lm_head", 1) |
|
|
| def benchmark_decode_lm_head(self): |
| self.out = self.kernel.gemm(self.pw, self.x) |
|
|
| def verify_decode_lm_head(self): |
| return (self.xs @ self.ref_w.T).to(torch.bfloat16) |
|
|
| def setup_prefill_mlp(self): |
| self._setup_shape("mlp", 512) |
| self.x = self.x * 0.2 |
| xq = self.kernel.quantize_reference(self.x.cpu()).to(self.device) |
| self.xs = xq.to(torch.float16) |
|
|
| def benchmark_prefill_mlp(self): |
| self.out = self.kernel.gemm(self.pw, self.x) |
|
|
| def verify_prefill_mlp(self): |
| return (self.xs @ self.ref_w.T).to(torch.bfloat16) |
|
|
|
|
| class NvFP4RawGemvBenchmark(Benchmark): |
| seed: int = 42 |
|
|
| def _setup_shape(self, name): |
| n, k = SHAPES[name] |
| torch.manual_seed(self.seed) |
| w = torch.randn(n, k, dtype=torch.bfloat16) * 0.02 |
| pw = self.kernel.pack(w.to(self.device)) |
| self.qw = pw.qweight |
| self.sf = pw.sf_rowmajor |
| self.alpha = (1.0 / pw.global_scale).float().reshape(1) |
| self.x = torch.randn(1, k, dtype=torch.bfloat16, device=self.device) * 0.125 |
| self.ref_w = _dequant(pw) |
| self.xs = self.x.to(torch.float16) |
|
|
| def setup_gemv_qkv(self): |
| self._setup_shape("qkv") |
|
|
| def benchmark_gemv_qkv(self): |
| self.out = self.kernel._ops.ops.nvfp4_gemv(self.x, self.qw, self.sf, self.alpha) |
|
|
| def verify_gemv_qkv(self): |
| return (self.xs @ self.ref_w.T).to(torch.bfloat16) |
|
|
| def setup_gemv_o(self): |
| self._setup_shape("o") |
|
|
| def benchmark_gemv_o(self): |
| self.out = self.kernel._ops.ops.nvfp4_gemv(self.x, self.qw, self.sf, self.alpha) |
|
|
| def verify_gemv_o(self): |
| return (self.xs @ self.ref_w.T).to(torch.bfloat16) |
|
|
| def setup_gemv_mlp(self): |
| self._setup_shape("mlp") |
|
|
| def benchmark_gemv_mlp(self): |
| self.out = self.kernel._ops.ops.nvfp4_gemv(self.x, self.qw, self.sf, self.alpha) |
|
|
| def verify_gemv_mlp(self): |
| return (self.xs @ self.ref_w.T).to(torch.bfloat16) |
|
|
| def setup_gemv_down(self): |
| self._setup_shape("down") |
|
|
| def benchmark_gemv_down(self): |
| self.out = self.kernel._ops.ops.nvfp4_gemv(self.x, self.qw, self.sf, self.alpha) |
|
|
| def verify_gemv_down(self): |
| return (self.xs @ self.ref_w.T).to(torch.bfloat16) |
|
|
| def setup_gemv_lm_head(self): |
| self._setup_shape("lm_head") |
|
|
| def benchmark_gemv_lm_head(self): |
| self.out = self.kernel._ops.ops.nvfp4_gemv(self.x, self.qw, self.sf, self.alpha) |
|
|
| def verify_gemv_lm_head(self): |
| return (self.xs @ self.ref_w.T).to(torch.bfloat16) |
|
|
