Skip to content

feat: fp and cuda-graph support#346

Open
devillove084 wants to merge 9 commits into
NVlabs:mainfrom
devillove084:fp-and-graph-support
Open

feat: fp and cuda-graph support#346
devillove084 wants to merge 9 commits into
NVlabs:mainfrom
devillove084:fp-and-graph-support

Conversation

@devillove084

@devillove084 devillove084 commented Jul 6, 2026

Copy link
Copy Markdown

Summary

Adds FP8 (E4M3/E5M2), FP6 (E3M2/E2M3), and FP4 (E2M1) type support through the full compiler pipeline (types → dialect → MIR translation → PTX lowering), plus a complete CUDA Graph API with stream capture, manual node construction, executable graph update, and an auto-warmup strategy inspired by llama.cpp. This is foundational work for Blackwell (sm_100a+) tensor core operations and low-overhead kernel launch.

Changes

Sub-byte float type system (4 crates)

  • DeviceExternType gains BFloat16, Float8E4M3, Float8E5M2, Float4E2M1, Float6E3M2, Float6E2M3 — all sub-byte float types map to i8/i16 in LLVM IR since no native types exist in pliron/LLVM
  • All match arms updated: module.rs, ops.rs, lower.rs, export_test.rs

FP8/FP6/FP4 tcgen05 device intrinsics + compiler pipeline (5 crates)

  • Tcgen05ElementType extended with E4M3, E5M2, E2M3, E3M2, E2M1 + idesc_value() method for kind-dependent idesc encoding (enum discriminants are identifiers; build() maps to correct PTX ISA bit patterns per kind family)
  • New device intrinsics: tcgen05_mma_ws_e4m3/e5m2/e2m3/e3m2/e2m1(), plus new_e4m3()/new_e5m2()/new_e2m3()/new_e3m2()/new_e2m1() descriptor builders
  • Dialect ops: Tcgen05MmaWsE2M3Op, Tcgen05MmaWsE3M2Op, Tcgen05MmaWsE2M1Op registered in dialect-nvvm
  • MIR translation: emit_tcgen05_mma_ws_e2m3/e3m2/e2m1() in mir-importer + dispatch entries for all f8f6f4 types
  • PTX lowering: all f8f6f4-family types (FP8/FP6/FP4) use kind::f8f6f4; F16/BF16 use kind::f16; TF32 uses kind::tf32 — the element type is encoded in the instruction descriptor (idesc), not the PTX mnemonic

CUDA Graph API (1 crate, 1 new module)

  • Stream capture: CudaStreamCaptureExt trait — begin_capture, end_capture, is_capturing
  • Graph lifecycle: CudaGraph (RAII for CUgraph), CudaGraphExec (instantiate/launch/upload/destroy)
  • Manual node construction: CudaGraphNode, add_kernel_node, add_empty_node, add_dependencies
  • Kernel parameters: KernelNodeParams with aligned push_param<T>() builder
  • Graph update: CudaGraphExec::update() returning GraphUpdateResult, set_kernel_node_params() for hot param updates
  • Auto-warmup (llama.cpp-inspired): CachedGraphExec with GraphStrategy (Direct / AutoCapture / Prebuilt), FNV-hash-based property change detection, two-phase warmup protocol
  • Thread safety: CaptureModeGuard RAII for cuThreadExchangeStreamCaptureMode
  • 19 integration tests covering lifecycle, capture, node construction, update, and multi-launch

Testing

# idesc encoding (cuda-device) — 7 passed
test idesc_value_matches_ptx_isa_encoding ... ok
test build_encodes_atype_btype_via_idesc_value ... ok
test fp8_constructors_use_f8f6f4_encoding ... ok
...

# DeviceExternType (llvm-export) — 13 passed (incl. existing FP8)
test device_extern_fp4_fp6_types_emit_as_i8 ... ok
test device_extern_all_subbyte_float_types_roundtrip ... ok
...

# PTX lowering kind verification (mir-lower) — 10 passed
test test_tcgen05_mma_ws_e4m3_lowers_to_f8f6f4_kind ... ok
test test_tcgen05_mma_ws_e2m1_lowers_to_f8f6f4_kind ... ok
test test_all_f8f6f4_types_share_same_ptx_mnemonic ... ok
test test_scalar_vs_tensor_core_fp4_instruction_count ... ok
...

# CUDA Graph integration (cuda-core, --test-threads=1) — 19 passed
test captured_graph_launches_and_syncs ... ok
test graph_update_succeeds_for_identical_topology ... ok
...

# Full workspace
cargo test --workspace   # all green (1 pre-existing GPU VRAM test excluded)

22 new tests added across 3 test suites. The test_scalar_vs_tensor_core_fp4_instruction_count test demonstrates that a single tcgen05.mma.ws.kind::f8f6f4 instruction replaces ~8.4M scalar instructions for a 128×256×32 FP4 tile (~1,000,000× speedup).

  • cargo test --workspace passes
  • cargo oxide run <example> — N/A (library-level additions, no new examples)

These storage-only types map to i8/i16 in LLVM IR since pliron
lacks dedicated types for them. FP8 and BF16 are used by Blackwell
tensor cores (tcgen05) and device-extern FFI declarations.
Introduces CudaGraph (RAII wrapper for CUgraph), CudaGraphExec
(instantiated executable graph), CaptureMode, CaptureModeGuard,
and CudaStreamCaptureExt for stream-based graph capture.

Graph tests require serial execution due to shared primary
context; run with --test-threads=1.
- Tcgen05ElementType gains E4M3 and E5M2 variants for FP8 tensor cores
- New tcgen05_mma_ws_e4m3 / tcgen05_mma_ws_e5m2 device intrinsics
- CudaGraph: add_kernel_node, add_empty_node, add_dependencies
- CudaGraphNode: RAII wrapper for CUgraphNode
- KernelNodeParams: safe kernel param builder with push_param
- CudaGraphExec: update() and set_kernel_node_params()
- 19 graph integration tests (run with --test-threads=1)
- dialect-nvvm: Tcgen05MmaWsE4M3Op, Tcgen05MmaWsE5M2Op
- mir-importer: emit_tcgen05_mma_ws_e4m3, emit_tcgen05_mma_ws_e5m2
- mir-lower: PTX lowering via convert_mma_ws with kind::f16
  (FP8 element type is encoded in the idesc, not PTX mnemonic)
Inspired by llama.cpp's CUDA graph pattern:
- GraphStrategy: Direct, AutoCapture, Prebuilt
- AutoCapture warmup: first run direct, second triggers capture
- Property-change detection via FNV hash of data ptrs + shapes
- CachedGraphExec::launch_or_capture() single entry point
@devillove084 devillove084 changed the title Fp and graph support feat: fp and cuda-graph support Jul 6, 2026
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