sgl-project/sglang · error · ValueError

Shape mismatch: A K={k}, B K={b.shape[1] * 2}

Error message

Shape mismatch: A K={k}, B K={b.shape[1] * 2}

What it means

The fused kernel computes K from the activation as a.shape[1]*2 (FP4 elements are packed two per byte) and from the weight as b.shape[1]*2. If these disagree, the GEMM contraction dimensions are inconsistent and the kernel cannot proceed.

Source

Thrown at python/sglang/kernels/ops/quantization/nvfp4_gemm_swiglu_nvfp4_quant.py:2878

    if ab_dtype != "float4_e2m1fn" or c_dtype != "float4_e2m1fn":
        raise ValueError(
            "nvfp4_gemm_swiglu_nvfp4_quant currently supports NVFP4 input "
            "and output only"
        )
    if a.device.type != "cuda" or b.device.type != "cuda":
        raise ValueError("nvfp4_gemm_swiglu_nvfp4_quant requires CUDA tensors")

    major, minor = get_compute_capability(a.device)
    if major != 10:
        raise ValueError(
            f"nvfp4_gemm_swiglu_nvfp4_quant requires SM100, got SM{major}{minor}"
        )

    m = a.shape[0]
    k = a.shape[1] * 2
    n = b.shape[0]
    if b.shape[1] * 2 != k:
        raise ValueError(f"Shape mismatch: A K={k}, B K={b.shape[1] * 2}")
    if n % 2 != 0:
        raise ValueError(f"Interleaved FC1 N must be even, got {n}")

    l = 1
    n_out = n // 2
    if n_out % sf_vec_size != 0:
        raise ValueError(
            f"Output N={n_out} must be divisible by sf_vec_size={sf_vec_size}"
        )
    scale_n_out = n_out // sf_vec_size
    padded_m = _round_up(m, 128)
    padded_scale_n = _round_up(scale_n_out, 4)

    ab_dtype_cutlass = get_cutlass_dtype(ab_dtype)
    sf_dtype_cutlass = get_cutlass_dtype(sf_dtype)
    c_dtype_cutlass = get_cutlass_dtype(c_dtype)

    if m <= 128:

View on GitHub (pinned to 0132848349)

Solutions

  1. Verify a.shape[1] == b.shape[1] (packed K) before the call
  2. Check that b is in the expected K-major layout ([K_packed, N] per the kernel's expectation) and re-repack weights if needed
  3. Confirm the model config hidden_size matches the checkpoint

Example fix

// before
assert a.shape[1]*2 == b.shape[1]*2
// after
if a.shape[1]*2 != b.shape[1]*2:
    b = repack_b_for_nvfp4(b_raw)  # fix layout
out = nvfp4_gemm_swiglu_nvfp4_quant(a, b, ...)
Defensive patterns

Strategy: validation

Validate before calling

assert a.shape[1] == b.shape[1], (a.shape, b.shape)

Type guard

def k_dims_match(a, b): return a.shape[1]*2 == b.shape[1]*2

Prevention

When it happens

Trigger: Calling nvfp4_gemm_swiglu_nvfp4_quant where a.shape[1]*2 != b.shape[1]*2 — e.g. passing a packed-K activation with an unpacked weight or weights from a different hidden_size.

Common situations: Loading a checkpoint whose hidden_size doesn't match the config, passing scale tensors/weights prepared for a different layer, or accidentally transposing b so its packed-K axis lands in shape[0].

Related errors


AI-assisted analysis of sgl-project/sglang@0132848349 (2026-08-28). Data as JSON: /api/errors/0ff6f38be2131b96. Report an issue: GitHub.