Skip to content

Add mixed-signedness INT8 mma.sync variants (s8*u8, u8*s8)#334

Open
honeyspoon wants to merge 1 commit into
NVlabs:mainfrom
honeyspoon:upstream/add-mma-int8-mixed
Open

Add mixed-signedness INT8 mma.sync variants (s8*u8, u8*s8)#334
honeyspoon wants to merge 1 commit into
NVlabs:mainfrom
honeyspoon:upstream/add-mma-int8-mixed

Conversation

@honeyspoon

Copy link
Copy Markdown
Contributor

Summary

Adds four warp-level MMA intrinsics for mixed-signedness INT8 matrix multiply-accumulate, extending the same-signedness s8 variant merged in #329. Mixed-signedness is common in quantized inference where weights and activations use different signedness.

mma.sync.aligned.m16n8k32.row.col.s32.s8.u8.s32
    {%d0, %d1, %d2, %d3},
    {%a0, %a1, %a2, %a3},
    {%b0, %b1},
    {%c0, %c1, %c2, %c3};

New intrinsics:

  • mma_m16n8k32_s32_s8_u8 — signed A × unsigned B, m16n8k32 shape
  • mma_m16n8k32_s32_u8_s8 — unsigned A × signed B, m16n8k32 shape
  • mma_m16n8k16_s32_s8_u8 — signed A × unsigned B, m16n8k16 shape
  • mma_m16n8k16_s32_u8_s8 — unsigned A × signed B, m16n8k16 shape

All variants require sm_80+ / PTX ISA 7.0 and use convergent inline assembly.

Part of tracking issue #274 (Phase 3).

What changed

  • cuda-device: Added 4 user-facing #[inline(never)] stubs in wmma.rs with full per-lane fragment documentation.
  • dialect-nvvm: Added 4 operations with 10/7 operands and 4 results respectively; satfinite attribute support; verifier tests for all variants including negative cases.
  • mir-importer: Added dispatch table entries and MMA fragment emitter with mixed-type PTX mnemonic construction.
  • mir-lower: Added lowering conversions using convergent inline asm. Mixed-signedness is encoded in the PTX type specifiers (.s8.u8 vs .u8.s8).
  • Pipeline detection for mixed-signedness INT8 MMA instructions.
  • Lowering tests for all 4 variants.

Safety and register contract

m16n8k32 (10-operand): C[s32 x4] + A[b32 x4] * B[b32 x2]  →  D[s32 x4]
m16n8k16  (7-operand): C[s32 x4] + A[b32 x2] * B[b32 x1]  →  D[s32 x4]

All 32 lanes must execute the same instruction with the same qualifiers, with no lane exited.

Testing

  • cargo fmt --check clean
  • cargo test -p dialect-nvvm --all-targets — verifier tests pass (including satfinite variants)
  • cargo test -p mir-lower --all-targets — lowering tests pass
  • cargo oxide run e2e on sm_80+ GPU

Checklist

  • Signed-off-by DCO trailer
  • No crates/cuda-bindings/ changes
  • SPDX headers on new source files

Adds four warp-level MMA intrinsics for mixed-signedness INT8 matrix
multiply-accumulate, where A and B operands have different signedness.
This extends the same-signedness INT8 s8 variant merged in NVlabs#329.

New intrinsics:
- mma_m16n8k32_s32_s8_u8: signed A * unsigned B, m16n8k32
- mma_m16n8k32_s32_u8_s8: unsigned A * signed B, m16n8k32
- mma_m16n8k16_s32_s8_u8: signed A * unsigned B, m16n8k16
- mma_m16n8k16_s32_u8_s8: unsigned A * signed B, m16n8k16

```ptx
mma.sync.aligned.m16n8k32.row.col.s32.s8.u8.s32
    {%d0, %d1, %d2, %d3},
    {%a0, %a1, %a2, %a3},
    {%b0, %b1},
    {%c0, %c1, %c2, %c3};
```

All variants require sm_80+ / PTX ISA 7.0 and use convergent inline
assembly. Implements the full 4-layer pattern with pipeline detection,
verifier tests, and lowering tests including satfinite variant coverage.

Part of tracking issue NVlabs#274 (Phase 3).

Signed-off-by: abder <bobmatt911@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant