Quantization and kernel policies

Configure quantization during loading, layer matching, weight precision, and execution kernels with a kernel policy YAML file.

Configure quantization and kernel selection in one kernel policy YAML file. quantization chooses the precision for each layer, and the top-level rules section selects execution kernels. After loading an unquantized FP16, BF16, or FP32 checkpoint, the system applies PTQ to the weights according to your configuration.

Load a kernel policy

Save the configuration in the next section as kernel-policy.yaml. In a Python environment with PhyAI and humming-kernels installed, specify the file through an environment variable:

PHYAI_KERNEL_CONFIG=./kernel-policy.yaml python run_model.py

Replace run_model.py with your model’s launch script and keep its existing arguments, including the checkpoint path. Run the model in BF16 or FP16.

In Python, pass the configuration containing KernelConfig to EngineArgs(config=config):

from phyai.engine_config import EngineConfig, KernelConfig

config = EngineConfig(
    kernel=KernelConfig(config_path="kernel-policy.yaml")
)

In a server YAML file, add this top-level field:

kernel_policy: ./kernel-policy.yaml

Relative paths in a server YAML file resolve from that file’s directory. Relative paths in environment variables and Python resolve from the process working directory. Choose one policy source for each entry point; do not set both kernel_policy and config.kernel.config_path.

PHYAI_KERNEL_CONFIG overrides Python and server YAML settings when the Engine starts. Clear an old environment setting before switching to either of those entry points:

unset PHYAI_KERNEL_CONFIG

Customize quantized layers and precision

The following complete policy uses INT4 for MLP weights and INT8 for attention projection weights, while keeping layer 0 at its original precision. The example assumes layer prefixes start with model.layers.<layer-number>. Replace them with your model’s prefixes.

schema: phyai.kernel/v1
profile: static

quantization:
  method: rtn
  stage: load
  default: null
  rules:
    - match:
        glob: "model.layers.0.*"
      scheme: null
    - match:
        glob: "model.layers.*.mlp.*"
      weight:
        dtype: int4
        block_shape: [1, 128]
      input: null
    - match:
        glob: "model.layers.*.self_attn.*"
      weight:
        dtype: int8
        block_shape: [1, 128]
      input: null

rules:
  - id: humming-linear
    match:
      op: gemm
      quant.layout: humming
    restrict_to: "humming.gemm.*"

quantization.rules are checked from top to bottom and stop at the first match. Put exclusions first. scheme: null skips a layer, and default: null keeps the configured precision for unmatched layers. PTQ during loading currently uses method: rtn and stage: load.

Each quantization rule’s match must use exactly one selector:

Selector	What it matches
name	The full layer prefix, or its last segment
glob	A wildcard pattern for the layer prefix, such as model.layers.*.mlp.*
regex	A regular expression for the layer prefix; add ^ and $ to require a full match
module_cls	A substring of the layer class name, such as RowParallelLinear

These names are the prefix values passed when creating PhyAI Linear layers. They differ from checkpoint tensor keys and may differ from named_modules() paths. Check where your model creates its Linear layers to confirm their prefixes. If the model fuses q_proj/k_proj/v_proj, match the fused qkv_proj layer. Likewise, match gate_up_proj for fused gate_proj/up_proj weights.

Common precision fields:

Field	Usage
weight.dtype	int4, int8, fp8_e4m3, fp8_e5m2, mxfp4, or nvfp4
weight.block_shape: [1, 128]	Group size 128; specifying this selects block granularity automatically
weight.granularity: per_channel	One set of weight quantization parameters per output channel; omit block_shape
weight.granularity: per_tensor	One set of quantization parameters for the whole weight tensor; omit block_shape
input: null	Keep BF16/FP16 activations
input.dtype	The activation type for dynamic A8; use a weight/activation pair from the table below

Common combinations are listed below. null means input: null, which keeps BF16/FP16 activations. For A8, also set granularity: per_channel and dynamic: true.

