diff --git a/README.md b/README.md index 6dfeebde8..5f1757d8c 100644 --- a/README.md +++ b/README.md @@ -5,10 +5,6 @@ SKaiNET logo -### Vision - -SKaiNET aims to democratize "Edge AI / On-device AI" by bridging the gap between high-level application development and low-level hardware optimization. We believe AI should be portable, type-safe, and developer-friendly, enabling seamless intelligence in everything from mobile apps to IoT devices without sacrificing performance. - > For architecture details see [ARCHITECTURE.md](ARCHITECTURE.md). --- @@ -89,7 +85,6 @@ SKaiNET is a modular ecosystem. While this repository contains the core engine, |---|---| | Examples and sample projects | [SKaiNET-examples](https://github.com/SKaiNET-developers/SKaiNET-examples) | | Interactive notebooks | [SKaiNET-notebook](https://github.com/SKaiNET-developers/SKaiNET-notebook) | -| LLM inference (Llama, Gemma, Qwen) | [SKaiNET-transformers](https://github.com/SKaiNET-developers/SKaiNET-transformers) | --- @@ -106,11 +101,6 @@ SKaiNET is a modular ecosystem. While this repository contains the core engine, - **SDPA & Gather**: High-performance Scaled Dot-Product Attention and indexing operations. - **TurboQuant**: Runtime KV-cache compression (~8x at 4-bit) for long-context LLM inference. Presets: `safe-lowbit`, `balanced`, `experimental-max`. See `TurboQuantUsage` for integration guide. -### Agentic AI Infrastructure - -- **ComputeGraph**: Unified framework for defining agentic workflows and tool-calling loops. -- Java facade: `JavaAgentLoop` (in `skainet-lang-java`) - ### Neural Network DSL - **Sequential**: `nn { input(); dense(); relu(); dense() }` @@ -125,10 +115,6 @@ SKaiNET is a modular ecosystem. While this repository contains the core engine, - Formats: GGUF, ONNX, SafeTensors, JSON, Image (JPEG, PNG) - Type-safe transform DSL: resize, crop, normalize, toTensor -### Java 21+ Support - -- `SKaiNET` entry point, `TensorJavaOps`, builder-pattern model definition -- Maven BOM (`sk.ainet:skainet-bom`) for one-line version management ### Edge AI: Arduino / C99 Export diff --git a/TURBOQUANT_METAL.md b/TURBOQUANT_METAL.md deleted file mode 100644 index ff0f5e707..000000000 --- a/TURBOQUANT_METAL.md +++ /dev/null @@ -1,325 +0,0 @@ -# TurboQuant Metal Backend — Implementation Task - -> Covers TQ-023 (Metal/Apple Silicon backend) and TQ-024 (Fused dequant+attention kernels) -> Status: TODO — requires Metal Shading Language + Kotlin/Native interop - ---- - -## Objective - -Implement TurboQuant KV-cache compression and decompression as Metal -compute shaders for Apple Silicon, enabling zero-copy unified-memory -KV cache and fused dequant+attention execution. - -## Why Metal - -- Apple Silicon unified memory eliminates CPU↔GPU copies for KV cache -- Metal Performance Shaders (MPS) provides optimized SDPA primitives -- Most on-device inference for SKaiNET targets macOS/iOS (Apple Silicon) -- TurboQuant decode is embarrassingly parallel — ideal for GPU compute - -## Prerequisites - -All prerequisites are complete: -- [x] TurboQuant encoding types (`TensorEncoding.TurboQuantPolar`, `TurboQuantPolarQjl`) -- [x] CPU reference kernels (rotation, quantize, bit-pack, QJL, codec) -- [x] `KvCacheStore` interface with `TurboQuantKvCacheStore` -- [x] `CompressedKvAttention` bridge with `RAW_STORAGE` extension point -- [x] `Placement` model with `DeviceKind.GPU`, `MemoryDomain.UNIFIED` -- [x] `BufferHandle.DeviceResident` for backend-managed buffers - -## Scope - -### In scope -- Metal compute shaders for TurboQuant encode/decode -- Fused dequant+SDPA Metal kernel -- Unified-memory KV cache (no CPU↔GPU copy) -- Kotlin/Native Metal interop for macOS/iOS targets -- Integration with existing `TensorOps.scaledDotProductAttention()` - -### Out of scope -- General-purpose Metal backend for all TensorOps (separate effort) -- CUDA/Vulkan backends -- Training support (inference only) - ---- - -## Architecture - -### Module structure - -``` -skainet-backends/ - skainet-backend-metal/ # New module - build.gradle.kts # KMP config: macosArm64, iosArm64 - src/ - commonMain/kotlin/sk/ainet/exec/metal/ - MetalTurboQuantOps.kt # Public API - MetalKvCacheStore.kt # Metal-backed KvCacheStore - MetalBufferPool.kt # MTLBuffer lifecycle management - nativeMain/kotlin/sk/ainet/exec/metal/ - MetalDevice.kt # MTLDevice + command queue wrapper - MetalShaderLibrary.kt # Compile & cache .metal shaders - MetalBufferHandle.kt # BufferHandle.DeviceResident for Metal - nativeMain/resources/ - turboquant.metal # Metal compute shaders - nativeTest/ - MetalTurboQuantOpsTest.kt # Correctness vs CPU reference -``` - -### Key interfaces to implement - -```kotlin -// MetalKvCacheStore: KvCacheStore backed by MTLBuffer in unified memory -class MetalKvCacheStore( - config: KvCacheConfig, - keyConfig: TurboQuantConfig, - valueConfig: TurboQuantConfig, - device: MetalDevice -) : KvCacheStore { - // KV data lives in MTLBuffer (unified memory) - // appendToken: GPU-side TurboQuant encode - // readKeys/readValues: GPU-side decode or zero-copy raw access -} - -// MetalTurboQuantOps: dispatch TurboQuant kernels to Metal GPU -class MetalTurboQuantOps(device: MetalDevice) { - fun encode(input: MTLBuffer, config: TurboQuantConfig): MTLBuffer - fun decode(encoded: MTLBuffer, config: TurboQuantConfig): MTLBuffer - fun fusedDequantAttention( - query: MTLBuffer, keyCache: MTLBuffer, valueCache: MTLBuffer, - config: TurboQuantConfig, scale: Float - ): MTLBuffer -} -``` - -### Integration with CompressedKvAttention - -The `RAW_STORAGE` dequant strategy in `CompressedKvAttention` is the -extension point. The Metal backend: -1. Returns raw `TensorStorage` with `BufferHandle.DeviceResident` pointing to MTLBuffer -2. The Metal SDPA kernel reads compressed K/V directly and fuses dequant - -```kotlin -// In MetalAttentionOps (extends or replaces scaledDotProductAttention) -override fun scaledDotProductAttention(query, key, value, mask, scale, causal): Tensor { - val keyStorage = compressedKv.loadKeyStorageRaw(layer) - if (keyStorage.buffer is BufferHandle.DeviceResident) { - // Dispatch fused Metal kernel - return metalOps.fusedDequantAttention(query, keyStorage, valueStorage, ...) - } - // Fallback to CPU - return super.scaledDotProductAttention(query, key, value, mask, scale, causal) -} -``` - ---- - -## Metal Shaders - -### File: `turboquant.metal` - -```metal -// Required compute kernels: - -// 1. turboquant_encode -// Per-thread: rotate → quantize → pack one head's vector -// Threadgroup: shared memory for Walsh-Hadamard butterfly -kernel void turboquant_encode( - device const float* input [[buffer(0)]], // [numHeads, headDim] - device uchar* packed_output [[buffer(1)]], // packed codes - device half* scales_output [[buffer(2)]], // per-group scales - constant TQParams& params [[buffer(3)]], // bits, headDim, seed - uint tid [[thread_position_in_grid]] -); - -// 2. turboquant_decode -// Per-thread: unpack → dequantize → inverse rotate one head's vector -kernel void turboquant_decode( - device const uchar* packed_input [[buffer(0)]], - device const half* scales_input [[buffer(1)]], - device float* output [[buffer(2)]], - constant TQParams& params [[buffer(3)]], - uint tid [[thread_position_in_grid]] -); - -// 3. turboquant_fused_sdpa (highest value kernel) -// Fuses: KV dequant + Q@K^T scaling + softmax + @V -// Avoids materializing decompressed K/V in global memory -kernel void turboquant_fused_sdpa( - device const float* query [[buffer(0)]], // [nHeads, seqLen, headDim] - device const uchar* key_packed [[buffer(1)]], // compressed keys - device const half* key_scales [[buffer(2)]], - device const uchar* value_packed [[buffer(3)]], // compressed values - device const half* value_scales [[buffer(4)]], - device float* output [[buffer(5)]], // [nHeads, seqLen, headDim] - constant SDPAParams& params [[buffer(6)]], - uint2 tid [[thread_position_in_grid]], - uint2 tgid [[threadgroup_position_in_grid]] -); - -// 4. walsh_hadamard_transform -// Threadgroup-cooperative WHT for rotation stage -// Uses threadgroup memory for butterfly communication -kernel void walsh_hadamard_transform( - device float* data [[buffer(0)]], - constant uint& log2_n [[buffer(1)]], - uint tid [[thread_position_in_threadgroup]], - uint tg_size [[threads_per_threadgroup]], - threadgroup float* shared [[threadgroup(0)]] -); -``` - -### Shader parameters - -```metal -struct TQParams { - uint bits; // 2, 3, 4, or 8 - uint headDim; // dimension per head - uint numHeads; // heads in this batch - uint seed; // rotation seed - uint groupSize; // quantization group size (32) - bool useQjl; // whether QJL residual is present - uint residualBits; // QJL residual bits (1-4) -}; - -struct SDPAParams { - uint nHeads; - uint nKVHeads; - uint seqLen; - uint kvLen; - uint headDim; - float scale; // 1/sqrt(headDim) - uint keyBits; - uint valueBits; - bool causal; -}; -``` - ---- - -## Implementation Plan - -### Phase 1: Metal infrastructure (no TurboQuant yet) - -| Task | Description | Files | -|---|---|---| -| M-001 | Create `skainet-backend-metal` module | `build.gradle.kts`, `settings.gradle.kts` | -| M-002 | `MetalDevice` wrapper (MTLDevice, command queue) | `MetalDevice.kt` | -| M-003 | `MetalShaderLibrary` (compile .metal, cache pipelines) | `MetalShaderLibrary.kt` | -| M-004 | `MetalBufferHandle` → `BufferHandle.DeviceResident` | `MetalBufferHandle.kt` | -| M-005 | `MetalBufferPool` (reusable MTLBuffer pool) | `MetalBufferPool.kt` | -| M-006 | Kotlin/Native cinterop for Metal.framework | `metal.def`, build config | - -### Phase 2: TurboQuant encode/decode shaders - -| Task | Description | Files | -|---|---|---| -| M-010 | `turboquant_encode` shader | `turboquant.metal` | -| M-011 | `turboquant_decode` shader | `turboquant.metal` | -| M-012 | `walsh_hadamard_transform` cooperative shader | `turboquant.metal` | -| M-013 | `MetalTurboQuantOps` Kotlin dispatch | `MetalTurboQuantOps.kt` | -| M-014 | Correctness tests vs CPU reference | `MetalTurboQuantOpsTest.kt` | - -### Phase 3: Metal KV cache store - -| Task | Description | Files | -|---|---|---| -| M-020 | `MetalKvCacheStore` with unified-memory buffers | `MetalKvCacheStore.kt` | -| M-021 | GPU-side append (encode on GPU, no CPU round-trip) | shader + Kotlin | -| M-022 | GPU-side read (decode on GPU for raw access) | shader + Kotlin | -| M-023 | Integration with `CompressedKvAttention.RAW_STORAGE` | bridge code | - -### Phase 4: Fused dequant+SDPA - -| Task | Description | Files | -|---|---|---| -| M-030 | `turboquant_fused_sdpa` shader | `turboquant.metal` | -| M-031 | Tiled attention with on-the-fly dequant | shader optimization | -| M-032 | Causal mask support in fused kernel | shader | -| M-033 | GQA (grouped-query attention) support | shader | -| M-034 | End-to-end benchmark vs CPU decode+SDPA | benchmark suite | - -### Phase 5: Integration & optimization - -| Task | Description | Files | -|---|---|---| -| M-040 | Wire Metal backend into `PlatformCpuOpsFactory` for macOS/iOS | factory impl | -| M-041 | Fallback to CPU when Metal unavailable | graceful degradation | -| M-042 | Unified-memory placement resolution in `MemoryPlanner` | planner update | -| M-043 | `@KvCache(device = GPU)` annotation handling | annotation processor | -| M-044 | Performance tuning: threadgroup sizes, occupancy | shader tuning | - ---- - -## Kotlin/Native Metal Interop - -### cinterop definition (`metal.def`) - -``` -language = Objective-C -headers = Metal/Metal.h MetalPerformanceShaders/MetalPerformanceShaders.h -compilerOpts = -framework Metal -framework MetalPerformanceShaders -linkerOpts = -framework Metal -framework MetalPerformanceShaders -framework Foundation -``` - -### Key ObjC types to bridge - -| Metal Type | Kotlin Usage | -|---|---| -| `MTLDevice` | GPU device handle | -| `MTLCommandQueue` | Serial command submission | -| `MTLCommandBuffer` | Batch of GPU commands | -| `MTLComputeCommandEncoder` | Dispatch compute kernels | -| `MTLBuffer` | GPU/unified memory buffer | -| `MTLComputePipelineState` | Compiled shader pipeline | -| `MTLLibrary` | Compiled shader library | - -### Unified memory pattern - -```kotlin -// Allocate in unified memory — accessible from both CPU and GPU -val buffer = device.newBuffer( - length = sizeInBytes, - options = MTLResourceStorageModeShared // unified memory -) - -// CPU can read/write directly (no copy needed) -val ptr = buffer.contents() - -// GPU kernel reads/writes same memory -encoder.setBuffer(buffer, offset = 0, index = 0) -encoder.dispatchThreads(...) -``` - ---- - -## Performance Targets - -| Metric | CPU Reference | Metal Target | -|---|---|---| -| TurboQuant encode (128d, 4-bit) | ~10 μs | < 2 μs | -| TurboQuant decode (128d, 4-bit) | ~8 μs | < 1 μs | -| Fused dequant+SDPA (8 heads, 128d, 1024 KV) | N/A (separate) | < 100 μs | -| KV cache memory (4-bit vs FP32) | 8x compression | 8x compression | -| CPU↔GPU copies for KV cache | N/A | 0 (unified memory) | - -## Acceptance Criteria - -- [ ] Metal shaders compile and run on Apple Silicon (M1+) -- [ ] Encode/decode correctness matches CPU reference within tolerance -- [ ] Fused dequant+SDPA produces correct attention output -- [ ] Zero CPU↔GPU copies for KV cache in unified memory mode -- [ ] Graceful fallback to CPU when Metal is unavailable -- [ ] Benchmark shows meaningful speedup over CPU reference path -- [ ] Works on both macOS (macosArm64) and iOS (iosArm64) - -## References - -- [Metal Shading Language Spec](https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf) -- [Metal Best Practices Guide](https://developer.apple.com/library/archive/documentation/3DDrawing/Conceptual/MTLBestPracticesGuide/) -- [MPSGraph Documentation](https://developer.apple.com/documentation/metalperformanceshadersgraph) -- [TurboQuant paper (arXiv)](https://arxiv.org/html/2504.19874v1) -- SKaiNET existing backend: `skainet-backends/skainet-backend-cpu/` -- SKaiNET CPU SIMD kernels: `JvmQuantizedVectorKernels.kt`, `JvmTurboQuantKernels.kt` -- SKaiNET TurboQuant reference: `skainet-lang/.../ops/turboquant/` diff --git a/build-logic/convention/build.gradle.kts b/build-logic/convention/build.gradle.kts index 589e1036d..8cb5a7a2b 100644 --- a/build-logic/convention/build.gradle.kts +++ b/build-logic/convention/build.gradle.kts @@ -35,5 +35,9 @@ gradlePlugin { id = "sk.ainet.documentation" implementationClass = "DocumentationPlugin" } + register("SKaiNetBomCoverage") { + id = "sk.ainet.transformers.bom-coverage" + implementationClass = "sk.ainet.buildlogic.bom.BomCoveragePlugin" + } } } diff --git a/build-logic/convention/src/main/kotlin/DocumentationPlugin.kt b/build-logic/convention/src/main/kotlin/DocumentationPlugin.kt index 0a17efa15..972f0c06a 100644 --- a/build-logic/convention/src/main/kotlin/DocumentationPlugin.kt +++ b/build-logic/convention/src/main/kotlin/DocumentationPlugin.kt @@ -1,9 +1,6 @@ import org.gradle.api.Plugin import org.gradle.api.Project import org.gradle.api.Action -import org.gradle.kotlin.dsl.named -import org.gradle.kotlin.dsl.register -import org.gradle.kotlin.dsl.configureEach class DocumentationPlugin : Plugin { override fun apply(project: Project) { diff --git a/build-logic/convention/src/main/kotlin/sk/ainet/buildlogic/bom/BomCoverageExtension.kt b/build-logic/convention/src/main/kotlin/sk/ainet/buildlogic/bom/BomCoverageExtension.kt new file mode 100644 index 000000000..5921383e1 --- /dev/null +++ b/build-logic/convention/src/main/kotlin/sk/ainet/buildlogic/bom/BomCoverageExtension.kt @@ -0,0 +1,7 @@ +package sk.ainet.buildlogic.bom + +import org.gradle.api.provider.SetProperty + +abstract class BomCoverageExtension { + abstract val excludePublished: SetProperty +} diff --git a/build-logic/convention/src/main/kotlin/sk/ainet/buildlogic/bom/BomCoveragePlugin.kt b/build-logic/convention/src/main/kotlin/sk/ainet/buildlogic/bom/BomCoveragePlugin.kt new file mode 100644 index 000000000..97e8f0c9d --- /dev/null +++ b/build-logic/convention/src/main/kotlin/sk/ainet/buildlogic/bom/BomCoveragePlugin.kt @@ -0,0 +1,45 @@ +package sk.ainet.buildlogic.bom + +import org.gradle.api.GradleException +import org.gradle.api.Plugin +import org.gradle.api.Project + +private const val PUBLISH_PLUGIN_ID = "com.vanniktech.maven.publish" + +class BomCoveragePlugin : Plugin { + override fun apply(project: Project) { + require(project.plugins.hasPlugin("java-platform")) { + "sk.ainet.transformers.bom-coverage requires the java-platform plugin to be applied first" + } + + val ext = project.extensions.create( + "bomCoverage", + BomCoverageExtension::class.java, + ) + ext.excludePublished.convention(emptySet()) + + project.rootProject.subprojects + .filter { it.path != project.path } + .forEach { project.evaluationDependsOn(it.path) } + + project.afterEvaluate { + val excluded = ext.excludePublished.get() + project.path + val publishedPaths = project.rootProject.subprojects + .filter { it.plugins.hasPlugin(PUBLISH_PLUGIN_ID) } + .map { it.path } + .filterNot { it in excluded } + .sorted() + + if (publishedPaths.isEmpty()) { + throw GradleException( + "[bom-coverage] No published subprojects found for ${project.path}. " + + "At least one sibling must apply '$PUBLISH_PLUGIN_ID'." + ) + } + + project.dependencies.constraints { + publishedPaths.forEach { add("api", project.project(it)) } + } + } + } +} diff --git a/docs/whisper-iree-issues/ISSUE-A-conv1d-infer-outputs.md b/docs/whisper-iree-issues/ISSUE-A-conv1d-infer-outputs.md deleted file mode 100644 index 4e13807a0..000000000 --- a/docs/whisper-iree-issues/ISSUE-A-conv1d-infer-outputs.md +++ /dev/null @@ -1,72 +0,0 @@ -# Conv1d/2d/3dOperation.inferOutputs() echoes input shape instead of computing output shape - -## Problem - -`Conv1dOperation.inferOutputs()` in `TensorOperations.kt` (line ~439) returns the -input tensor's shape as the output shape, ignoring weight shape, stride, padding, -and dilation: - -```kotlin -override fun inferOutputs(inputs: List): List { - require(inputs.size >= 2) { "Conv1d operation requires at least 2 inputs" } - val outputShape = inputs[0].shape // <-- BUG: just copies input shape - return listOf(TensorSpec("conv1d_output", outputShape, inputs[0].dtype, ...)) -} -``` - -Conv2dOperation (line ~471) and Conv3dOperation (line ~503) have the identical bug. - -## Expected - -```kotlin -override fun inferOutputs(inputs: List): List { - val inShape = inputs[0].shape // [N, Cin, L] - val wShape = inputs[1].shape // [Cout, Cin/g, K] - val stride = (parameters["stride"] as? Int) ?: 1 - val padding = (parameters["padding"] as? Int) ?: 0 - val dilation = (parameters["dilation"] as? Int) ?: 1 - val outShape = if (inShape != null && wShape != null && inShape.size == 3 && wShape.size == 3) - listOf(inShape[0], wShape[0], - (inShape[2] + 2*padding - dilation*(wShape[2]-1) - 1)/stride + 1) - else null - return listOf(TensorSpec("conv1d_output", outShape, inputs[0].dtype, ...)) -} -``` - -The formula already exists in `VoidTensorOps.calculateConv1dShape()` (line ~747). -`ConvShapeUtils` was added to the JAR but `inferOutputs()` does not call it yet. - -## Impact - -When the StableHLO converter calls `inferOutputs()` to determine the MLIR output -type, it gets the wrong shape. For Whisper's first conv1d: - -``` -Input: [1, 80, 3000] Weight: [384, 80, 3] stride=1 padding=1 -Actual: [1, 80, 3000] ← wrong (echoed input) -Expect: [1, 384, 3000] ← correct -``` - -This produces `tensor` in the MLIR (12 occurrences), which `iree-compile` -rejects. - -## Parameters are available - -PR #532 stores stride/padding/dilation in `operation.parameters`: - -```kotlin -// RecordingExecution.kt:238-261 -val params = mapOf("stride" to stride, "padding" to padding, "dilation" to dilation, "groups" to groups) -record(Conv1dOperation(params), ...) -``` - -Verified by test: `assertEquals(1, recorded.operation.parameters["stride"])` - -## Suggested fix - -Extract `ConvShapeUtils` calls into all three `inferOutputs()` methods. -Single PR covering conv1d/2d/3d since the bug and fix are identical. - -## Test - -See `Conv1dTapeToHloTest.kt` — asserts `tensor scaledDotProductAttention(...) = - base.scaledDotProductAttention(query, key, value, mask, scale, causal) -``` - -Fix (same pattern as conv1d in PR #532): -```kotlin -override fun scaledDotProductAttention( - query, key, value, mask, scale, causal -): Tensor { - val out = base.scaledDotProductAttention(query, key, value, mask, scale, causal) - val params = mapOf("scale" to scale, "causal" to causal) - record(ScaledDotProductAttentionOperation(params), - listOfNotNull(query, key, value, mask), listOf(out)) - return out -} -``` - -### 2. TensorOperations: add ScaledDotProductAttentionOperation - -**File:** `skainet-lang-core/.../tensor/ops/TensorOperations.kt` - -```kotlin -class ScaledDotProductAttentionOperation( - parameters: Map = emptyMap() -) : BaseOperation("scaledDotProductAttention", "nn", parameters) { - override fun inferOutputs(inputs: List): List { - // Output shape = query shape: [batch, nHeads, seqLen, headDim] - return listOf(TensorSpec("sdpa_output", inputs[0].shape, inputs[0].dtype)) - } -} -``` - -### 3. StableHLO converter: decompose SDPA - -**File:** `skainet-compile-hlo/.../converters/NeuralNetOperationsConverter.kt` - -Register "scaledDotProductAttention" and decompose into: -```mlir -// scores = Q @ K.T (batched matmul with K transposed) -%scores = stablehlo.dot_general %query, %key, - batching_dims = [0, 1] x [0, 1], - contracting_dims = [3] x [3] - : (tensor, tensor) -> tensor - -// scale -%scaled = stablehlo.multiply %scores, %scale_splat - -// optional mask (additive) -%masked = stablehlo.add %scaled, %mask // if mask != null - -// softmax over last dim -%weights = stablehlo softmax ... - -// output = weights @ V (batched matmul) -%output = stablehlo.dot_general %weights, %value, - batching_dims = [0, 1] x [0, 1], - contracting_dims = [3] x [2] -``` - -Note: `contracting_dims = [3] x [3]` for Q@K.T because we contract -headDim of Q (last dim) with headDim of K (also last dim). This is -different from standard matmul where you contract last of A with -second-to-last of B — here K is NOT pre-transposed. - -## Test - -```kotlin -val ctx = DefaultGraphExecutionContext.tape(baseOps = VoidTensorOps()) -val q = ctx.fromFloatArray(Shape(1, 6, 4, 64), ...) // [batch, heads, seq, headDim] -val k = ctx.fromFloatArray(Shape(1, 6, 4, 64), ...) -val v = ctx.fromFloatArray(Shape(1, 6, 4, 64), ...) - -val (tape, out) = ctx.record { - ctx.ops.scaledDotProductAttention(q, k, v) -} - -val graph = tape!!.toComputeGraph(synthesizeExternalInputs = true) -val module = StableHloConverterFactory.createExtended().convert(graph, "test_sdpa") - -// Should contain dot_general for Q@K.T and weights@V -assertTrue(module.content.contains("stablehlo.dot_general")) -assertFalse(module.content.contains("dense<0.0>")) // no zero constants -``` diff --git a/docs/whisper-iree-issues/README.md b/docs/whisper-iree-issues/README.md deleted file mode 100644 index d10af1a42..000000000 --- a/docs/whisper-iree-issues/README.md +++ /dev/null @@ -1,37 +0,0 @@ -# SKaiNET Upstream Issues — Whisper IREE Pipeline - -Two issues block the native SKaiNET DSL → StableHLO → IREE compilation path. - -## Issue A: Conv1dOperation.inferOutputs echoes input shape - -`Conv1dOperation.inferOutputs()` returns `inputs[0].shape` instead of -computing `[batch, outChannels, outLength]`. Same bug in Conv2d/Conv3d. - -**File:** `skainet-lang/skainet-lang-core/.../tensor/ops/TensorOperations.kt` -**Fix:** Use `ConvShapeUtils` (already in JAR) from `inferOutputs()`. - -## Issue B: toComputeGraph loses edge wiring and op types - -`tape.toComputeGraph()` produces nodes where: -- Binary ops (add, matmul, subtract, ...) have wrong input edge count -- Some ops have `operation.type = "trace"` instead of recognized names - -157 of 296 Whisper encoder nodes emit "Unsupported ... arity" in MLIR. - -**File:** `skainet-compile/skainet-compile-dag/.../tape/extensions.kt` or -`DefaultExecutionTape.toComputeGraph()` - -## Test - -`Conv1dTapeToHloTest.kt` is a KMP commonTest that: -1. Builds a tape-recording context -2. Runs conv1d → gelu → add through `ctx.ops` -3. Converts tape to ComputeGraph -4. Exports to StableHLO MLIR -5. Asserts: no `tensor