sgl-project/sglang · error · ValueError

Fused QK-Norm + RoPE kernel only supports float16/bfloat16,

Error message

Fused QK-Norm + RoPE kernel only supports float16/bfloat16, but got {img_q.dtype}

What it means

JoyImage's fused QK-Norm + RoPE path only accepts half-precision query tensors. The Triton/fused kernel that applies RMSNorm and rotary position embeddings is compiled for float16/bfloat16 layouts, so any other dtype (float32, float64) is rejected before the kernel launch.

Source

Thrown at python/sglang/multimodal_gen/runtime/models/dits/joy_image.py:259

        ) = self.txt_mod(vec)

        # Image attention
        img_modulated = self.fused_modulate_img_norm1(
            img, shift=img_mod1_shift, scale=img_mod1_scale
        )
        img_qkv, _ = self.img_attn_qkv(img_modulated)
        img_q, img_k, img_v = rearrange(
            img_qkv, "B L (K H D) -> K B L H D", K=3, H=self.local_heads_num
        )

        if vis_freqs_cis is None:
            raise ValueError(
                "vis_freqs_cis is required for fused QK-Norm + RoPE kernel"
            )
        if not (isinstance(vis_freqs_cis, torch.Tensor) and vis_freqs_cis.dim() == 2):
            raise ValueError("vis_freqs_cis must be a 2D cos_sin_cache tensor")
        if img_q.dtype not in (torch.float16, torch.bfloat16):
            raise ValueError(
                f"Fused QK-Norm + RoPE kernel only supports float16/bfloat16, but got {img_q.dtype}"
            )
        img_q = img_q.contiguous()
        img_k = img_k.contiguous()
        img_q, img_k = apply_qk_norm_with_optional_rope(
            q=img_q,
            k=img_k,
            q_norm=self.img_attn_q_norm,
            k_norm=self.img_attn_k_norm,
            head_dim=img_q.shape[-1],
            cos_sin_cache=vis_freqs_cis,
            is_neox=False,
            allow_inplace=True,
        )
        img_q, img_k = img_q.to(img_v), img_k.to(img_v)

        # Text attention
        txt_modulated = self.fused_modulate_txt_norm1(

View on GitHub (pinned to 0132848349)

Solutions

  1. Cast model/inputs to bfloat16 or float16 before forward: model.to(torch.bfloat16), img_q = img_q.to(torch.bfloat16)
  2. Ensure autocast/sampler dtype matches the configured dtype (e.g. --dtype bfloat16 in server args)
  3. If fp32 is required, bypass the fused apply_qk_norm_with_optional_rope path with an unfused reference implementation (compute q-norm/k-norm + RoPE in eager PyTorch)

Example fix

// before
out = dit(hidden_states, encoder_hidden_states=cond, ...)  # img_q is float32

// after
dit = dit.to(torch.bfloat16)
hidden_states = hidden_states.to(torch.bfloat16)
out = dit(hidden_states, encoder_hidden_states=cond.to(torch.bfloat16), ...)
Defensive patterns

Strategy: validation

Validate before calling

assert img_q.dtype in (torch.float16, torch.bfloat16), f"need fp16/bf16, got {img_q.dtype}"

Type guard

def is_half_precision(t: torch.Tensor) -> bool:\n    return t.dtype in (torch.float16, torch.bfloat16)

Try / catch

try:\n    out = dit(...)\nexcept ValueError as e:\n    if 'float16/bfloat16' in str(e):\n        dit = dit.to(torch.bfloat16); out = dit(...)\n    else:\n        raise

Prevention

When it happens

Trigger: Calling JoyImage forward where the img_q tensor (image-branch query after qkv split) has a dtype other than torch.float16 or torch.bfloat16, e.g. running the DiT in float32 (debugging, CPU fallback) or with autocast disabled while model weights were kept in fp32.

Common situations: Developer forces float32 for numerical debugging or runs on hardware without bf16; upcasting tensors before forward; mixed setups where text branch is fp16 but vision tensors were cast to fp32.

Related errors


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