weight.dtype	Weight granularity	input.dtype, or input: null
int4	block_shape: [1, 128]	null, int8, fp8_e4m3
int8	block_shape: [1, 128]	null, int8
fp8_e4m3	granularity: per_channel	null, fp8_e4m3
fp8_e5m2	granularity: per_channel	null, fp8_e5m2
fp8_e4m3	block_shape: [128, 128]	fp8_e4m3
fp8_e5m2	block_shape: [128, 128]	fp8_e5m2
mxfp4	block_shape: [1, 32]	null
nvfp4	block_shape: [1, 16]	null

For example, replace a rule’s weight and input fields with these settings to use FP8 E4M3 block A8:

weight:
  dtype: fp8_e4m3
  block_shape: [128, 128]
input:
  dtype: fp8_e4m3
  granularity: per_channel
  dynamic: true

A8 still accepts and returns BF16/FP16 tensors. Activations are quantized dynamically per token during each call. FP8 A8 requires NVIDIA SM89 or newer; the other combinations in the table require SM80 or newer when running in BF16. MXFP4 and FP8 E5M2 A8 currently require BF16 inputs and outputs. PTQ uses symmetric weight quantization, and A8 activations must use dynamic quantization.

After changing quantization, create a new Engine and reload the weights. Editing the YAML does not change the precision of an existing instance. If you omit the entire quantization section, the system uses the checkpoint and model’s existing configuration. Load an already quantized checkpoint through its existing loading path.

Customize execution kernels

The top-level rules section configures execution kernels. You can limit candidates by operation, role, dtype, device, and call shape. For example, append this attention rule to the top-level rules list in the file above:

- id: paged-attention-fa2
  priority: 10
  match:
    op: attention_paged
    device.arch: sm90
  restrict_to: flashinfer.attention_paged.fa2

Quantization rules and kernel rules use different matching methods:

Configuration	Selection behavior
quantization.rules	Matches layer prefixes/classes; the first match applies
Top-level rules	Matches call conditions and chooses the highest priority; simultaneous matches at the same priority raise an error
Top-level overrides	Takes precedence over rules; use: full-kernel-ID forces a kernel and raises an error if it is incompatible

All fields in a kernel’s match must match. Common examples include dtype.input: bf16, quant.format: int4, device.arch: ">=sm80", and shape.M: "<=64". Here, role is an operation role. Layer prefixes belong in quantization rules; kernel match does not accept their glob or regex selectors.

Each ordinary kernel rule must choose one candidate setting:

• restrict_to: A full kernel ID or an ID wildcard pattern, such as "humming.gemm.*".
• prefer: An ordered list of full kernel IDs, such as [flashinfer.gemm.bf16, torch.gemm.bf16]. Other ordinary implementations are excluded from the candidates; compatible reference fallbacks are allowed by default.

To disable reference fallbacks, set defaults: {fallback: error} at the top level of the policy. Humming quantized weights require a compatible Humming kernel. Forcing an ordinary BF16 GEMM for them will fail.

params passes parameters that the selected kernel supports. For example, add this to the Humming rule:

params:
  compute_config:
    use_batch_invariant: true

Parameters from different rules are not merged. To change the weight format, edit quantization. Changing only prefer, restrict_to, or params does not requantize the weights.

Check the configuration

List the GEMM kernel IDs in your installed version:

python -c 'from phyai.kernel import build_catalog; print("\n".join(row.kernel_id for row in build_catalog().impls("gemm")))'

Then parse the file and check that a specific layer matches the settings you intended:

from phyai.kernel import build_catalog, load_policy

policy = load_policy("kernel-policy.yaml", build_catalog())
prefix = "model.layers.1.self_attn.o_proj"
print(policy.quant_plan.resolve(prefix))

With the custom policy above, this layer should show INT8. Changing the layer number to 0 returns None. Parsing does not verify checkpoint contents, GPU capabilities, or model accuracy. Follow it with loading and inference through your normal model entry point.

Problem	What to check
Editing the YAML has no effect	Whether PHYAI_KERNEL_CONFIG is still set, and whether you created a new Engine
A target layer is not quantized	Whether its prefix matches, or an earlier scheme: null rule excludes it
A weight shape is unsupported	The local dimensions after TP partitioning and the group/block size; group size 128 requires local K to be divisible by 128
No kernel is available	GPU architecture, BF16/FP16 restrictions, and whether rules/overrides exclude compatible implementations