sgl-project/sglang · error · ValueError

num_heads must be divisible by num_epi_subtiles

Error message

num_heads must be divisible by num_epi_subtiles

What it means

The fp8 paged MQA logits CuteDSL kernel's TMA-based epilogue splits the head dimension into num_epi_subtiles for SMEM staging, which requires the split to be exact. The kernel __init__ raises ValueError when num_epi_subtiles > 1 and num_heads % num_epi_subtiles != 0.

Source

Thrown at python/sglang/kernels/ops/attention/cutedsl_fp8_paged_mqa_logits.py:250

        self.smem_subpartition_opt = smem_subpartition_opt
        self.max_w_in_reg = max_w_in_reg
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.next_n = next_n
        self.N = next_n * num_heads
        self.num_sms = num_sms
        self.enable_pdl = enable_pdl
        self.num_epi_subtiles = num_epi_subtiles
        self.epi_dtype = epi_dtype
        self.epi_bytes = 2 if epi_dtype == cutlass.Float16 else 4
        # sW stage stride padded to 128-byte SMEM alignment for TMA bulk copy.
        # Without padding, e.g. fp16 + N=32 gives 64B per stage, so stage 1
        # at +64 would be misaligned (TMA requires 128-byte aligned SMEM dest).
        w_stage_bytes = self.N * self.epi_bytes
        self.w_stage_stride = ((w_stage_bytes + 127) // 128 * 128) // self.epi_bytes
        self.output_dtype = output_dtype
        if num_epi_subtiles > 1 and num_heads % num_epi_subtiles != 0:
            raise ValueError("num_heads must be divisible by num_epi_subtiles")
        if (num_heads // num_epi_subtiles) % 4 != 0:
            raise ValueError(
                "num_heads // num_epi_subtiles must be divisible by 4 (FMA unroll granularity)"
            )
        self.num_groups = 2

        self.num_math_threads = 256
        self.num_specialized_threads = 128
        self.threads_per_cta = 384
        self.num_math_warps = 8
        self.tma_warp_base = 8
        self.umma_warp_base = 10

        self.num_q_stages = 3  # 3 stages for Q pipelining across batch sequences

        # TMEM: 512 columns total, each group needs N columns per UMMA stage
        # max_umma_stages = 512 // (2 * N)
        TMEM_COLS = 512

View on GitHub (pinned to 0132848349)

Solutions

1. Pick num_epi_subtiles as a power-of-2 divisor of num_heads (1, 2, 4, 8...)
2. Use num_epi_subtiles=1 to bypass the constraint (with possible perf cost)
3. Add an assertion at the caller: assert num_heads % num_epi_subtiles == 0 before constructing the kernel

Example fix

# before
kernel = MqaLogitsKernel(num_heads=48, num_epi_subtiles=5)  # ValueError
# after
kernel = MqaLogitsKernel(num_heads=48, num_epi_subtiles=4)  # 48 % 4 == 0

assert num_epi_subtiles == 1 or num_heads % num_epi_subtiles == 0, (
    f'num_heads ({num_heads}) must be divisible by num_epi_subtiles ({num_epi_subtiles})')

When it happens

Trigger: Instantiating the MQA logits kernel class with num_epi_subtiles > 1 while num_heads is not a multiple of it — e.g. num_heads=32, num_epi_subtiles=6; or a caller computing num_epi_subtiles from SMEM budget without aligning it to the model's head count.

Common situations: Tuning a new GQA/MQA model config (odd head counts like 48 with an epi subtile of 5) where the epilogue tiling was tuned for a different model; passing num_epi_subtiles derived from tile-size heuristics rather than a divisor of num_heads.

Related errors

• num_heads // num_epi_subtiles must be divisible by 4 (FMA un
• Unexpected initial_state_source shape: {initial_state_source
• Unexpected A_log shape: {A_log.shape}; expected numel={HV}
• Unexpected dt_bias shape: {dt_bias.shape}; expected numel={H
• Unexpected a shape for varlen: {a.shape}

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