Quantize attention projections (Q, K, V, O) to FP8 E4M3 with static per-tensor scales separately
Graft the FP8 attention weights + scales into the MXFP4 checkpoint

The result uses the standard mxfp4 quant_method with an additional attn_quantization_config field to signal that attention layers should be loaded with FP8 instead of MXFP4.

Model Details

Property	Value
Base model	openai/gpt-oss-120b
Architecture	GptOssForCausalLM (MoE, 128 experts, top-4)
Hidden size	2880
Layers	36 (sliding + full attention)
MoE expert weights	MXFP4 (unchanged from upstream)
Attention projections (Q, K, V, O)	FP8 E4M3, static per-tensor scales
Context length	131072 (YaRN rope scaling)
Checkpoint format	14 safetensor shards (~60GB total)

Quantization Recipe

Step 1: FP8 Attention Calibration

The BF16 base model attention projections were quantized to FP8 E4M3 with static per-tensor weight and input scales using llmcompressor calibration.

Step 2: Weight Grafting

The graft script takes the original MXFP4 checkpoint (MoE experts as packed _blocks/_scales) and replaces the BF16 attention projection weights with the FP8 weights + scales from the calibrated model.

For each safetensor shard containing attention weights:

Replace .weight tensors for Q/K/V/O projections with FP8 E4M3 versions
Inject .weight_scale and .input_scale tensors alongside them
Shards without attention weights are copied verbatim

config.json

json

{
  "quantization_config": {
    "quant_method": "mxfp4",
    "modules_to_not_convert": [
      "model.layers.*.mlp.router",
      "model.embed_tokens",
      "lm_head"
    ]
  },
  "attn_quantization_config": {
    "quant_method": "fp8",
    "activation_scheme": "static"
  }
}

Scripts

FP8 Attention Quantization Script

quantize_fp8_attn.py (click to expand)

python

"""
Quantize GPT-OSS 120B MXFP4 attention projections (q/k/v/o) to FP8 static.

Uses llmcompressor oneshot with wikitext calibration data.
The MLP experts remain MXFP4; only self_attn.{q,k,v,o}_proj get FP8.
"""

import os
import logging

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
logger = logging.getLogger(__name__)

from llmcompressor.modifiers.quantization import QuantizationModifier
from llmcompressor import oneshot

MODEL_PATH = "/path/to/gpt-oss-120b-mxfp4"
OUTPUT_DIR = "/path/to/gpt-oss-120b-mxfp4-fp8attn"

# FP8 static quantization on attention projection layers only.
# "FP8" preset = weights + input_activations in float8_e4m3fn,
# symmetric, per-tensor, static scales via minmax observer.
#
# compressed_tensors _match_name does exact string match on the full module path,
# so we must use "re:" prefix for regex matching against paths like
# "model.layers.0.self_attn.q_proj".
recipe = QuantizationModifier(
    targets=[
        "re:.*self_attn\\.q_proj$",
        "re:.*self_attn\\.k_proj$",
        "re:.*self_attn\\.v_proj$",
        "re:.*self_attn\\.o_proj$",
    ],
    scheme="FP8",
)

logger.info("Starting FP8 quantization of attention projections")
logger.info("Model: %s", MODEL_PATH)
logger.info("Output: %s", OUTPUT_DIR)

model = oneshot(
    model=MODEL_PATH,
    recipe=recipe,
    dataset="wikitext",
    dataset_config_name="wikitext-2-raw-v1",
    splits={"calibration": "test"},
    num_calibration_samples=256,
    max_seq_length=2048,
    output_dir=OUTPUT_DIR,
    trust_remote_code_model=True,
    save_compressed=True,
)

logger.info("Done. Quantized model saved to %s", OUTPUT_DIR)

Graft Script

graft_fp8_attn.py (click to expand)

python

"""
Tensor surgery: Graft FP8 attention projections into the MXFP4 model.

Takes the original MXFP4 model (experts as _blocks/_scales) and replaces
the bf16 attention projection weights with FP8 e4m3fn weights + scales
from the llmcompressor-calibrated model.

Result: MXFP4 experts + FP8 static attention, single compressed-tensors config.
"""

import os
import json
import shutil
import logging

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
log = logging.getLogger(__name__)

import torch
from safetensors import safe_open
from safetensors.torch import save_file

MXFP4_DIR = "/path/to/gpt-oss-120b-mxfp4"
FP8_DIR = "/path/to/gpt-oss-120b-mxfp4-fp8attn"
OUTPUT_DIR = "/path/to/gpt-oss-120b-mxfp4-fp8attn-v2"

ATTN_PROJS = ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.o_proj")
SKIP_DIRS = ("original", "metal")

# Step 1: Collect FP8 attention tensors from the calibrated model
log.info("Loading FP8 attention tensors from %s", FP8_DIR)
with open(os.path.join(FP8_DIR, "model.safetensors.index.json")) as f:
    fp8_index = json.load(f)

