[BUG][PVC/BMG] DiffusionGemma XPU segfault (CCS write fault) in model.generate() — CPU PASS, XPU FAIL

[weishi-deng]

[BUG][PVC/BMG] DiffusionGemma XPU segfault (CCS write fault) in model.generate() — CPU PASS, XPU FAIL

Tracking issue for the DiffusionGemma-26B-A4B XPU failure, same device-side fault class as OmniVoice #5115.
Model is in the Q3'26 client validation set (bucket: XPU op/kernel fail, model-driven — not a PyTorch framework defect, not OOM, not CUDA-bias).

Environment

• GPU: Intel XPU, compute-runtime (drm_neo); PVC/BMG class, 4× visible on the eval host (same host as [BUG][PVC/BMG] gather kernel out-of-bound assert in OmniVoice TTS model.generate() — CPU PASS, XPU FAIL #5115)
• torch: 2.15.0.dev20260824+xpu (latest nightly, bn_nightly env)
• transformers: 5.16.1
• bitsandbytes: 0.50.1 (NF4 4-bit load)
• Python: 3.10 (conda env bn_nightly)
• OS: Linux x86_64
• Model package: google/diffusiongemma-26B-A4B-it (DiffusionGemmaForBlockDiffusion)

Minimal Reproducer

infer.py (official API, bitsandbytes NF4 4-bit load). The caching_allocator_warmup monkeypatch below
works around a separate transformers loader bug (it allocates a full BF16 buffer during warmup and OOMs on a
48 GB card); with it, weights load cleanly (1047/1047 tensors, ~13 GB) and the crash moves to generate().

#!/usr/bin/env python3
"""Official API with bitsandbytes 4-bit loading.
Source: https://huggingface.co/google/diffusiongemma-26B-A4B-it
"""
import json
from pathlib import Path

import torch
import transformers.modeling_utils as modeling_utils

OUT = Path(__file__).resolve().parent
MODEL_ID = "google/diffusiongemma-26B-A4B-it"

# Transformers warmup incorrectly requests a full-size BF16 buffer for this
# quantized model before loading weights. Warmup only affects load speed.
modeling_utils.caching_allocator_warmup = lambda *args, **kwargs: None


def main():
    from transformers import (
        AutoProcessor,
        BitsAndBytesConfig,
        DiffusionGemmaForBlockDiffusion,
    )

    processor = AutoProcessor.from_pretrained(MODEL_ID)
    quantization_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True,
    )
    model = DiffusionGemmaForBlockDiffusion.from_pretrained(
        MODEL_ID,
        quantization_config=quantization_config,
        dtype=torch.bfloat16,
        device_map="xpu",
    ).eval()

    message = [{"role": "user", "content": "Why is the sky blue?"}]
    inputs = processor.apply_chat_template(
        message, tokenize=True, add_generation_prompt=True,
        return_dict=True, return_tensors="pt",
    ).to("xpu")
    output = model.generate(**inputs, max_new_tokens=8)   # <-- segfault here
    text = processor.decode(output[0], skip_special_tokens=False)
    print(f"Output: {text}")


if __name__ == "__main__":
    main()

CPU control (proves XPU-specific): same MODEL_ID + prompt, BF16 on CPU (infer_cpu.py,
device_map="cpu", 52 GB host RAM of 436 GB). generate() completes and returns coherent text
("The sky is blue because of a phenomenon called Rayleigh scattering..."). The model/transformers code path is correct.

Expected Behavior

model.generate() completes on XPU and produces coherent text, exactly as it does on CPU.

Actual Behavior

Process aborts with a GPU-side segfault during the first generate() forward pass (after weights load successfully):

Segmentation fault from GPU at 0xff00000bdaed4000, ctx_id: 1 (CCS) type: 0 (NotPresent), level: 1 (PDE), access: 1 (Write), banned: 1, aborting.
Segmentation fault from GPU at 0xff00000bdaed4000, ctx_id: 1 (CCS) type: 0 (NotPresent), level: 1 (PDE), access: 1 (Write), banned: 1, aborting.
Abort was called at 288 line in file:
./shared/source/os_interface/linux/drm_neo.cpp

CCS = compute command streamer; Write fault, page-not-present. This is the GPU driver / compute-runtime layer
(Intel drm_neo.cpp), not Python/PyTorch.

What We Ruled Out

Test	Result
CPU generation (BF16)	✅ Always passes, coherent output
OOM as root cause	❌ Not the cause — weights load 1047/1047 tensors (~13 GB); the loader warmup OOM is a separate transformers bug with a workaround (see Reproducer)
Model / transformers code path	✅ Verified correct end-to-end on CPU
FP16 / BF16 compute dtype	N/A — model is BnB NF4 with bf16 compute; crash is at device layer, not dtype
TORCH_COMPILE_DISABLE=1	❌ Not yet tried (suspected contributor — default torch.compile path)
Different transformers version	❌ Not yet tried

Analysis

The abort originates from the Intel compute-runtime (drm_neo.cpp) CCS write fault — i.e. a device-side
addressing/dispatch fault below PyTorch. This is the same failure class as OmniVoice #5115 (CCS write
abort / device-side fault), strongly suggesting a shared root cause in a SYCL kernel dispatch or a missing
dependency barrier between kernel launches for this model's forward pass (block-diffusion + attention).

Isolating the exact op is hard because:

1. The crash is inside model.generate(), which dispatches hundreds of kernels.
2. Any Python-level instrumentation (hooks, .cpu(), thread spawn) changes kernel scheduling and prevents reproduction.

Request

1. Help identify which SYCL kernel dispatch triggers the CCS write fault (debug tracing of the drm_neo CCS write path).
2. Determine whether this shares the root cause with OmniVoice [BUG][PVC/BMG] gather kernel out-of-bound assert in OmniVoice TTS model.generate() — CPU PASS, XPU FAIL #5115 (same device-side fault class).
3. Suggest an op-level workaround (e.g., TORCH_COMPILE_DISABLE=1, or isolating the first forward op that triggers the CCS write fault) so the model can be unblocked on XPU.