flashinfer.quantization.silu_and_mul_nvfp4_quantize¶
- flashinfer.quantization.silu_and_mul_nvfp4_quantize(input: Tensor, global_scale: Tensor, sf_vec_size: int = 16, is_sf_swizzled_layout: bool = True, is_sf_8x4_layout: bool = False, enable_pdl: bool | None = None) Tuple[Tensor, Tensor]¶
Apply SwiGLU and NVFP4 quantization in one CuTe-DSL kernel.
Computes
silu(input[..., :K]) * input[..., K:]before quantization. Requires CuTe-DSL and an SM100+ GPU.- Parameters:
input (torch.Tensor) – Contiguous fp16/bf16 tensor of shape
[..., 2K]. Leading dimensions are flattened into M.global_scale (torch.Tensor) – Float32 scale of shape
[1], typically(448 * 6) / silu_and_mul(input).abs().max().sf_vec_size (int) – Scale vector size. Only 16 is supported.
is_sf_swizzled_layout (bool) – Use a swizzled scale layout. Defaults to True.
is_sf_8x4_layout (bool) – Use the 8x4 rather than 128x4 swizzled layout.
enable_pdl (bool, optional) – Enable Programmatic Dependent Launch. Auto-detected when None.
- Returns:
Packed FP4 values of shape
[M, K/2]and their scale factors.- Return type:
Tuple[torch.Tensor, torch.Tensor]