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]