[Issue]: Inferencing GPT-OSS 20B with FlashAttentionV2 in ROCm7.0 on MI300X results in unusable & degenerate repetitive output

[geozhai]

Problem Description

Used the instructions in the ROCm documentation to attempt to inference GPT-OSS 20B with CK FlashAttentionV2

Initially, in the ROCm6.4.3 container, I tried with the "eager mode" execution for attention, and it worked. The output was coherent, and the throughput was about ~30 tokens per second. However, when I try to inference with FlashAttentionV2, it hangs indefinitely, no output is generated

Then with the ROCm7.0 release, I tried the same code snippet with FlashAttentionV2 with the same prompt, and I am getting the following output with a lower throughput of ~21 tokens per second:

today is from I is from I have I and I with I and I and I and I and I and I and I and I and I and I and I and I and I and
I and I and I and I and I and I and I and I and I and I and I and I and I and I and I and I and I and I and I and I and I 
and I and I and I and I and I and I and I and I and I and I and I and I and I and I and I and I and I and I and I and I 
and I and I and I and I and I and I and I and I and I and I and I and I and I and I and I and I and i and i and i and i 
and i and i and i and i and i and i and i and i and i and i and i and i and i and i and i and i and i and i and i and i 
and i and i and i and i and i and i and i and i and i and i and i and i and i and i and i and i and i and i and i and i 
and i and i and i and i and i and i and i and i and i and

OR

today is i is i is i i i is i of i i j i of i j i of i j of i j i of i j of i j i of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j as i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i j of i

The issue: unusable & degenerate repetitive output and FlashAttentionV2 resulting in lower throughput than eager mode execution. Please fix

Operating System

Ubuntu 22.04.5 LTS (Jammy Jellyfish)

CPU

AMD Ryzen 9 7950X 16-Core Processor

GPU

MI300X

ROCm Version

7.0.51831-a3e329ad8

ROCm Component

composable_kernel

Steps to Reproduce

Code snippet to reproduce:

from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig

tok = AutoTokenizer.from_pretrained(model_path, 
                                    use_fast = True, 
                                    trust_remote_code=True)

config = AutoConfig.from_pretrained(model_path)

gpt_oss = AutoModelForCausalLM.from_pretrained(model_path, 
                                                  device_map=device,
                                                  config = config, 
                                                  offload_folder=None,
                                                  low_cpu_mem_usage=True,
                                                  torch_dtype="auto",
                                                  attn_implementation="flash_attention_2",
                                                  trust_remote_code=True, 
                                                  local_files_only=True).eval()

prompt = "today is the day of the Lord, and we are the"
inputs = tok(prompt, return_tensors="pt").to(gpt_oss.device)
gen = gpt_oss.generate(**inputs, max_new_tokens=256)
print(tok.decode(gen[0], skip_special_tokens=True))

(Optional for Linux users) Output of /opt/rocm/bin/rocminfo --support

No response

Additional Information

Python version 3.10.12 --> this is what is shipped with the ROCm 7.0 container, the specific tag is "rocm7.0_ubuntu_22.04_vllm_0.10.1_instinct_20250915"

Versioning of relevant libraries:

transformers -> /opt/venv/lib/python3.10/site-packages/transformers/__init__.py 4.56.1
accelerate -> /opt/venv/lib/python3.10/site-packages/accelerate/__init__.py 1.10.1
safetensors -> /opt/venv/lib/python3.10/site-packages/safetensors/__init__.py 0.6.2
kernels -> /opt/venv/lib/python3.10/site-packages/kernels/__init__.py 0.10.1
triton -> /opt/venv/lib/python3.10/site-packages/triton/__init__.py 3.4.0+rocm7.0.2.git56765e8c
flash_attn -> /opt/venv/lib/python3.10/site-packages/flash_attn/__init__.py 2.8.0.post2

[juliantorr-es]

ROCm FlashAttention V2 / GPT-OSS 20B Degenerate Output — Root Cause Analysis

Issue

#158

GPT-OSS 20B with FlashAttentionV2 on ROCm 7.0 / MI300X produces degenerate repetitive
output. Eager mode works correctly. The output pattern:

today is from I is from I have I and I with I and I and I...

Model Architecture (GPT-OSS 20B)

• MoE: 20.91B params, 3.61B active
• GQA: 64 query heads, 8 KV heads (ratio 8:1)
• Hidden: 2880, Head dim: 64
• Alternating full_attention / sliding_attention layers (24 total)
• Sliding window: 128
• Attention bias: true
• YaRN RoPE (factor=32, original_max=4096)
• MXFP4 quantized

Kernel Chain

HuggingFace GPT_OSS_ATTENTION.forward()
  → ALL_ATTENTION_FUNCTIONS["flash_attention_2"]
    → flash_attn_varlen_func(q, k, v, causal=True, window_size=(128,0))
      → flash_attn_gpu.varlen_fwd()  [torch custom op]
        → mha_fwd()  [CK wrapper: csrc/flash_attn_ck/mha_fwd.cpp]
          → mask_info::decode("b:128,0", seqlen_q, seqlen_k)
            → GenericAttentionMask<masking=true, local=true>(y=129, x=1, y_total=seqlen_q, x_total=seqlen_k)
              → ComposedAttention::LogitsMask()
                → IsOutOfBound(i_y, i_x)

