"""`tile_kernels.quant` -- FP8/FP4 casting, fused SwiGLU, scale-factor layout. The compiled functions take `num_per_channels` and write into caller-provided `out` buffers; the references take `block_size` and return fresh tensors. The signatures here are the compiled ones, because that is what the 4 call sites in the model are written against. The adapters are thin but not vacuous -- three things genuinely differ: * `num_per_channels=n` means a per-token block of `(1, n)`, not `(n, n)`. * `swiglu_forward_and_per_token_cast` fuses activation and cast; the reference has them separate, so it is activation followed by cast. * `cast_back`'s `fmt` is the *output* dtype (`'bf16'`/`'fp32'`), which is easy to misread as the quantised format. Verified against the compiled kernels: `per_token_cast`, `cast_back` and `swiglu_forward_and_per_token_cast` are bitwise identical on the SM90 recipe (`scripts/tile_kernels_patch/compare_with_cuda.py`). """ from __future__ import annotations from collections.abc import Sequence from typing import Any import torch from deepinfer.tk_compat import use_references if not use_references(): from tile_kernels import quant as _compiled from tile_kernels import torch as _ref_raw from deepinfer.tk_compat._sync import reference_module # Calls through `_ref` may sync; see `_sync.py` for why that is the reference's design. _ref = reference_module(_ref_raw) def per_token_cast( x: torch.Tensor | tuple[torch.Tensor, torch.Tensor], fmt: str, num_per_channels: int, x_block_size: tuple[int, int] | None = None, use_tma_aligned_col_major_sf: bool = False, round_sf: bool = False, use_packed_ue8m0: bool = False, use_e4m3_sf: bool = False, stochastic_cast: bool = False, out: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Cast to FP8/FP4 with per-token scaling factors.""" if not use_references(): return _compiled.per_token_cast( x, fmt, num_per_channels, x_block_size, use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0, use_e4m3_sf, stochastic_cast, out, ) if stochastic_cast: raise NotImplementedError( "stochastic_cast has no torch reference; it changes rounding, so silently using the " "deterministic path would produce plausible but different numbers" ) result = _ref.cast( x, fmt, block_size=(1, num_per_channels), x_block_size=x_block_size, round_sf=round_sf, use_tma_aligned_col_major_sf=use_tma_aligned_col_major_sf, use_packed_ue8m0=use_packed_ue8m0, use_e4m3_sf=use_e4m3_sf, ) if out is not None: # Honour the out-buffer contract: the call sites allocate these and may hold views. out[0].copy_(result[0]) out[1].copy_(result[1]) return out[0], out[1] return result def cast_back( x: tuple[torch.Tensor, torch.Tensor], fmt: str, x_block_size: tuple[int, int], x_special_fmt: str | None = None, ) -> torch.Tensor: """Dequantise an FP8/FP4 tensor. `fmt` is the output dtype, not the quantised format.""" if not use_references(): return _compiled.cast_back(x, fmt, x_block_size, x_special_fmt) if x_special_fmt is not None: raise NotImplementedError(f"x_special_fmt={x_special_fmt!r} has no torch reference") return _ref.cast_back(x, fmt, block_size=x_block_size) def swiglu_forward( x: torch.Tensor, fmt: str, psum_num_tokens_per_expert: torch.Tensor | None = None, topk_weights: torch.Tensor | None = None, routed_scaling_factor: float | None = None, clamp_value: float | None = None, clamped_count: torch.Tensor | None = None, ) -> torch.Tensor: """SwiGLU activation. `fmt` is unused by the compiled path here; kept for signature parity.""" if not use_references(): return _compiled.swiglu_forward( x, fmt, psum_num_tokens_per_expert, topk_weights, routed_scaling_factor, clamp_value, clamped_count, ) return _ref.swiglu_forward( x, psum_num_tokens_per_expert, topk_weights, routed_scaling_factor, clamp_value, clamped_count, ) def swiglu_forward_and_per_token_cast( x: torch.Tensor, fmt: str, num_per_channels: int, psum_num_tokens_per_expert: torch.Tensor | None = None, topk_weights: torch.Tensor | None = None, routed_scaling_factor: float | None = None, use_tma_aligned_col_major_sf: bool = False, round_sf: bool = False, use_packed_ue8m0: bool = False, clamp_value: float | None = None, clamped_count: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Fused SwiGLU + per-token cast.""" if not use_references(): return _compiled.swiglu_forward_and_per_token_cast( x, fmt, num_per_channels, psum_num_tokens_per_expert, topk_weights, routed_scaling_factor, use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0, clamp_value, clamped_count, ) activated = _ref.swiglu_forward( x, psum_num_tokens_per_expert, topk_weights, routed_scaling_factor, clamp_value, clamped_count, ) return _ref.cast( activated, fmt, block_size=(1, num_per_channels), round_sf=round_sf, use_tma_aligned_col_major_sf=use_tma_aligned_col_major_sf, use_packed_ue8m0=use_packed_ue8m0, ) def batched_transpose_weight_sf( sf: Sequence[torch.Tensor], block_size: int, transpose_sf1s: Sequence[bool], outs: Sequence[tuple[torch.Tensor, torch.Tensor] | None], ) -> list[tuple[torch.Tensor, torch.Tensor]]: """Transpose the scale factors of several weights in one launch.""" if not use_references(): return _compiled.batched_transpose_weight_sf(sf, block_size, transpose_sf1s, outs) # The reference allocates its outputs, so the `outs` slots cannot be filled in place. Check # rather than ignore: a caller that pre-allocated would read stale memory. if any(o is not None for o in outs): raise NotImplementedError("the torch reference allocates its own outputs; pre-allocated `outs` cannot be filled") return _ref.batched_transpose_weight_sf(sf, block_size, transpose_sf1s) def unpack_from_e2m1fn_x2(x: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.Tensor: """Decode packed FP4 pairs. Same on both paths -- it is host-side bit manipulation.""" return _ref.unpack_from_e2m1fn_x2(x, out_dtype) def backend() -> str: return "ref" if use_references() else "cuda" def _passthrough(name: str) -> Any: """Resolve a name this module does not adapt from whichever backend is active.""" if not use_references(): return getattr(_compiled, name) return getattr(_ref, name)