flagos-ai / flagos-ai/FlagTree
[BUG][TLE][PPU] PPU0010 async AIU load rejects i32 operands and lacks B32 lowering
- Dominant language
- Python
- Stars
- 350
- Forks
- 149
- Avg merge
- 2d 4h
- Merged PRs (30d)
- 81
Description
## Summary
On PPU0010, an `int32` block-pointer load marked with `tle.load(..., is_async=True)` is promoted to an AIU load, but the current V1 lowering only accepts 8-bit and 16-bit element widths.
On current `main` after [#1026](https://github.com/flagos-ai/FlagTree/pull/1026), compilation fails at:
```cpp
assert((elementSizeInBytes == 2 || elementSizeInBytes == 1) &&
"AIU load only supports b16 and b8 element type on PPU0010");
```
Before this assertion was added, the same `int32` load was incorrectly lowered to a `.b16` AIU instruction, causing device-side data corruption.
This appears to be a missing FlagTree B32 lowering path rather than a PPU0010 hardware limitation. ACTLIZE provides a PPU0010 `.b32` AIU load implementation for both transposed and non-transposed layouts.
## Environment
| Component | Value |
|---|---|
| Device | T-Head ZW810E / PPU0010 |
| FlagTree | `main`, commit `214a92baaec674eed39efd99c90cf6ae304f61ad` |
| Triton API | 3.6 |
| PyTorch | 2.10 |
| TLE API | `triton.experimental.tle.language` |
## Minimal reproducer
Save the following as `repro_int32_aiu_load.py` and run:
```bash
python repro_int32_aiu_load.py
```
```python
import torch
import triton
import triton.language as tl
from triton.experimental.tle import language as tle
@triton.jit
def int32_aiu_bitwise_probe(packed_ptr, output_ptr):
packed_block = tl.make_block_ptr(
base=packed_ptr,
shape=(16, 16),
strides=(16, 1),
offsets=(0, 0),
block_shape=(16, 16),
order=(1, 0),
)
packed = tle.load(packed_block, is_async=True)
# Simulate extracting one INT4 value from an INT32-packed weight.
unpacked = ((packed >> 4) & 0xF).to(tl.bfloat16)
offsets_m = tl.arange(0, 16)
offsets_n = tl.arange(0, 16)
tl.store(
output_ptr + offsets_m[:, None] * 16 + offsets_n[None, :],
unpacked,
)
torch.manual_seed(0)
packed = torch.randint(
-(1 << 31),
1 << 31,
(16, 16),
device="cuda",
dtype=torch.int32,
)
output = torch.empty((16, 16), device="cuda", dtype=torch.bfloat16)
int32_aiu_bitwise_probe[(1,)](
packed,
output,
num_warps=4,
num_stages=2,
)
torch.cuda.synchronize()
reference = (
((packed.cpu().to(torch.int64) >> 4) & 0xF)
.to(torch.bfloat16)
)
torch.testing.assert_close(
output.cpu(),
reference,
rtol=0,
atol=0,
)
```
## Actual behavior
### Current `main`
The async load is promoted to an AIU load, but PPU0010 V1 lowering rejects the four-byte element width:
```text
Assertion failed:
AIU load only supports b16 and b8 element type on PPU0010
```
The promotion pass currently promotes every asynchronous tensor-pointer load without checking the element width:
- [`TlePromoteAsyncLoadToAIU.cpp`](https://github.com/flagos-ai/FlagTree/blob/214a92baaec674eed39efd99c90cf6ae304f61ad/third_party/ppu/lib/TritonPPUGPUTransforms/TlePromoteAsyncLoadToAIU.cpp#L47-L65)
The failure is introduced later by the V1 LLVM lowering:
- [`LoadStoreOpToLLVM.cpp`](https://github.com/flagos-ai/FlagTree/blob/214a92baaec674eed39efd99c90cf6ae304f61ad/third_party/ppu/lib/TritonPPUGPUToLLVM/LoadStoreOpToLLVM.cpp#L637-L704)
### Behavior before the assertion
With the earlier V1 lowering path, the same four-byte operand was accepted but incorrectly emitted:
```text
ppu.cp.async.aiu.bulk.tensor.shared.global.padz.swzl.zfill.2d.b16
```
The device-side validation then reported:
```text
Mismatched elements: 207 / 256 (80.9%)
Greatest absolute difference: 15
```
Therefore, simply accepting `int32` while selecting `.b16` is not valid.
## Expected behavior
A four-byte PPU0010 asynchronous AIU load should generate a B32 instruction:
```text
ppu.cp.async.aiu.bulk.tensor.shared.global.padz.swzl.zfill.2d.b32
```
The bitwise extraction in the reproducer should then be bit-exact against the CPU reference.
## Evidence that PPU0010 supports B32 AIU load
ACTLIZE 1.0.0 contains explicit PPU0010 implementations for 32-bit AIU copies:
- [PPU0010 non-transposed B32 AIU load](https://github.com/t-head/actlize/blob/129651181ed29ec3d3e61df786f610926b739a9b/include/cute/arch/copy_ppu0010_aiu.hpp#L123-L154)
- [PPU0010 transposed B32 AIU load](https://github.com/t-head/actlize/blob/129651181ed29ec3d3e61df786f610926b739a9b/include/cute/arch/copy_ppu0010_aiu.hpp#L156-L189)
These implementations select 32-bit operands with:
```cpp
sizeof_bits::value == 32
```
and emit:
```text
ppu.cp.async.aiu.bulk.tensor.shared.global.padz.swzl.2d.b32
ppu.cp.async.aiu.bulk.tensor.shared.global.padz.linear.2d.b32
```
ACTLIZE also implements logical 4-bit PPU0010 AIU loads using a B8 transport:
- [PPU0010 4-bit AIU load](https://github.com/t-head/actlize/blob/129651181ed29ec3d3e61df786f610926b739a9b/include/cute/arch/copy_ppu0010_aiu.hpp#L259-L291)
The official PPU1.0 mixed-input example uses BF16 activation and INT4 weights:
- [Mixed-input GEMM documentation](https://github.com/t-head/actlize/blob/129651181ed29ec3d3e61df786f610926b739a9b/examples/16_ppu_mixed_dtype_gemm/README.md#L1-L33)
- [BF16 × INT4 PPU0010 configuration](https://github.com/t-head/actlize/blob/129651181ed29ec3d3e61df786f610926b739a9b/examples/16_ppu_mixed_dtype_gemm/16_ppu_mixed_dtype_gemm.cu#L84-L130)
## Relation to PRs #976 and #1026
[#976](https://github.com/flagos-ai/FlagTree/pull/976) introduced the initial PPU0010 V1 B8 AIU support and made part of the V1 swizzle addressing element-width-aware.
[#1026](https://github.com/flagos-ai/FlagTree/pull/1026) superseded #976 and was merged into `main`. It added the explicit B8/B16 restriction and removed FP32 from the AIU promotion dtype tests.
The B8 work in these PRs appears to provide most of the width-aware structure needed for a B32 extension, but the instruction selection and four-byte correctness path are still missing.
## Why this matters
GPTQ/Marlin-style W4A16 kernels commonly store eight INT4 values in each `int32` container. A kernel may need to:
1. asynchronously move the packed `int32` tile;
2. extract the INT4 nibbles in registers;
3. apply scales;
4. execute BF16 MMA.
The packed INT32 value is only a storage representation. This request does not require an INT32 MMA instruction.
Without B32 AIU load support, such kernels must either use ordinary loads or reinterpret the storage as B8 and implement a separate byte-oriented unpacking path.
## Suggested direction
- Extend PPU0010 V1 AIU instruction selection to cover:
```text
1-byte element -> .b8
2-byte element -> .b16
4-byte element -> .b32
```
- Verify that the generalized V1 swizzle pointer calculation correctly handles `elemBytes == 4`.
- Add a compile-time test requiring `.2d.b32` and rejecting `.b8/.b16` for an `i32` source.
- Add device correctness tests for:
- plain INT32 AIU copy;
- INT32 AIU load followed by shift-and-mask;
- both supported layout orders.
- Preserve the existing B8 and B16 regression tests.
- If other element widths remain unsupported, reject them during AIU promotion or with a structured compiler diagnostic rather than promoting them and failing in a late assertion.
Contributor guide
Research direction
Start with TlePromoteAsyncLoadToAIU.cpp and the PPU0010 V1 lowering in LoadStoreOpToLLVM.cpp, then run repro_int32_aiu_load.py on the stated environment. Trace the element-width selection and swizzle addressing for both layout orders. Done means i32 async loads emit .2d.b32, the bitwise probe matches the CPU reference, unsupported widths fail clearly, and existing B8/B16 behavior remains intact.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python, pytorch
- Domain
- compilers, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 42/100