Root Cause: Bottom-Right Mask Alignment During Decode

File: csrc/flash_attn_ck/mha_fwd.cpp, line 219-222:

if (is_causal) {
    window_size_right = 0;
    std::string mask_identify = "b:" + std::to_string(window_size_left) + "," + "0";
    mask = mask_info::decode(mask_identify, seqlen_q, seqlen_k);
}

The prefix "b:" selects bottom-right mask alignment (is_top_left=false).

During Prefill (works)

seqlen_q ≈ seqlen_k (both ~prompt_len). Bottom-right alignment:

x_tmp = x_total - y_total ≈ 0
y_tmp = y_total - x_total ≈ 0
x = 1 + right_size + x_tmp = 1 + 0 + 0 = 1
y = 1 + left_size + y_tmp = 1 + 128 + 0 = 129

Produces valid range [i_y - 128, i_y) — correct for causal sliding window of 128.
Identical to top-left alignment when seqlen_q == seqlen_k.

During Decode (broken)

seqlen_q=1 (single generated token), seqlen_k=total_prompt_len.

x_tmp = x_total - y_total = seqlen_k - 1  (large positive)
y_tmp = y_total - x_total = 1 - seqlen_k  (large negative)
x = 1 + 0 + (seqlen_k - 1) = seqlen_k
y = 1 + 128 + (1 - seqlen_k) = 130 - seqlen_k  (negative for seqlen_k > 130)

For seqlen_k=256: y = -126. For i_y=0 (single query):

x_start = -(-126) + 0 + 1 = 127
x_end = min(0 + 256, 256) = 256
return i_x < 127 → mask i_x=0..126

Result: the first 127 KV positions are invisible. The model can only attend to
tokens 127-255 during generation. This corrupts every decode step.

Why "today is" appears in output

Prefill (all prompt tokens simultaneously) works because seqlen_q ≈ seqlen_k,
and bottom-right ≈ top-left. The model correctly processes the prompt. But each
autoregressive decode step computes attention against a corrupted mask.

Why This Wasn't Caught

• Standard LLaMA-style models with GQA don't use sliding window, so they test
  only full causal or no mask
• Without sliding_window (left=128), left_size = -1 which takes the default
  (y_total-1 for top-left or x_total-1 for bottom-right), masking the alignment
  difference even during decode
• The sliding_window=128 parameter exposes the alignment bug specifically
• Upstream (NVIDIA CUDA) flash-attn uses top-left alignment. The ROCm port
  switched to bottom-right, likely because ROCm CK's internal mask conventions
  differ, but the asymmetry between prefill and decode was missed

Fix

In csrc/flash_attn_ck/mha_fwd.cpp, change "b:" to "t:" on line 222:

// Before:
std::string mask_identify = "b:" + std::to_string(window_size_left) + "," + "0";
// After:
std::string mask_identify = "t:" + std::to_string(window_size_left) + "," + "0";

This selects top-left alignment, matching HuggingFace causal attention convention.

Alternatively, use the correct alignment based on the context:

// More correct: pass alignment parameter
bool is_top_left = true;  // HuggingFace causal convention
std::string prefix = is_top_left ? "t:" : "b:";
std::string mask_identify = prefix + std::to_string(window_size_left) + "," + "0";

The same fix should be applied in mha_bwd.cpp for the backward pass alignment.

Verification

After the fix, for causal sliding window during decode (seqlen_q=1, seqlen_k=256):

// Top-left alignment
x_tmp = 0
y_tmp = 0
x = 1 + 0 + 0 = 1
y = 1 + 128 + 0 = 129

// For i_y=0 (single query token):
x_start = -129 + 0 + 1 = -128
x_end = min(0 + 1, 256) = 1
return i_x < -128 || i_x >= 1  → mask i_x >= 1

Result: query at position 256 sees KV positions [0, 256) (the entire prefix).
With window, the effective range is [256-128, 256) = [128, 256). Correct.

Additional Context

The same file also has this at line 216:

if (seqlen_q == 1 && !alibi_slopes_.has_value()) { is_causal = false; }

This disables causal masking for single-query decode when there's no alibi! This
means even without the sliding window bug, generation with flash_attn on MI300X
disables the causal mask entirely during decode. The model sees ALL KV positions
equally, including FUTURE tokens (positions > query position aren't in the cache
yet, but the last N tokens are).

Wait — this depends on the FlashAttention API. In the varlen variant, cu_seqlens
prevent this from triggering. But in the batched variant (flash_attn_func), this
would disable causal during decode. This is a separate issue.

Actually, re-reading the issue, the user runs model.generate() which internally
switches between prefill and decode. HuggingFace uses flash_attn_varlen_func
for prefill and a different path for decode. The decode path may or may not hit
the is_causal = false branch depending on the transform version.

In summary:

1. Primary cause: "b:" vs "t:" alignment mismatch in the CK mask wrapper.
   This affects ALL decoding steps for any model with causal + window_size.
2. Secondary issue: line 216 disables causal for single-query decode in the non-varlen path.