name: nvfp4a16_attnbf16
scheme: NVFP4A16  # FP4 weights, bf16 activations — skips activation quant for much lower KLD
engine: llmcompressor

# Variant of `nvfp4a16.yaml` that additionally keeps the ENTIRE self-attention
# block in bf16, leaving only the MLPs quantised — the same lever measured on
# ThinkingCap for FP8 (landmine 36), applied to the W4A16 path.
#
# The motivating fact is architectural, not scheme-specific: `attn_output_gate:
# true` fuses the attention output gate into `q_proj`, making it
# [2*heads*head_dim, hidden] instead of [heads*head_dim, hidden]. Half that
# tensor is a multiplicative per-head gate on what attention writes into the
# residual stream, and only the full_attention layers (a quarter of the stack)
# carry long-range retrieval. Quantisation error on a multiplicative gate
# behaves worse than on an additive projection.
#
# Cost/benefit differs from the FP8 case and should be re-measured, not assumed:
# at W4A16 the attention block would otherwise compress 4x rather than 2x, so
# holding it in bf16 costs proportionally more on-disk than the +4.6% measured
# for FP8. The KLD win is expected to be larger too, since NVFP4A16's error
# floor is 2-3x FP8's. Ship whichever the measured KLD/size tradeoff justifies —
# `nvfp4a16.yaml` remains the standard build.

calibration:
  dataset: neuralmagic/calibration
  config: LLM
  split: train
  num_samples: 128
  max_seq_length: 2048

ignore:
  - lm_head
  - "re:.*visual.*"
  - "re:.*linear_attn.*"   # entire SSM block kept in bf16 — same rationale as the standard build
  - "re:.*self_attn.*"     # THE VARIANT: q/k/v/o_proj too, incl. the output gate fused into q_proj
  - "re:.*mtp.*"
  # Note: dense base. On MoE bases also add "re:.*mlp.gate$" and
  # "re:.*mlp.shared_expert_gate$" — no-ops on dense models.

export:
  save_compressed: true