fp8_tensors = {}
fp8_shards_loaded = set()

for key, shard in fp8_index["weight_map"].items():
    if not any(proj in key for proj in ATTN_PROJS):
        continue
    if shard not in fp8_shards_loaded:
        shard_path = os.path.join(FP8_DIR, shard)
        with safe_open(shard_path, framework="pt") as sf:
            for k in sf.keys():
                if any(proj in k for proj in ATTN_PROJS):
                    fp8_tensors[k] = sf.get_tensor(k)
        fp8_shards_loaded.add(shard)

log.info("Loaded %d FP8 attention tensors (weights + scales)", len(fp8_tensors))

# Step 2: Create output dir, copy non-safetensor files
os.makedirs(OUTPUT_DIR, exist_ok=True)

for fname in os.listdir(MXFP4_DIR):
    if fname in SKIP_DIRS:
        continue
    src = os.path.join(MXFP4_DIR, fname)
    if os.path.isfile(src) and not fname.endswith(".safetensors"):
        shutil.copy2(src, os.path.join(OUTPUT_DIR, fname))
        log.info("Copied %s", fname)

# Step 3: Process each MXFP4 shard
with open(os.path.join(MXFP4_DIR, "model.safetensors.index.json")) as f:
    mxfp4_index = json.load(f)

shard_keys = {}
for key, shard in mxfp4_index["weight_map"].items():
    shard_keys.setdefault(shard, []).append(key)

new_weight_map = {}

for shard_name, keys in sorted(shard_keys.items()):
    src_path = os.path.join(MXFP4_DIR, shard_name)
    dst_path = os.path.join(OUTPUT_DIR, shard_name)

    has_attn = any(any(proj in k for proj in ATTN_PROJS) for k in keys)

    if not has_attn:
        shutil.copy2(src_path, dst_path)
        for k in keys:
            new_weight_map[k] = shard_name
        log.info("Copied shard %s (no attn weights)", shard_name)
        continue

    tensors = {}
    with safe_open(src_path, framework="pt") as sf:
        for k in sf.keys():
            tensors[k] = sf.get_tensor(k)

    replaced = 0
    added = 0

    for k in list(tensors.keys()):
        if k in fp8_tensors:
            tensors[k] = fp8_tensors[k]
            replaced += 1
        if k.endswith(".weight") and any(proj in k for proj in ATTN_PROJS):
            base = k.removesuffix(".weight")
            for suffix in (".weight_scale", ".input_scale"):
                scale_key = base + suffix
                if scale_key in fp8_tensors and scale_key not in tensors:
                    tensors[scale_key] = fp8_tensors[scale_key]
                    added += 1

    save_file(tensors, dst_path)
    for k in tensors:
        new_weight_map[k] = shard_name
    log.info("Wrote shard %s (replaced %d, added %d scale tensors)", shard_name, replaced, added)

# Step 4: Write updated safetensors index
new_index = {
    "metadata": mxfp4_index.get("metadata", {}),
    "weight_map": dict(sorted(new_weight_map.items())),
}
with open(os.path.join(OUTPUT_DIR, "model.safetensors.index.json"), "w") as f:
    json.dump(new_index, f, indent=2)
log.info("Wrote safetensors index (%d keys)", len(new_weight_map))

# Step 5: Write config.json with dual quantization config
with open(os.path.join(MXFP4_DIR, "config.json")) as f:
    config = json.load(f)

config["quantization_config"] = {
    "modules_to_not_convert": [
        "model.layers.*.mlp.router",
        "model.embed_tokens",
        "lm_head"
    ],
    "quant_method": "mxfp4"
}
config["attn_quantization_config"] = {
    "quant_method": "fp8",
    "activation_scheme": "static"
}

with open(os.path.join(OUTPUT_DIR, "config.json"), "w") as f:
    json.dump(config, f, indent=2)
log.info("Wrote config.json with dual MXFP4+FP8 quantization config")

Evaluation

Metric	BF16	FP8 (this model)
GSM8K accuracy	0.848	0.834
Throughput (4xH200)	1x	~1.5x

Serving

Requires vLLM with GPT-OSS model support that reads attn_quantization_config from the HuggingFace config and overrides the attention layer quant_config with Fp8Config. The MoE layers use the standard MXFP4 path.

vllm serve jasperg662/gpt-oss-120b-mxfp4-fp8attn-v2 --tensor-parallel-size 4

Difference from v1

	v1 (llmcompressor)	v2 (graft)
Tool	llmcompressor	Custom graft script
quant_method	compressed-tensors	mxfp4 + attn_quantization_config
Checkpoint size	73 shards	14 shards

Source

Base model: openai/gpt-oss-120b (MXFP4 MoE)
FP8 attention weights grafted from separately quantized BF16 model