{"record":{"id":"09a467403ccb6a54","repo":"sgl-project/sglang","slug":"mxfp8-fused-prologue-requires-head-dim-aligned-q-k","errorCode":null,"errorMessage":"MXFP8 fused prologue requires head_dim-aligned Q/K/V.","messagePattern":"MXFP8 fused prologue requires head_dim-aligned Q/K/V\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/inkling_attn_prologue.py","lineNumber":92,"sourceCode":"    dq: int,\n    dkv: int,\n    draft_token_num: int,\n    activation: str | None = None,\n    use_residual: bool = True,\n    do_store: bool = True,\n    mxfp8_quant: bool = False,\n    sfk: torch.Tensor | None = None,\n    sfv: torch.Tensor | None = None,\n    page_size: int = 128,\n    log_scaling_tau: torch.Tensor | None = None,\n) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:\n    \"\"\"Returns fresh contiguous (q_normed, k_normed, v_conv) [T, dq/dkv];\n    KV rows are also scattered into k_buf/v_buf at ``loc`` (the attention call\n    should pass save_kv_cache=False).\"\"\"\n    t = qkvr.shape[0]\n    if mxfp8_quant:\n        if dq % 128 != 0 or dkv % 128 != 0:\n            raise ValueError(\"MXFP8 fused prologue requires head_dim-aligned Q/K/V.\")\n        if sfk is None or sfv is None:\n            raise ValueError(\"MXFP8 fused prologue requires K/V scale buffers.\")\n        sf_shape = (k_buf.shape[0] // page_size, dkv // 128, 32, page_size // 32, 4)\n        if sfk.shape != sf_shape or sfv.shape != sf_shape:\n            raise ValueError(\n                \"MXFP8 fused prologue requires interleaved K/V scale buffers \"\n                f\"with shape {sf_shape}, got {tuple(sfk.shape)} and {tuple(sfv.shape)}.\"\n            )\n        if not sfk.is_contiguous() or not sfv.is_contiguous():\n            raise ValueError(\n                \"MXFP8 fused prologue requires contiguous interleaved SFK/SFV.\"\n            )\n        q_out = torch.empty(t, dq, dtype=torch.float8_e4m3fn, device=qkvr.device)\n        sfq_u8 = torch.empty(\n            (t, dq // 128, 128 // 32), dtype=torch.uint8, device=qkvr.device\n        )\n        sfk_u8 = sfk.view(torch.uint8)\n        sfv_u8 = sfv.view(torch.uint8)","sourceCodeStart":74,"sourceCodeEnd":110,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/inkling_attn_prologue.py#L74-L110","documentation":"inkling_attn_prologue_verify is the verification (target-verify) variant of the fused attention prologue; when mxfp8_quant=True it quantizes Q/K/V to MXFP8, which requires the 128-byte-blocked microscaling format — hence dq and dkv must both be multiples of 128. The check dq % 128 != 0 or dkv % 128 != 0 rejects non-aligned head dims before allocating float8_e4m3 buffers, because scale factors are computed per 32-element block within 128-element superblocks.","triggerScenarios":"Calling inkling_attn_prologue_verify with mxfp8_quant=True on a model whose head dims (dq, dkv) are not multiples of 128 — e.g. dq=64, dkv=128, or dq=96 — such as smaller Inkling configurations or head_dim-64 attention layers routed into the fused MXFP8 path.","commonSituations":"Enabling MXFP8 quantization for a model with head_dim 64/80/96; a config knob like quantization_kv_cache=mxfp8 applied globally across layers with heterogeneous head dims; new model checkpoints with non-128-aligned dims tested against the fused prologue.","solutions":["Disable mxfp8_quant (fall back to the non-quantized prologue) for this layer/model","Use a model configuration with head dims that are multiples of 128 (e.g. 128 or 256)","Gate the MXFP8 path on (dq % 128 == 0 and dkv % 128 == 0) in the caller so it degrades gracefully"],"exampleFix":"# before\nq, k, v = inkling_attn_prologue_verify(qkvr, ..., mxfp8_quant=True)\n# after\nuse_mxfp8 = (dq % 128 == 0 and dkv % 128 == 0)\nq, k, v = inkling_attn_prologue_verify(qkvr, ..., mxfp8_quant=use_mxfp8)","handlingStrategy":"validation","validationCode":"mxfp8_quant = mxfp8_quant and dq % 128 == 0 and dkv % 128 == 0","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Check head-dim alignment at model load time and disable MXFP8 once per layer, not per call","Document 128-alignment as a hard requirement of the MXFP8 fused prologue"],"tags":["mxfp8","quantization","head-dim","attention-prologue","inkling"],"backgroundTag":"quantization-shape-constraint","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}