flagos-ai / flagos-ai/FlagTree

[MetaX] ICE in ConvertTritonGPUToLLVM when BLOCK_SIZE_N=1: PassManager::run failed

Open
#1,179 0 comments 0 reactions 1 assignee Claimed by @zhzhcookie View on GitHub
AABS bug flagos2.2-rc2 metax P0
Dominant language
Python
Stars
350
Forks
149
Avg merge
2d 4h
Merged PRs (30d)
81

Description

## Summary

On the MetaX (mcTriton) backend, a valid Triton GEMM kernel with `BLOCK_SIZE_N = 1` fails to compile with an internal compiler error during the `ConvertTritonGPUToLLVM` MLIR pass. The same kernel compiles and runs correctly with `BLOCK_SIZE_N >= 16`. This looks like an ICE-on-valid in the MetaX Triton backend, not a problem with the kernel source.

## Environment

| Item | Value |
| --- | --- |
| FlagTree / Triton backend | MetaX (mcTriton) |
| Triton path | `/opt/flagtree/triton` |
| Backend compiler | `/opt/flagtree/triton/backends/metax/compiler.py` (line 368, `make_mlir`) |
| Device | MetaX C550 |
| Host driver | 3.8.30 |
| Python | 3.12 |
| Container image | `harbor.baai.ac.cn/flagos-app/vllm0.20.2-metax-maca3.7.2.1:2.1.2-0.2.2rc2.post1_gdb28502.d20260915` |

## Minimal reproduction

```python
import torch
import triton
import triton.language as tl

@triton.jit
def gemm_bn1_kernel(A, B, C, M, N, K,
BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
pid_m = tl.program_id(0)
pid_n = tl.program_id(1)

rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
rk = tl.arange(0, BLOCK_K)

acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
for k in range(0, K, BLOCK_K):
A_ptrs = A + rm[:, None] * K + (k + rk)[None, :]
B_ptrs = B + (k + rk)[:, None] * N + rn[None, :]
a = tl.load(A_ptrs, mask=rm[:, None] < M, other=0.0)
b = tl.load(B_ptrs, mask=rn[None, :] < N, other=0.0)
acc += tl.dot(a, b)

C_ptrs = C + rm[:, None] * N + rn[None, :]
tl.store(C_ptrs, acc.to(tl.bfloat16), mask=rm[:, None] < M)

def run(BM, BN, BK, num_warps):
M, N, K = 128, 64, 32
a = torch.randn((M, K), device='cuda', dtype=torch.bfloat16)
b = torch.randn((K, N), device='cuda', dtype=torch.bfloat16)
c = torch.empty((M, N), device='cuda', dtype=torch.bfloat16)
grid = (triton.cdiv(M, BM), triton.cdiv(N, BN))
print(f"--- BM={BM}, BN={BN}, BK={BK}, warps={num_warps} ---")
try:
gemm_bn1_kernel[grid](a, b, c, M, N, K,
BLOCK_M=BM, BLOCK_N=BN, BLOCK_K=BK,
num_warps=num_warps)
torch.cuda.synchronize()
print(" OK")
except Exception as e:
print(f" FAIL: {type(e).__name__}: {str(e)[:200]}")

if __name__ == "__main__":
run(128, 64, 32, 4) # expected: OK
run(128, 1, 32, 4) # expected: OK, actual: FAIL
run(128, 1, 64, 8) # expected: OK, actual: FAIL
run(64, 1, 32, 4) # expected: OK, actual: FAIL
run(32, 1, 32, 2) # expected: OK, actual: FAIL
```

## Expected behavior

All configurations compile and run. `BLOCK_SIZE_N=1` is a legal Triton configuration, even if inefficient.

## Actual behavior

Configurations with `BLOCK_SIZE_N = 1` fail with:

```text
RuntimeError: PassManager::run failed
error: Failures have been detected while processing an MLIR pass pipeline
note: Pipeline failed while executing [`ConvertTritonGPUToLLVM` on 'builtin.module' operation]: reproducer generated at `std::errs, please share the reproducer above with Triton project.`
```

The failure originates in `/opt/flagtree/triton/backends/metax/compiler.py:368` (`make_mlir`), i.e. inside the MetaX Triton compilation pipeline, not in the kernel source.

## Stack trace (excerpt)

```text
File "/opt/flagtree/triton/runtime/jit.py", line 906, in _do_compile
kernel = self.compile(src, target=target, options=options.__dict__)
File "/opt/flagtree/triton/compiler/compiler.py", line 346, in compile
next_module = compile_ir(module, metadata)
File "/opt/flagtree/triton/backends/metax/compiler.py", line 470, in
stages["mlir"] = lambda src, metadata: self.make_mlir(src, metadata, options, capability)
File "/opt/flagtree/triton/backends/metax/compiler.py", line 368, in make_mlir
pm.run(mod, "make_mlir")
RuntimeError: PassManager::run failed
```

## Analysis

- The kernel is valid Triton. `BLOCK_SIZE_N=1` is semantically allowed.
- The failure happens in `ConvertTritonGPUToLLVM`, a compiler pass, not in the frontend.
- The same kernel compiles and runs with `BLOCK_SIZE_N >= 16`, and also on NVIDIA/AMD Triton backends.
- This is therefore an ICE-on-valid in the MetaX Triton backend, not a kernel bug.
- The issue is reachable from real workloads: FlagGems' linear op can emit `BLOCK_SIZE_N=1` for small-N shapes (e.g. MoE expert gate / shared expert projections), which then hits this compiler bug.

## Impact

Blocks running MoE models (e.g. Qwen3.6-35B-A3B) on MetaX via vLLM + FlagGems, because the linear kernel for the expert gate/shared experts is compiled with `BLOCK_SIZE_N=1`.

## Request

Please fix `ConvertTritonGPUToLLVM` in the MetaX backend to handle `BLOCK_SIZE_N=1` correctly, or at minimum emit a clear diagnostic instead of `PassManager::run failed`.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.