feat: naive attention kernel

gpu_test/conftest.py

@@ -118,7 +118,10 @@ def _cleanup_orphans(self) -> None:
 
  def _launch(self) -> None:
  """Find the cheapest suitable offer and launch an instance."""
- query = f"num_gpus=1 rentable=True rented=False compute_cap>=700 dph<={MAX_COST_PER_HOUR}"
+ query = (
+ f"num_gpus=1 rentable=True rented=False compute_cap>=700"
+ f" reliability2>=0.95 inet_up>=100 dph<={MAX_COST_PER_HOUR}"
+ )
  offers = self.sdk.search_offers(query=query, order="dph", limit=5)
 
  if not offers:

gpu_test/test_kernels.py

@@ -4,6 +4,7 @@
 
 from typing import TYPE_CHECKING
 
+import numpy as np
 import pytest
 
 if TYPE_CHECKING:
@@ -583,3 +584,157 @@ def test_float_to_int_conversion(kernel_runner: KernelRunner) -> None:
  forth_source=("\\! kernel main\n\\! param DATA i64[256]\n7.9 F>S\n0 CELLS DATA + !"),
  )
  assert result[0] == 7
+
+
+# --- Attention ---
+
+_ATTENTION_KERNEL = """\
+\\! kernel attention
+\\! param Q f64[{n}]
+\\! param K f64[{n}]
+\\! param V f64[{n}]
+\\! param O f64[{n}]
+\\! param SEQ_LEN i64
+\\! param HEAD_DIM i64
+\\! shared SCORES f64[{seq_len}]
+\\! shared SCRATCH f64[{seq_len}]
+BID-X
+TID-X
+0.0
+HEAD_DIM 0 DO
+ 2 PICK HEAD_DIM * I + CELLS Q + F@
+ 2 PICK HEAD_DIM * I + CELLS K + F@
+ F* F+
+LOOP
+HEAD_DIM S>F FSQRT F/
+OVER 3 PICK >
+IF DROP -1.0e30 THEN
+OVER CELLS SCORES + SF!
+BARRIER
+TID-X 0= IF
+ 0 CELLS SCORES + SF@
+ SEQ_LEN 1 DO I CELLS SCORES + SF@ FMAX LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F- FEXP
+OVER CELLS SCORES + SF!
+BARRIER
+TID-X 0= IF
+ 0.0
+ SEQ_LEN 0 DO I CELLS SCORES + SF@ F+ LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F/
+OVER CELLS SCORES + SF!
+BARRIER
+DUP BEGIN DUP HEAD_DIM < WHILE
+ 0.0
+ SEQ_LEN 0 DO
+ I CELLS SCORES + SF@
+ I HEAD_DIM * 3 PICK + CELLS V + F@
+ F* F+
+ LOOP
+ OVER 4 PICK HEAD_DIM * + CELLS O + F!
+ BDIM-X +
+REPEAT
+DROP DROP DROP
+"""
+
+
+def _attention_reference(q: np.ndarray, k: np.ndarray, v: np.ndarray, seq_len: int) -> list[float]:
+ """Compute scaled dot-product attention with causal mask (NumPy reference)."""
+ head_dim = q.shape[1]
+ scores = q @ k.T / np.sqrt(head_dim)
+ causal_mask = np.triu(np.ones((seq_len, seq_len), dtype=bool), k=1)
+ scores[causal_mask] = -1e30
+ exp_scores = np.exp(scores - scores.max(axis=1, keepdims=True))
+ attn = exp_scores / exp_scores.sum(axis=1, keepdims=True)
+ return (attn @ v).flatten().tolist()
+
+
+def test_naive_attention_f64(kernel_runner: KernelRunner) -> None:
+ """Naive scaled dot-product attention with causal mask.
+
+ O = softmax(Q @ K^T / sqrt(d_k)) @ V, seq_len=4, head_dim=4.
+ One block per query row, one thread per key position.
+ """
+ seq_len, head_dim = 4, 4
+
+ q = np.array(
+ [
+ [1.0, 0.0, 1.0, 0.0],
+ [0.0, 1.0, 0.0, 1.0],
+ [1.0, 1.0, 0.0, 0.0],
+ [0.0, 0.0, 1.0, 1.0],
+ ]
+ )
+ k = np.array(
+ [
+ [1.0, 0.0, 0.0, 1.0],
+ [0.0, 1.0, 1.0, 0.0],
+ [1.0, 1.0, 0.0, 0.0],
+ [0.0, 0.0, 1.0, 1.0],
+ ]
+ )
+ v = np.array(
+ [
+ [1.0, 2.0, 3.0, 4.0],
+ [5.0, 6.0, 7.0, 8.0],
+ [9.0, 10.0, 11.0, 12.0],
+ [13.0, 14.0, 15.0, 16.0],
+ ]
+ )
+
+ expected = _attention_reference(q, k, v, seq_len)
+ n = seq_len * head_dim
+
+ result = kernel_runner.run(
+ forth_source=_ATTENTION_KERNEL.format(n=n, seq_len=seq_len),
+ params={
+ "Q": q.flatten().tolist(),
+ "K": k.flatten().tolist(),
+ "V": v.flatten().tolist(),
+ "SEQ_LEN": seq_len,
+ "HEAD_DIM": head_dim,
+ },
+ grid=(seq_len, 1, 1),
+ block=(seq_len, 1, 1),
+ output_param=3,
+ output_count=n,
+ )
+ assert result == [pytest.approx(v) for v in expected]
+
+
+def test_naive_attention_f64_16x64(kernel_runner: KernelRunner) -> None:
+ """Naive scaled dot-product attention, seq_len=16, head_dim=64."""
+ seq_len, head_dim = 16, 64
+
+ rng = np.random.default_rng(42)
+ q = rng.standard_normal((seq_len, head_dim))
+ k = rng.standard_normal((seq_len, head_dim))
+ v = rng.standard_normal((seq_len, head_dim))
+
+ expected = _attention_reference(q, k, v, seq_len)
+ n = seq_len * head_dim
+
+ result = kernel_runner.run(
+ forth_source=_ATTENTION_KERNEL.format(n=n, seq_len=seq_len),
+ params={
+ "Q": q.flatten().tolist(),
+ "K": k.flatten().tolist(),
+ "V": v.flatten().tolist(),
+ "SEQ_LEN": seq_len,
+ "HEAD_DIM": head_dim,
+ },
+ grid=(seq_len, 1, 1),
+ block=(seq_len, 1, 1),
+ output_param=3,
+ output_count=n,
+ )
+ assert result == [pytest.approx(v) for v in expected]

lib/Bitcode/CMakeLists.txt

@@ -0,0 +1,5 @@
+set(WARPFORTH_LIBDEVICE_PATH
+ "${CMAKE_CURRENT_BINARY_DIR}/libdevice.10.bc" PARENT_SCOPE)
+
+configure_file(libdevice.10.bc
+ "${CMAKE_CURRENT_BINARY_DIR}/libdevice.10.bc" COPYONLY)

lib/CMakeLists.txt

@@ -1,3 +1,4 @@
 add_subdirectory(Dialect)
+add_subdirectory(Bitcode)
 add_subdirectory(Conversion)
 add_subdirectory(Translation)

lib/Conversion/CMakeLists.txt

@@ -16,5 +16,9 @@ add_mlir_library(MLIRConversionPasses
  MLIRTransforms
 )
 
+target_compile_definitions(obj.MLIRConversionPasses PRIVATE
+ WARPFORTH_LIBDEVICE_PATH="${WARPFORTH_LIBDEVICE_PATH}"
+)
+
 add_subdirectory(ForthToMemRef)
 add_subdirectory(ForthToGPU)

lib/Conversion/Passes.cpp

@@ -33,7 +33,10 @@ void buildWarpForthPipeline(OpPassManager &pm) {
  pm.addPass(createCanonicalizerPass());
 
  // Stage 4: Attach NVVM target to GPU modules (sm_70 = Volta architecture)
- pm.addPass(createGpuNVVMAttachTarget());
+ GpuNVVMAttachTargetOptions nvvmOptions;
+ nvvmOptions.chip = "sm_70";
+ nvvmOptions.linkLibs.push_back(WARPFORTH_LIBDEVICE_PATH);
+ pm.addPass(createGpuNVVMAttachTarget(nvvmOptions));
 
  // Stage 5: Lower GPU to NVVM with bare pointers
  ConvertGpuOpsToNVVMOpsOptions gpuToNVVMOptions;

pyproject.toml

@@ -4,6 +4,7 @@ version = "0.1.0"
 requires-python = ">=3.11"
 dependencies = [
  "lit>=18.1.0",
+ "numpy",
  "pytest",
  "vastai-sdk",
 ]

test/Pipeline/attention.forth

@@ -0,0 +1,80 @@
+\ RUN: %warpforth-translate --forth-to-mlir %s | %warpforth-opt --warpforth-pipeline | %FileCheck %s
+
+\ Verify that a naive attention kernel with shared memory and float intrinsics
+\ survives the full pipeline to gpu.binary.
+\ CHECK: gpu.binary @warpforth_module
+
+\! kernel attention
+\! param Q f64[16]
+\! param K f64[16]
+\! param V f64[16]
+\! param O f64[16]
+\! param SEQ_LEN i64
+\! param HEAD_DIM i64
+\! shared SCORES f64[4]
+\! shared SCRATCH f64[4]
+
+\ row = BID-X, t = TID-X
+BID-X
+TID-X
+
+\ --- Dot product: Q[row,:] . K[t,:] ---
+0.0
+HEAD_DIM 0 DO
+ 2 PICK HEAD_DIM * I + CELLS Q + F@
+ 2 PICK HEAD_DIM * I + CELLS K + F@
+ F* F+
+LOOP
+HEAD_DIM S>F FSQRT F/
+
+\ --- Causal mask: if t > row, score = -inf ---
+OVER 3 PICK >
+IF DROP -1.0e30 THEN
+
+\ --- Store score to shared memory ---
+OVER CELLS SCORES + SF!
+BARRIER
+
+\ --- Softmax: max reduction (thread 0) ---
+TID-X 0= IF
+ 0 CELLS SCORES + SF@
+ SEQ_LEN 1 DO I CELLS SCORES + SF@ FMAX LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+
+\ --- Softmax: exp(score - max) ---
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F- FEXP
+OVER CELLS SCORES + SF!
+BARRIER
+
+\ --- Softmax: sum reduction (thread 0) ---
+TID-X 0= IF
+ 0.0
+ SEQ_LEN 0 DO I CELLS SCORES + SF@ F+ LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+
+\ --- Softmax: normalize ---
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F/
+OVER CELLS SCORES + SF!
+BARRIER
+
+\ --- V accumulation: O[row,col] = sum_j SCORES[j] * V[j*HD + col] ---
+\ Stride over head_dim columns: col = t, t+BDIM-X, t+2*BDIM-X, ...
+DUP BEGIN DUP HEAD_DIM < WHILE
+ 0.0
+ SEQ_LEN 0 DO
+ I CELLS SCORES + SF@
+ I HEAD_DIM * 3 PICK + CELLS V + F@
+ F* F+
+ LOOP
+ OVER 4 PICK HEAD_DIM * + CELLS O + F!
+ BDIM-X +
+REPEAT
+DROP DROP DROP

gpu_test/conftest.py

@@ -118,7 +118,10 @@ def _cleanup_orphans(self) -> None:
 
  def _launch(self) -> None:
  """Find the cheapest suitable offer and launch an instance."""
- query = f"num_gpus=1 rentable=True rented=False compute_cap>=700 dph<={MAX_COST_PER_HOUR}"
+ query = (
+ f"num_gpus=1 rentable=True rented=False compute_cap>=700"
+ f" reliability2>=0.95 inet_up>=100 dph<={MAX_COST_PER_HOUR}"
+ )
  offers = self.sdk.search_offers(query=query, order="dph", limit=5)
 
  if not offers:

gpu_test/test_kernels.py

@@ -4,6 +4,7 @@
 
 from typing import TYPE_CHECKING
 
+import numpy as np
 import pytest
 
 if TYPE_CHECKING:
@@ -583,3 +584,157 @@ def test_float_to_int_conversion(kernel_runner: KernelRunner) -> None:
  forth_source=("\\! kernel main\n\\! param DATA i64[256]\n7.9 F>S\n0 CELLS DATA + !"),
  )
  assert result[0] == 7
+
+
+# --- Attention ---
+
+_ATTENTION_KERNEL = """\
+\\! kernel attention
+\\! param Q f64[{n}]
+\\! param K f64[{n}]
+\\! param V f64[{n}]
+\\! param O f64[{n}]
+\\! param SEQ_LEN i64
+\\! param HEAD_DIM i64
+\\! shared SCORES f64[{seq_len}]
+\\! shared SCRATCH f64[{seq_len}]
+BID-X
+TID-X
+0.0
+HEAD_DIM 0 DO
+ 2 PICK HEAD_DIM * I + CELLS Q + F@
+ 2 PICK HEAD_DIM * I + CELLS K + F@
+ F* F+
+LOOP
+HEAD_DIM S>F FSQRT F/
+OVER 3 PICK >
+IF DROP -1.0e30 THEN
+OVER CELLS SCORES + SF!
+BARRIER
+TID-X 0= IF
+ 0 CELLS SCORES + SF@
+ SEQ_LEN 1 DO I CELLS SCORES + SF@ FMAX LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F- FEXP
+OVER CELLS SCORES + SF!
+BARRIER
+TID-X 0= IF
+ 0.0
+ SEQ_LEN 0 DO I CELLS SCORES + SF@ F+ LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F/
+OVER CELLS SCORES + SF!
+BARRIER
+DUP BEGIN DUP HEAD_DIM < WHILE
+ 0.0
+ SEQ_LEN 0 DO
+ I CELLS SCORES + SF@
+ I HEAD_DIM * 3 PICK + CELLS V + F@
+ F* F+
+ LOOP
+ OVER 4 PICK HEAD_DIM * + CELLS O + F!
+ BDIM-X +
+REPEAT
+DROP DROP DROP
+"""
+
+
+def _attention_reference(q: np.ndarray, k: np.ndarray, v: np.ndarray, seq_len: int) -> list[float]:
+ """Compute scaled dot-product attention with causal mask (NumPy reference)."""
+ head_dim = q.shape[1]
+ scores = q @ k.T / np.sqrt(head_dim)
+ causal_mask = np.triu(np.ones((seq_len, seq_len), dtype=bool), k=1)
+ scores[causal_mask] = -1e30
+ exp_scores = np.exp(scores - scores.max(axis=1, keepdims=True))
+ attn = exp_scores / exp_scores.sum(axis=1, keepdims=True)
+ return (attn @ v).flatten().tolist()
+
+
+def test_naive_attention_f64(kernel_runner: KernelRunner) -> None:
+ """Naive scaled dot-product attention with causal mask.
+
+ O = softmax(Q @ K^T / sqrt(d_k)) @ V, seq_len=4, head_dim=4.
+ One block per query row, one thread per key position.
+ """
+ seq_len, head_dim = 4, 4
+
+ q = np.array(
+ [
+ [1.0, 0.0, 1.0, 0.0],
+ [0.0, 1.0, 0.0, 1.0],
+ [1.0, 1.0, 0.0, 0.0],
+ [0.0, 0.0, 1.0, 1.0],
+ ]
+ )
+ k = np.array(
+ [
+ [1.0, 0.0, 0.0, 1.0],
+ [0.0, 1.0, 1.0, 0.0],
+ [1.0, 1.0, 0.0, 0.0],
+ [0.0, 0.0, 1.0, 1.0],
+ ]
+ )
+ v = np.array(
+ [
+ [1.0, 2.0, 3.0, 4.0],
+ [5.0, 6.0, 7.0, 8.0],
+ [9.0, 10.0, 11.0, 12.0],
+ [13.0, 14.0, 15.0, 16.0],
+ ]
+ )
+
+ expected = _attention_reference(q, k, v, seq_len)
+ n = seq_len * head_dim
+
+ result = kernel_runner.run(
+ forth_source=_ATTENTION_KERNEL.format(n=n, seq_len=seq_len),
+ params={
+ "Q": q.flatten().tolist(),
+ "K": k.flatten().tolist(),
+ "V": v.flatten().tolist(),
+ "SEQ_LEN": seq_len,
+ "HEAD_DIM": head_dim,
+ },
+ grid=(seq_len, 1, 1),
+ block=(seq_len, 1, 1),
+ output_param=3,
+ output_count=n,
+ )
+ assert result == [pytest.approx(v) for v in expected]
+
+
+def test_naive_attention_f64_16x64(kernel_runner: KernelRunner) -> None:
+ """Naive scaled dot-product attention, seq_len=16, head_dim=64."""
+ seq_len, head_dim = 16, 64
+
+ rng = np.random.default_rng(42)
+ q = rng.standard_normal((seq_len, head_dim))
+ k = rng.standard_normal((seq_len, head_dim))
+ v = rng.standard_normal((seq_len, head_dim))
+
+ expected = _attention_reference(q, k, v, seq_len)
+ n = seq_len * head_dim
+
+ result = kernel_runner.run(
+ forth_source=_ATTENTION_KERNEL.format(n=n, seq_len=seq_len),
+ params={
+ "Q": q.flatten().tolist(),
+ "K": k.flatten().tolist(),
+ "V": v.flatten().tolist(),
+ "SEQ_LEN": seq_len,
+ "HEAD_DIM": head_dim,
+ },
+ grid=(seq_len, 1, 1),
+ block=(seq_len, 1, 1),
+ output_param=3,
+ output_count=n,
+ )
+ assert result == [pytest.approx(v) for v in expected]

lib/Bitcode/CMakeLists.txt

@@ -0,0 +1,5 @@
+set(WARPFORTH_LIBDEVICE_PATH
+ "${CMAKE_CURRENT_BINARY_DIR}/libdevice.10.bc" PARENT_SCOPE)
+
+configure_file(libdevice.10.bc
+ "${CMAKE_CURRENT_BINARY_DIR}/libdevice.10.bc" COPYONLY)

lib/CMakeLists.txt

@@ -1,3 +1,4 @@
 add_subdirectory(Dialect)
+add_subdirectory(Bitcode)
 add_subdirectory(Conversion)
 add_subdirectory(Translation)

lib/Conversion/CMakeLists.txt

@@ -16,5 +16,9 @@ add_mlir_library(MLIRConversionPasses
  MLIRTransforms
 )
 
+target_compile_definitions(obj.MLIRConversionPasses PRIVATE
+ WARPFORTH_LIBDEVICE_PATH="${WARPFORTH_LIBDEVICE_PATH}"
+)
+
 add_subdirectory(ForthToMemRef)
 add_subdirectory(ForthToGPU)

lib/Conversion/Passes.cpp

@@ -33,7 +33,10 @@ void buildWarpForthPipeline(OpPassManager &pm) {
  pm.addPass(createCanonicalizerPass());
 
  // Stage 4: Attach NVVM target to GPU modules (sm_70 = Volta architecture)
- pm.addPass(createGpuNVVMAttachTarget());
+ GpuNVVMAttachTargetOptions nvvmOptions;
+ nvvmOptions.chip = "sm_70";
+ nvvmOptions.linkLibs.push_back(WARPFORTH_LIBDEVICE_PATH);
+ pm.addPass(createGpuNVVMAttachTarget(nvvmOptions));
 
  // Stage 5: Lower GPU to NVVM with bare pointers
  ConvertGpuOpsToNVVMOpsOptions gpuToNVVMOptions;

pyproject.toml

@@ -4,6 +4,7 @@ version = "0.1.0"
 requires-python = ">=3.11"
 dependencies = [
  "lit>=18.1.0",
+ "numpy",
  "pytest",
  "vastai-sdk",
 ]

test/Pipeline/attention.forth

@@ -0,0 +1,80 @@
+\ RUN: %warpforth-translate --forth-to-mlir %s | %warpforth-opt --warpforth-pipeline | %FileCheck %s
+
+\ Verify that a naive attention kernel with shared memory and float intrinsics
+\ survives the full pipeline to gpu.binary.
+\ CHECK: gpu.binary @warpforth_module
+
+\! kernel attention
+\! param Q f64[16]
+\! param K f64[16]
+\! param V f64[16]
+\! param O f64[16]
+\! param SEQ_LEN i64
+\! param HEAD_DIM i64
+\! shared SCORES f64[4]
+\! shared SCRATCH f64[4]
+
+\ row = BID-X, t = TID-X
+BID-X
+TID-X
+
+\ --- Dot product: Q[row,:] . K[t,:] ---
+0.0
+HEAD_DIM 0 DO
+ 2 PICK HEAD_DIM * I + CELLS Q + F@
+ 2 PICK HEAD_DIM * I + CELLS K + F@
+ F* F+
+LOOP
+HEAD_DIM S>F FSQRT F/
+
+\ --- Causal mask: if t > row, score = -inf ---
+OVER 3 PICK >
+IF DROP -1.0e30 THEN
+
+\ --- Store score to shared memory ---
+OVER CELLS SCORES + SF!
+BARRIER
+
+\ --- Softmax: max reduction (thread 0) ---
+TID-X 0= IF
+ 0 CELLS SCORES + SF@
+ SEQ_LEN 1 DO I CELLS SCORES + SF@ FMAX LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+
+\ --- Softmax: exp(score - max) ---
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F- FEXP
+OVER CELLS SCORES + SF!
+BARRIER
+
+\ --- Softmax: sum reduction (thread 0) ---
+TID-X 0= IF
+ 0.0
+ SEQ_LEN 0 DO I CELLS SCORES + SF@ F+ LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+
+\ --- Softmax: normalize ---
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F/
+OVER CELLS SCORES + SF!
+BARRIER
+
+\ --- V accumulation: O[row,col] = sum_j SCORES[j] * V[j*HD + col] ---
+\ Stride over head_dim columns: col = t, t+BDIM-X, t+2*BDIM-X, ...
+DUP BEGIN DUP HEAD_DIM < WHILE
+ 0.0
+ SEQ_LEN 0 DO
+ I CELLS SCORES + SF@
+ I HEAD_DIM * 3 PICK + CELLS V + F@
+ F* F+
+ LOOP
+ OVER 4 PICK HEAD_DIM * + CELLS O + F!
+ BDIM-X +
+REPEAT
+DROP DROP DROP

gpu_test/conftest.py

@@ -118,7 +118,10 @@ def _cleanup_orphans(self) -> None:
 
  def _launch(self) -> None:
  """Find the cheapest suitable offer and launch an instance."""
- query = f"num_gpus=1 rentable=True rented=False compute_cap>=700 dph<={MAX_COST_PER_HOUR}"
+ query = (
+ f"num_gpus=1 rentable=True rented=False compute_cap>=700"
+ f" reliability2>=0.95 inet_up>=100 dph<={MAX_COST_PER_HOUR}"
+ )
  offers = self.sdk.search_offers(query=query, order="dph", limit=5)
 
  if not offers:

gpu_test/test_kernels.py

@@ -4,6 +4,7 @@
 
 from typing import TYPE_CHECKING
 
+import numpy as np
 import pytest
 
 if TYPE_CHECKING:
@@ -583,3 +584,157 @@ def test_float_to_int_conversion(kernel_runner: KernelRunner) -> None:
  forth_source=("\\! kernel main\n\\! param DATA i64[256]\n7.9 F>S\n0 CELLS DATA + !"),
  )
  assert result[0] == 7
+
+
+# --- Attention ---
+
+_ATTENTION_KERNEL = """\
+\\! kernel attention
+\\! param Q f64[{n}]
+\\! param K f64[{n}]
+\\! param V f64[{n}]
+\\! param O f64[{n}]
+\\! param SEQ_LEN i64
+\\! param HEAD_DIM i64
+\\! shared SCORES f64[{seq_len}]
+\\! shared SCRATCH f64[{seq_len}]
+BID-X
+TID-X
+0.0
+HEAD_DIM 0 DO
+ 2 PICK HEAD_DIM * I + CELLS Q + F@
+ 2 PICK HEAD_DIM * I + CELLS K + F@
+ F* F+
+LOOP
+HEAD_DIM S>F FSQRT F/
+OVER 3 PICK >
+IF DROP -1.0e30 THEN
+OVER CELLS SCORES + SF!
+BARRIER
+TID-X 0= IF
+ 0 CELLS SCORES + SF@
+ SEQ_LEN 1 DO I CELLS SCORES + SF@ FMAX LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F- FEXP
+OVER CELLS SCORES + SF!
+BARRIER
+TID-X 0= IF
+ 0.0
+ SEQ_LEN 0 DO I CELLS SCORES + SF@ F+ LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F/
+OVER CELLS SCORES + SF!
+BARRIER
+DUP BEGIN DUP HEAD_DIM < WHILE
+ 0.0
+ SEQ_LEN 0 DO
+ I CELLS SCORES + SF@
+ I HEAD_DIM * 3 PICK + CELLS V + F@
+ F* F+
+ LOOP
+ OVER 4 PICK HEAD_DIM * + CELLS O + F!
+ BDIM-X +
+REPEAT
+DROP DROP DROP
+"""
+
+
+def _attention_reference(q: np.ndarray, k: np.ndarray, v: np.ndarray, seq_len: int) -> list[float]:
+ """Compute scaled dot-product attention with causal mask (NumPy reference)."""
+ head_dim = q.shape[1]
+ scores = q @ k.T / np.sqrt(head_dim)
+ causal_mask = np.triu(np.ones((seq_len, seq_len), dtype=bool), k=1)
+ scores[causal_mask] = -1e30
+ exp_scores = np.exp(scores - scores.max(axis=1, keepdims=True))
+ attn = exp_scores / exp_scores.sum(axis=1, keepdims=True)
+ return (attn @ v).flatten().tolist()
+
+
+def test_naive_attention_f64(kernel_runner: KernelRunner) -> None:
+ """Naive scaled dot-product attention with causal mask.
+
+ O = softmax(Q @ K^T / sqrt(d_k)) @ V, seq_len=4, head_dim=4.
+ One block per query row, one thread per key position.
+ """
+ seq_len, head_dim = 4, 4
+
+ q = np.array(
+ [
+ [1.0, 0.0, 1.0, 0.0],
+ [0.0, 1.0, 0.0, 1.0],
+ [1.0, 1.0, 0.0, 0.0],
+ [0.0, 0.0, 1.0, 1.0],
+ ]
+ )
+ k = np.array(
+ [
+ [1.0, 0.0, 0.0, 1.0],
+ [0.0, 1.0, 1.0, 0.0],
+ [1.0, 1.0, 0.0, 0.0],
+ [0.0, 0.0, 1.0, 1.0],
+ ]
+ )
+ v = np.array(
+ [
+ [1.0, 2.0, 3.0, 4.0],
+ [5.0, 6.0, 7.0, 8.0],
+ [9.0, 10.0, 11.0, 12.0],
+ [13.0, 14.0, 15.0, 16.0],
+ ]
+ )
+
+ expected = _attention_reference(q, k, v, seq_len)
+ n = seq_len * head_dim
+
+ result = kernel_runner.run(
+ forth_source=_ATTENTION_KERNEL.format(n=n, seq_len=seq_len),
+ params={
+ "Q": q.flatten().tolist(),
+ "K": k.flatten().tolist(),
+ "V": v.flatten().tolist(),
+ "SEQ_LEN": seq_len,
+ "HEAD_DIM": head_dim,
+ },
+ grid=(seq_len, 1, 1),
+ block=(seq_len, 1, 1),
+ output_param=3,
+ output_count=n,
+ )
+ assert result == [pytest.approx(v) for v in expected]
+
+
+def test_naive_attention_f64_16x64(kernel_runner: KernelRunner) -> None:
+ """Naive scaled dot-product attention, seq_len=16, head_dim=64."""
+ seq_len, head_dim = 16, 64
+
+ rng = np.random.default_rng(42)
+ q = rng.standard_normal((seq_len, head_dim))
+ k = rng.standard_normal((seq_len, head_dim))
+ v = rng.standard_normal((seq_len, head_dim))
+
+ expected = _attention_reference(q, k, v, seq_len)
+ n = seq_len * head_dim
+
+ result = kernel_runner.run(
+ forth_source=_ATTENTION_KERNEL.format(n=n, seq_len=seq_len),
+ params={
+ "Q": q.flatten().tolist(),
+ "K": k.flatten().tolist(),
+ "V": v.flatten().tolist(),
+ "SEQ_LEN": seq_len,
+ "HEAD_DIM": head_dim,
+ },
+ grid=(seq_len, 1, 1),
+ block=(seq_len, 1, 1),
+ output_param=3,
+ output_count=n,
+ )
+ assert result == [pytest.approx(v) for v in expected]

lib/Bitcode/CMakeLists.txt

@@ -0,0 +1,5 @@
+set(WARPFORTH_LIBDEVICE_PATH
+ "${CMAKE_CURRENT_BINARY_DIR}/libdevice.10.bc" PARENT_SCOPE)
+
+configure_file(libdevice.10.bc
+ "${CMAKE_CURRENT_BINARY_DIR}/libdevice.10.bc" COPYONLY)

lib/CMakeLists.txt

@@ -1,3 +1,4 @@
 add_subdirectory(Dialect)
+add_subdirectory(Bitcode)
 add_subdirectory(Conversion)
 add_subdirectory(Translation)

lib/Conversion/CMakeLists.txt

@@ -16,5 +16,9 @@ add_mlir_library(MLIRConversionPasses
  MLIRTransforms
 )
 
+target_compile_definitions(obj.MLIRConversionPasses PRIVATE
+ WARPFORTH_LIBDEVICE_PATH="${WARPFORTH_LIBDEVICE_PATH}"
+)
+
 add_subdirectory(ForthToMemRef)
 add_subdirectory(ForthToGPU)

lib/Conversion/Passes.cpp

@@ -33,7 +33,10 @@ void buildWarpForthPipeline(OpPassManager &pm) {
  pm.addPass(createCanonicalizerPass());
 
  // Stage 4: Attach NVVM target to GPU modules (sm_70 = Volta architecture)
- pm.addPass(createGpuNVVMAttachTarget());
+ GpuNVVMAttachTargetOptions nvvmOptions;
+ nvvmOptions.chip = "sm_70";
+ nvvmOptions.linkLibs.push_back(WARPFORTH_LIBDEVICE_PATH);
+ pm.addPass(createGpuNVVMAttachTarget(nvvmOptions));
 
  // Stage 5: Lower GPU to NVVM with bare pointers
  ConvertGpuOpsToNVVMOpsOptions gpuToNVVMOptions;

pyproject.toml

@@ -4,6 +4,7 @@ version = "0.1.0"
 requires-python = ">=3.11"
 dependencies = [
  "lit>=18.1.0",
+ "numpy",
  "pytest",
  "vastai-sdk",
 ]

test/Pipeline/attention.forth

@@ -0,0 +1,80 @@
+\ RUN: %warpforth-translate --forth-to-mlir %s | %warpforth-opt --warpforth-pipeline | %FileCheck %s
+
+\ Verify that a naive attention kernel with shared memory and float intrinsics
+\ survives the full pipeline to gpu.binary.
+\ CHECK: gpu.binary @warpforth_module
+
+\! kernel attention
+\! param Q f64[16]
+\! param K f64[16]
+\! param V f64[16]
+\! param O f64[16]
+\! param SEQ_LEN i64
+\! param HEAD_DIM i64
+\! shared SCORES f64[4]
+\! shared SCRATCH f64[4]
+
+\ row = BID-X, t = TID-X
+BID-X
+TID-X
+
+\ --- Dot product: Q[row,:] . K[t,:] ---
+0.0
+HEAD_DIM 0 DO
+ 2 PICK HEAD_DIM * I + CELLS Q + F@
+ 2 PICK HEAD_DIM * I + CELLS K + F@
+ F* F+
+LOOP
+HEAD_DIM S>F FSQRT F/
+
+\ --- Causal mask: if t > row, score = -inf ---
+OVER 3 PICK >
+IF DROP -1.0e30 THEN
+
+\ --- Store score to shared memory ---
+OVER CELLS SCORES + SF!
+BARRIER
+
+\ --- Softmax: max reduction (thread 0) ---
+TID-X 0= IF
+ 0 CELLS SCORES + SF@
+ SEQ_LEN 1 DO I CELLS SCORES + SF@ FMAX LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+
+\ --- Softmax: exp(score - max) ---
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F- FEXP
+OVER CELLS SCORES + SF!
+BARRIER
+
+\ --- Softmax: sum reduction (thread 0) ---
+TID-X 0= IF
+ 0.0
+ SEQ_LEN 0 DO I CELLS SCORES + SF@ F+ LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+
+\ --- Softmax: normalize ---
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F/
+OVER CELLS SCORES + SF!
+BARRIER
+
+\ --- V accumulation: O[row,col] = sum_j SCORES[j] * V[j*HD + col] ---
+\ Stride over head_dim columns: col = t, t+BDIM-X, t+2*BDIM-X, ...
+DUP BEGIN DUP HEAD_DIM < WHILE
+ 0.0
+ SEQ_LEN 0 DO
+ I CELLS SCORES + SF@
+ I HEAD_DIM * 3 PICK + CELLS V + F@
+ F* F+
+ LOOP
+ OVER 4 PICK HEAD_DIM * + CELLS O + F!
+ BDIM-X +
+REPEAT
+DROP DROP DROP

gpu_test/conftest.py

@@ -118,7 +118,10 @@ def _cleanup_orphans(self) -> None:
 
  def _launch(self) -> None:
  """Find the cheapest suitable offer and launch an instance."""
- query = f"num_gpus=1 rentable=True rented=False compute_cap>=700 dph<={MAX_COST_PER_HOUR}"
+ query = (
+ f"num_gpus=1 rentable=True rented=False compute_cap>=700"
+ f" reliability2>=0.95 inet_up>=100 dph<={MAX_COST_PER_HOUR}"
+ )
  offers = self.sdk.search_offers(query=query, order="dph", limit=5)
 
  if not offers:

gpu_test/test_kernels.py

@@ -4,6 +4,7 @@
 
 from typing import TYPE_CHECKING
 
+import numpy as np
 import pytest
 
 if TYPE_CHECKING:
@@ -583,3 +584,157 @@ def test_float_to_int_conversion(kernel_runner: KernelRunner) -> None:
  forth_source=("\\! kernel main\n\\! param DATA i64[256]\n7.9 F>S\n0 CELLS DATA + !"),
  )
  assert result[0] == 7
+
+
+# --- Attention ---
+
+_ATTENTION_KERNEL = """\
+\\! kernel attention
+\\! param Q f64[{n}]
+\\! param K f64[{n}]
+\\! param V f64[{n}]
+\\! param O f64[{n}]
+\\! param SEQ_LEN i64
+\\! param HEAD_DIM i64
+\\! shared SCORES f64[{seq_len}]
+\\! shared SCRATCH f64[{seq_len}]
+BID-X
+TID-X
+0.0
+HEAD_DIM 0 DO
+ 2 PICK HEAD_DIM * I + CELLS Q + F@
+ 2 PICK HEAD_DIM * I + CELLS K + F@
+ F* F+
+LOOP
+HEAD_DIM S>F FSQRT F/
+OVER 3 PICK >
+IF DROP -1.0e30 THEN
+OVER CELLS SCORES + SF!
+BARRIER
+TID-X 0= IF
+ 0 CELLS SCORES + SF@
+ SEQ_LEN 1 DO I CELLS SCORES + SF@ FMAX LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F- FEXP
+OVER CELLS SCORES + SF!
+BARRIER
+TID-X 0= IF
+ 0.0
+ SEQ_LEN 0 DO I CELLS SCORES + SF@ F+ LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F/
+OVER CELLS SCORES + SF!
+BARRIER
+DUP BEGIN DUP HEAD_DIM < WHILE
+ 0.0
+ SEQ_LEN 0 DO
+ I CELLS SCORES + SF@
+ I HEAD_DIM * 3 PICK + CELLS V + F@
+ F* F+
+ LOOP
+ OVER 4 PICK HEAD_DIM * + CELLS O + F!
+ BDIM-X +
+REPEAT
+DROP DROP DROP
+"""
+
+
+def _attention_reference(q: np.ndarray, k: np.ndarray, v: np.ndarray, seq_len: int) -> list[float]:
+ """Compute scaled dot-product attention with causal mask (NumPy reference)."""
+ head_dim = q.shape[1]
+ scores = q @ k.T / np.sqrt(head_dim)
+ causal_mask = np.triu(np.ones((seq_len, seq_len), dtype=bool), k=1)
+ scores[causal_mask] = -1e30
+ exp_scores = np.exp(scores - scores.max(axis=1, keepdims=True))
+ attn = exp_scores / exp_scores.sum(axis=1, keepdims=True)
+ return (attn @ v).flatten().tolist()
+
+
+def test_naive_attention_f64(kernel_runner: KernelRunner) -> None:
+ """Naive scaled dot-product attention with causal mask.
+
+ O = softmax(Q @ K^T / sqrt(d_k)) @ V, seq_len=4, head_dim=4.
+ One block per query row, one thread per key position.
+ """
+ seq_len, head_dim = 4, 4
+
+ q = np.array(
+ [
+ [1.0, 0.0, 1.0, 0.0],
+ [0.0, 1.0, 0.0, 1.0],
+ [1.0, 1.0, 0.0, 0.0],
+ [0.0, 0.0, 1.0, 1.0],
+ ]
+ )
+ k = np.array(
+ [
+ [1.0, 0.0, 0.0, 1.0],
+ [0.0, 1.0, 1.0, 0.0],
+ [1.0, 1.0, 0.0, 0.0],
+ [0.0, 0.0, 1.0, 1.0],
+ ]
+ )
+ v = np.array(
+ [
+ [1.0, 2.0, 3.0, 4.0],
+ [5.0, 6.0, 7.0, 8.0],
+ [9.0, 10.0, 11.0, 12.0],
+ [13.0, 14.0, 15.0, 16.0],
+ ]
+ )
+
+ expected = _attention_reference(q, k, v, seq_len)
+ n = seq_len * head_dim
+
+ result = kernel_runner.run(
+ forth_source=_ATTENTION_KERNEL.format(n=n, seq_len=seq_len),
+ params={
+ "Q": q.flatten().tolist(),
+ "K": k.flatten().tolist(),
+ "V": v.flatten().tolist(),
+ "SEQ_LEN": seq_len,
+ "HEAD_DIM": head_dim,
+ },
+ grid=(seq_len, 1, 1),
+ block=(seq_len, 1, 1),
+ output_param=3,
+ output_count=n,
+ )
+ assert result == [pytest.approx(v) for v in expected]
+
+
+def test_naive_attention_f64_16x64(kernel_runner: KernelRunner) -> None:
+ """Naive scaled dot-product attention, seq_len=16, head_dim=64."""
+ seq_len, head_dim = 16, 64
+
+ rng = np.random.default_rng(42)
+ q = rng.standard_normal((seq_len, head_dim))
+ k = rng.standard_normal((seq_len, head_dim))
+ v = rng.standard_normal((seq_len, head_dim))
+
+ expected = _attention_reference(q, k, v, seq_len)
+ n = seq_len * head_dim
+
+ result = kernel_runner.run(
+ forth_source=_ATTENTION_KERNEL.format(n=n, seq_len=seq_len),
+ params={
+ "Q": q.flatten().tolist(),
+ "K": k.flatten().tolist(),
+ "V": v.flatten().tolist(),
+ "SEQ_LEN": seq_len,
+ "HEAD_DIM": head_dim,
+ },
+ grid=(seq_len, 1, 1),
+ block=(seq_len, 1, 1),
+ output_param=3,
+ output_count=n,
+ )
+ assert result == [pytest.approx(v) for v in expected]

lib/Bitcode/CMakeLists.txt

@@ -0,0 +1,5 @@
+set(WARPFORTH_LIBDEVICE_PATH
+ "${CMAKE_CURRENT_BINARY_DIR}/libdevice.10.bc" PARENT_SCOPE)
+
+configure_file(libdevice.10.bc
+ "${CMAKE_CURRENT_BINARY_DIR}/libdevice.10.bc" COPYONLY)

lib/CMakeLists.txt

@@ -1,3 +1,4 @@
 add_subdirectory(Dialect)
+add_subdirectory(Bitcode)
 add_subdirectory(Conversion)
 add_subdirectory(Translation)

lib/Conversion/CMakeLists.txt

@@ -16,5 +16,9 @@ add_mlir_library(MLIRConversionPasses
  MLIRTransforms
 )
 
+target_compile_definitions(obj.MLIRConversionPasses PRIVATE
+ WARPFORTH_LIBDEVICE_PATH="${WARPFORTH_LIBDEVICE_PATH}"
+)
+
 add_subdirectory(ForthToMemRef)
 add_subdirectory(ForthToGPU)

lib/Conversion/Passes.cpp

@@ -33,7 +33,10 @@ void buildWarpForthPipeline(OpPassManager &pm) {
  pm.addPass(createCanonicalizerPass());
 
  // Stage 4: Attach NVVM target to GPU modules (sm_70 = Volta architecture)
- pm.addPass(createGpuNVVMAttachTarget());
+ GpuNVVMAttachTargetOptions nvvmOptions;
+ nvvmOptions.chip = "sm_70";
+ nvvmOptions.linkLibs.push_back(WARPFORTH_LIBDEVICE_PATH);
+ pm.addPass(createGpuNVVMAttachTarget(nvvmOptions));
 
  // Stage 5: Lower GPU to NVVM with bare pointers
  ConvertGpuOpsToNVVMOpsOptions gpuToNVVMOptions;

pyproject.toml

@@ -4,6 +4,7 @@ version = "0.1.0"
 requires-python = ">=3.11"
 dependencies = [
  "lit>=18.1.0",
+ "numpy",
  "pytest",
  "vastai-sdk",
 ]

test/Pipeline/attention.forth

@@ -0,0 +1,80 @@
+\ RUN: %warpforth-translate --forth-to-mlir %s | %warpforth-opt --warpforth-pipeline | %FileCheck %s
+
+\ Verify that a naive attention kernel with shared memory and float intrinsics
+\ survives the full pipeline to gpu.binary.
+\ CHECK: gpu.binary @warpforth_module
+
+\! kernel attention
+\! param Q f64[16]
+\! param K f64[16]
+\! param V f64[16]
+\! param O f64[16]
+\! param SEQ_LEN i64
+\! param HEAD_DIM i64
+\! shared SCORES f64[4]
+\! shared SCRATCH f64[4]
+
+\ row = BID-X, t = TID-X
+BID-X
+TID-X
+
+\ --- Dot product: Q[row,:] . K[t,:] ---
+0.0
+HEAD_DIM 0 DO
+ 2 PICK HEAD_DIM * I + CELLS Q + F@
+ 2 PICK HEAD_DIM * I + CELLS K + F@
+ F* F+
+LOOP
+HEAD_DIM S>F FSQRT F/
+
+\ --- Causal mask: if t > row, score = -inf ---
+OVER 3 PICK >
+IF DROP -1.0e30 THEN
+
+\ --- Store score to shared memory ---
+OVER CELLS SCORES + SF!
+BARRIER
+
+\ --- Softmax: max reduction (thread 0) ---
+TID-X 0= IF
+ 0 CELLS SCORES + SF@
+ SEQ_LEN 1 DO I CELLS SCORES + SF@ FMAX LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+
+\ --- Softmax: exp(score - max) ---
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F- FEXP
+OVER CELLS SCORES + SF!
+BARRIER
+
+\ --- Softmax: sum reduction (thread 0) ---
+TID-X 0= IF
+ 0.0
+ SEQ_LEN 0 DO I CELLS SCORES + SF@ F+ LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+
+\ --- Softmax: normalize ---
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F/
+OVER CELLS SCORES + SF!
+BARRIER
+
+\ --- V accumulation: O[row,col] = sum_j SCORES[j] * V[j*HD + col] ---
+\ Stride over head_dim columns: col = t, t+BDIM-X, t+2*BDIM-X, ...
+DUP BEGIN DUP HEAD_DIM < WHILE
+ 0.0
+ SEQ_LEN 0 DO
+ I CELLS SCORES + SF@
+ I HEAD_DIM * 3 PICK + CELLS V + F@
+ F* F+
+ LOOP
+ OVER 4 PICK HEAD_DIM * + CELLS O + F!
+ BDIM-X +
+REPEAT
+DROP DROP DROP

gpu_test/conftest.py

@@ -118,7 +118,10 @@ def _cleanup_orphans(self) -> None:
 
  def _launch(self) -> None:
  """Find the cheapest suitable offer and launch an instance."""
- query = f"num_gpus=1 rentable=True rented=False compute_cap>=700 dph<={MAX_COST_PER_HOUR}"
+ query = (
+ f"num_gpus=1 rentable=True rented=False compute_cap>=700"
+ f" reliability2>=0.95 inet_up>=100 dph<={MAX_COST_PER_HOUR}"
+ )
  offers = self.sdk.search_offers(query=query, order="dph", limit=5)
 
  if not offers:

gpu_test/test_kernels.py

@@ -4,6 +4,7 @@
 
 from typing import TYPE_CHECKING
 
+import numpy as np
 import pytest
 
 if TYPE_CHECKING:
@@ -583,3 +584,157 @@ def test_float_to_int_conversion(kernel_runner: KernelRunner) -> None:
  forth_source=("\\! kernel main\n\\! param DATA i64[256]\n7.9 F>S\n0 CELLS DATA + !"),
  )
  assert result[0] == 7
+
+
+# --- Attention ---
+
+_ATTENTION_KERNEL = """\
+\\! kernel attention
+\\! param Q f64[{n}]
+\\! param K f64[{n}]
+\\! param V f64[{n}]
+\\! param O f64[{n}]
+\\! param SEQ_LEN i64
+\\! param HEAD_DIM i64
+\\! shared SCORES f64[{seq_len}]
+\\! shared SCRATCH f64[{seq_len}]
+BID-X
+TID-X
+0.0
+HEAD_DIM 0 DO
+ 2 PICK HEAD_DIM * I + CELLS Q + F@
+ 2 PICK HEAD_DIM * I + CELLS K + F@
+ F* F+
+LOOP
+HEAD_DIM S>F FSQRT F/
+OVER 3 PICK >
+IF DROP -1.0e30 THEN
+OVER CELLS SCORES + SF!
+BARRIER
+TID-X 0= IF
+ 0 CELLS SCORES + SF@
+ SEQ_LEN 1 DO I CELLS SCORES + SF@ FMAX LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F- FEXP
+OVER CELLS SCORES + SF!
+BARRIER
+TID-X 0= IF
+ 0.0
+ SEQ_LEN 0 DO I CELLS SCORES + SF@ F+ LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F/
+OVER CELLS SCORES + SF!
+BARRIER
+DUP BEGIN DUP HEAD_DIM < WHILE
+ 0.0
+ SEQ_LEN 0 DO
+ I CELLS SCORES + SF@
+ I HEAD_DIM * 3 PICK + CELLS V + F@
+ F* F+
+ LOOP
+ OVER 4 PICK HEAD_DIM * + CELLS O + F!
+ BDIM-X +
+REPEAT
+DROP DROP DROP
+"""
+
+
+def _attention_reference(q: np.ndarray, k: np.ndarray, v: np.ndarray, seq_len: int) -> list[float]:
+ """Compute scaled dot-product attention with causal mask (NumPy reference)."""
+ head_dim = q.shape[1]
+ scores = q @ k.T / np.sqrt(head_dim)
+ causal_mask = np.triu(np.ones((seq_len, seq_len), dtype=bool), k=1)
+ scores[causal_mask] = -1e30
+ exp_scores = np.exp(scores - scores.max(axis=1, keepdims=True))
+ attn = exp_scores / exp_scores.sum(axis=1, keepdims=True)
+ return (attn @ v).flatten().tolist()
+
+
+def test_naive_attention_f64(kernel_runner: KernelRunner) -> None:
+ """Naive scaled dot-product attention with causal mask.
+
+ O = softmax(Q @ K^T / sqrt(d_k)) @ V, seq_len=4, head_dim=4.
+ One block per query row, one thread per key position.
+ """
+ seq_len, head_dim = 4, 4
+
+ q = np.array(
+ [
+ [1.0, 0.0, 1.0, 0.0],
+ [0.0, 1.0, 0.0, 1.0],
+ [1.0, 1.0, 0.0, 0.0],
+ [0.0, 0.0, 1.0, 1.0],
+ ]
+ )
+ k = np.array(
+ [
+ [1.0, 0.0, 0.0, 1.0],
+ [0.0, 1.0, 1.0, 0.0],
+ [1.0, 1.0, 0.0, 0.0],
+ [0.0, 0.0, 1.0, 1.0],
+ ]
+ )
+ v = np.array(
+ [
+ [1.0, 2.0, 3.0, 4.0],
+ [5.0, 6.0, 7.0, 8.0],
+ [9.0, 10.0, 11.0, 12.0],
+ [13.0, 14.0, 15.0, 16.0],
+ ]
+ )
+
+ expected = _attention_reference(q, k, v, seq_len)
+ n = seq_len * head_dim
+
+ result = kernel_runner.run(
+ forth_source=_ATTENTION_KERNEL.format(n=n, seq_len=seq_len),
+ params={
+ "Q": q.flatten().tolist(),
+ "K": k.flatten().tolist(),
+ "V": v.flatten().tolist(),
+ "SEQ_LEN": seq_len,
+ "HEAD_DIM": head_dim,
+ },
+ grid=(seq_len, 1, 1),
+ block=(seq_len, 1, 1),
+ output_param=3,
+ output_count=n,
+ )
+ assert result == [pytest.approx(v) for v in expected]
+
+
+def test_naive_attention_f64_16x64(kernel_runner: KernelRunner) -> None:
+ """Naive scaled dot-product attention, seq_len=16, head_dim=64."""
+ seq_len, head_dim = 16, 64
+
+ rng = np.random.default_rng(42)
+ q = rng.standard_normal((seq_len, head_dim))
+ k = rng.standard_normal((seq_len, head_dim))
+ v = rng.standard_normal((seq_len, head_dim))
+
+ expected = _attention_reference(q, k, v, seq_len)
+ n = seq_len * head_dim
+
+ result = kernel_runner.run(
+ forth_source=_ATTENTION_KERNEL.format(n=n, seq_len=seq_len),
+ params={
+ "Q": q.flatten().tolist(),
+ "K": k.flatten().tolist(),
+ "V": v.flatten().tolist(),
+ "SEQ_LEN": seq_len,
+ "HEAD_DIM": head_dim,
+ },
+ grid=(seq_len, 1, 1),
+ block=(seq_len, 1, 1),
+ output_param=3,
+ output_count=n,
+ )
+ assert result == [pytest.approx(v) for v in expected]

lib/Bitcode/CMakeLists.txt

@@ -0,0 +1,5 @@
+set(WARPFORTH_LIBDEVICE_PATH
+ "${CMAKE_CURRENT_BINARY_DIR}/libdevice.10.bc" PARENT_SCOPE)
+
+configure_file(libdevice.10.bc
+ "${CMAKE_CURRENT_BINARY_DIR}/libdevice.10.bc" COPYONLY)

lib/CMakeLists.txt

@@ -1,3 +1,4 @@
 add_subdirectory(Dialect)
+add_subdirectory(Bitcode)
 add_subdirectory(Conversion)
 add_subdirectory(Translation)

lib/Conversion/CMakeLists.txt

@@ -16,5 +16,9 @@ add_mlir_library(MLIRConversionPasses
  MLIRTransforms
 )
 
+target_compile_definitions(obj.MLIRConversionPasses PRIVATE
+ WARPFORTH_LIBDEVICE_PATH="${WARPFORTH_LIBDEVICE_PATH}"
+)
+
 add_subdirectory(ForthToMemRef)
 add_subdirectory(ForthToGPU)

lib/Conversion/Passes.cpp

@@ -33,7 +33,10 @@ void buildWarpForthPipeline(OpPassManager &pm) {
  pm.addPass(createCanonicalizerPass());
 
  // Stage 4: Attach NVVM target to GPU modules (sm_70 = Volta architecture)
- pm.addPass(createGpuNVVMAttachTarget());
+ GpuNVVMAttachTargetOptions nvvmOptions;
+ nvvmOptions.chip = "sm_70";
+ nvvmOptions.linkLibs.push_back(WARPFORTH_LIBDEVICE_PATH);
+ pm.addPass(createGpuNVVMAttachTarget(nvvmOptions));
 
  // Stage 5: Lower GPU to NVVM with bare pointers
  ConvertGpuOpsToNVVMOpsOptions gpuToNVVMOptions;

pyproject.toml

@@ -4,6 +4,7 @@ version = "0.1.0"
 requires-python = ">=3.11"
 dependencies = [
  "lit>=18.1.0",
+ "numpy",
  "pytest",
  "vastai-sdk",
 ]

test/Pipeline/attention.forth

@@ -0,0 +1,80 @@
+\ RUN: %warpforth-translate --forth-to-mlir %s | %warpforth-opt --warpforth-pipeline | %FileCheck %s
+
+\ Verify that a naive attention kernel with shared memory and float intrinsics
+\ survives the full pipeline to gpu.binary.
+\ CHECK: gpu.binary @warpforth_module
+
+\! kernel attention
+\! param Q f64[16]
+\! param K f64[16]
+\! param V f64[16]
+\! param O f64[16]
+\! param SEQ_LEN i64
+\! param HEAD_DIM i64
+\! shared SCORES f64[4]
+\! shared SCRATCH f64[4]
+
+\ row = BID-X, t = TID-X
+BID-X
+TID-X
+
+\ --- Dot product: Q[row,:] . K[t,:] ---
+0.0
+HEAD_DIM 0 DO
+ 2 PICK HEAD_DIM * I + CELLS Q + F@
+ 2 PICK HEAD_DIM * I + CELLS K + F@
+ F* F+
+LOOP
+HEAD_DIM S>F FSQRT F/
+
+\ --- Causal mask: if t > row, score = -inf ---
+OVER 3 PICK >
+IF DROP -1.0e30 THEN
+
+\ --- Store score to shared memory ---
+OVER CELLS SCORES + SF!
+BARRIER
+
+\ --- Softmax: max reduction (thread 0) ---
+TID-X 0= IF
+ 0 CELLS SCORES + SF@
+ SEQ_LEN 1 DO I CELLS SCORES + SF@ FMAX LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+
+\ --- Softmax: exp(score - max) ---
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F- FEXP
+OVER CELLS SCORES + SF!
+BARRIER
+
+\ --- Softmax: sum reduction (thread 0) ---
+TID-X 0= IF
+ 0.0
+ SEQ_LEN 0 DO I CELLS SCORES + SF@ F+ LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+
+\ --- Softmax: normalize ---
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F/
+OVER CELLS SCORES + SF!
+BARRIER
+
+\ --- V accumulation: O[row,col] = sum_j SCORES[j] * V[j*HD + col] ---
+\ Stride over head_dim columns: col = t, t+BDIM-X, t+2*BDIM-X, ...
+DUP BEGIN DUP HEAD_DIM < WHILE
+ 0.0
+ SEQ_LEN 0 DO
+ I CELLS SCORES + SF@
+ I HEAD_DIM * 3 PICK + CELLS V + F@
+ F* F+
+ LOOP
+ OVER 4 PICK HEAD_DIM * + CELLS O + F!
+ BDIM-X +
+REPEAT
+DROP DROP DROP

gpu_test/conftest.py

@@ -118,7 +118,10 @@ def _cleanup_orphans(self) -> None:
 
  def _launch(self) -> None:
  """Find the cheapest suitable offer and launch an instance."""
- query = f"num_gpus=1 rentable=True rented=False compute_cap>=700 dph<={MAX_COST_PER_HOUR}"
+ query = (
+ f"num_gpus=1 rentable=True rented=False compute_cap>=700"
+ f" reliability2>=0.95 inet_up>=100 dph<={MAX_COST_PER_HOUR}"
+ )
  offers = self.sdk.search_offers(query=query, order="dph", limit=5)
 
  if not offers:

gpu_test/test_kernels.py

@@ -4,6 +4,7 @@
 
 from typing import TYPE_CHECKING
 
+import numpy as np
 import pytest
 
 if TYPE_CHECKING:
@@ -583,3 +584,157 @@ def test_float_to_int_conversion(kernel_runner: KernelRunner) -> None:
  forth_source=("\\! kernel main\n\\! param DATA i64[256]\n7.9 F>S\n0 CELLS DATA + !"),
  )
  assert result[0] == 7
+
+
+# --- Attention ---
+
+_ATTENTION_KERNEL = """\
+\\! kernel attention
+\\! param Q f64[{n}]
+\\! param K f64[{n}]
+\\! param V f64[{n}]
+\\! param O f64[{n}]
+\\! param SEQ_LEN i64
+\\! param HEAD_DIM i64
+\\! shared SCORES f64[{seq_len}]
+\\! shared SCRATCH f64[{seq_len}]
+BID-X
+TID-X
+0.0
+HEAD_DIM 0 DO
+ 2 PICK HEAD_DIM * I + CELLS Q + F@
+ 2 PICK HEAD_DIM * I + CELLS K + F@
+ F* F+
+LOOP
+HEAD_DIM S>F FSQRT F/
+OVER 3 PICK >
+IF DROP -1.0e30 THEN
+OVER CELLS SCORES + SF!
+BARRIER
+TID-X 0= IF
+ 0 CELLS SCORES + SF@
+ SEQ_LEN 1 DO I CELLS SCORES + SF@ FMAX LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F- FEXP
+OVER CELLS SCORES + SF!
+BARRIER
+TID-X 0= IF
+ 0.0
+ SEQ_LEN 0 DO I CELLS SCORES + SF@ F+ LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F/
+OVER CELLS SCORES + SF!
+BARRIER
+DUP BEGIN DUP HEAD_DIM < WHILE
+ 0.0
+ SEQ_LEN 0 DO
+ I CELLS SCORES + SF@
+ I HEAD_DIM * 3 PICK + CELLS V + F@
+ F* F+
+ LOOP
+ OVER 4 PICK HEAD_DIM * + CELLS O + F!
+ BDIM-X +
+REPEAT
+DROP DROP DROP
+"""
+
+
+def _attention_reference(q: np.ndarray, k: np.ndarray, v: np.ndarray, seq_len: int) -> list[float]:
+ """Compute scaled dot-product attention with causal mask (NumPy reference)."""
+ head_dim = q.shape[1]
+ scores = q @ k.T / np.sqrt(head_dim)
+ causal_mask = np.triu(np.ones((seq_len, seq_len), dtype=bool), k=1)
+ scores[causal_mask] = -1e30
+ exp_scores = np.exp(scores - scores.max(axis=1, keepdims=True))
+ attn = exp_scores / exp_scores.sum(axis=1, keepdims=True)
+ return (attn @ v).flatten().tolist()
+
+
+def test_naive_attention_f64(kernel_runner: KernelRunner) -> None:
+ """Naive scaled dot-product attention with causal mask.
+
+ O = softmax(Q @ K^T / sqrt(d_k)) @ V, seq_len=4, head_dim=4.
+ One block per query row, one thread per key position.
+ """
+ seq_len, head_dim = 4, 4
+
+ q = np.array(
+ [
+ [1.0, 0.0, 1.0, 0.0],
+ [0.0, 1.0, 0.0, 1.0],
+ [1.0, 1.0, 0.0, 0.0],
+ [0.0, 0.0, 1.0, 1.0],
+ ]
+ )
+ k = np.array(
+ [
+ [1.0, 0.0, 0.0, 1.0],
+ [0.0, 1.0, 1.0, 0.0],
+ [1.0, 1.0, 0.0, 0.0],
+ [0.0, 0.0, 1.0, 1.0],
+ ]
+ )
+ v = np.array(
+ [
+ [1.0, 2.0, 3.0, 4.0],
+ [5.0, 6.0, 7.0, 8.0],
+ [9.0, 10.0, 11.0, 12.0],
+ [13.0, 14.0, 15.0, 16.0],
+ ]
+ )
+
+ expected = _attention_reference(q, k, v, seq_len)
+ n = seq_len * head_dim
+
+ result = kernel_runner.run(
+ forth_source=_ATTENTION_KERNEL.format(n=n, seq_len=seq_len),
+ params={
+ "Q": q.flatten().tolist(),
+ "K": k.flatten().tolist(),
+ "V": v.flatten().tolist(),
+ "SEQ_LEN": seq_len,
+ "HEAD_DIM": head_dim,
+ },
+ grid=(seq_len, 1, 1),
+ block=(seq_len, 1, 1),
+ output_param=3,
+ output_count=n,
+ )
+ assert result == [pytest.approx(v) for v in expected]
+
+
+def test_naive_attention_f64_16x64(kernel_runner: KernelRunner) -> None:
+ """Naive scaled dot-product attention, seq_len=16, head_dim=64."""
+ seq_len, head_dim = 16, 64
+
+ rng = np.random.default_rng(42)
+ q = rng.standard_normal((seq_len, head_dim))
+ k = rng.standard_normal((seq_len, head_dim))
+ v = rng.standard_normal((seq_len, head_dim))
+
+ expected = _attention_reference(q, k, v, seq_len)
+ n = seq_len * head_dim
+
+ result = kernel_runner.run(
+ forth_source=_ATTENTION_KERNEL.format(n=n, seq_len=seq_len),
+ params={
+ "Q": q.flatten().tolist(),
+ "K": k.flatten().tolist(),
+ "V": v.flatten().tolist(),
+ "SEQ_LEN": seq_len,
+ "HEAD_DIM": head_dim,
+ },
+ grid=(seq_len, 1, 1),
+ block=(seq_len, 1, 1),
+ output_param=3,
+ output_count=n,
+ )
+ assert result == [pytest.approx(v) for v in expected]

lib/Bitcode/CMakeLists.txt

@@ -0,0 +1,5 @@
+set(WARPFORTH_LIBDEVICE_PATH
+ "${CMAKE_CURRENT_BINARY_DIR}/libdevice.10.bc" PARENT_SCOPE)
+
+configure_file(libdevice.10.bc
+ "${CMAKE_CURRENT_BINARY_DIR}/libdevice.10.bc" COPYONLY)

lib/CMakeLists.txt

@@ -1,3 +1,4 @@
 add_subdirectory(Dialect)
+add_subdirectory(Bitcode)
 add_subdirectory(Conversion)
 add_subdirectory(Translation)

lib/Conversion/CMakeLists.txt

@@ -16,5 +16,9 @@ add_mlir_library(MLIRConversionPasses
  MLIRTransforms
 )
 
+target_compile_definitions(obj.MLIRConversionPasses PRIVATE
+ WARPFORTH_LIBDEVICE_PATH="${WARPFORTH_LIBDEVICE_PATH}"
+)
+
 add_subdirectory(ForthToMemRef)
 add_subdirectory(ForthToGPU)

lib/Conversion/Passes.cpp

@@ -33,7 +33,10 @@ void buildWarpForthPipeline(OpPassManager &pm) {
  pm.addPass(createCanonicalizerPass());
 
  // Stage 4: Attach NVVM target to GPU modules (sm_70 = Volta architecture)
- pm.addPass(createGpuNVVMAttachTarget());
+ GpuNVVMAttachTargetOptions nvvmOptions;
+ nvvmOptions.chip = "sm_70";
+ nvvmOptions.linkLibs.push_back(WARPFORTH_LIBDEVICE_PATH);
+ pm.addPass(createGpuNVVMAttachTarget(nvvmOptions));
 
  // Stage 5: Lower GPU to NVVM with bare pointers
  ConvertGpuOpsToNVVMOpsOptions gpuToNVVMOptions;

pyproject.toml

@@ -4,6 +4,7 @@ version = "0.1.0"
 requires-python = ">=3.11"
 dependencies = [
  "lit>=18.1.0",
+ "numpy",
  "pytest",
  "vastai-sdk",
 ]

test/Pipeline/attention.forth

@@ -0,0 +1,80 @@
+\ RUN: %warpforth-translate --forth-to-mlir %s | %warpforth-opt --warpforth-pipeline | %FileCheck %s
+
+\ Verify that a naive attention kernel with shared memory and float intrinsics
+\ survives the full pipeline to gpu.binary.
+\ CHECK: gpu.binary @warpforth_module
+
+\! kernel attention
+\! param Q f64[16]
+\! param K f64[16]
+\! param V f64[16]
+\! param O f64[16]
+\! param SEQ_LEN i64
+\! param HEAD_DIM i64
+\! shared SCORES f64[4]
+\! shared SCRATCH f64[4]
+
+\ row = BID-X, t = TID-X
+BID-X
+TID-X
+
+\ --- Dot product: Q[row,:] . K[t,:] ---
+0.0
+HEAD_DIM 0 DO
+ 2 PICK HEAD_DIM * I + CELLS Q + F@
+ 2 PICK HEAD_DIM * I + CELLS K + F@
+ F* F+
+LOOP
+HEAD_DIM S>F FSQRT F/
+
+\ --- Causal mask: if t > row, score = -inf ---
+OVER 3 PICK >
+IF DROP -1.0e30 THEN
+
+\ --- Store score to shared memory ---
+OVER CELLS SCORES + SF!
+BARRIER
+
+\ --- Softmax: max reduction (thread 0) ---
+TID-X 0= IF
+ 0 CELLS SCORES + SF@
+ SEQ_LEN 1 DO I CELLS SCORES + SF@ FMAX LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+
+\ --- Softmax: exp(score - max) ---
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F- FEXP
+OVER CELLS SCORES + SF!
+BARRIER
+
+\ --- Softmax: sum reduction (thread 0) ---
+TID-X 0= IF
+ 0.0
+ SEQ_LEN 0 DO I CELLS SCORES + SF@ F+ LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+
+\ --- Softmax: normalize ---
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F/
+OVER CELLS SCORES + SF!
+BARRIER
+
+\ --- V accumulation: O[row,col] = sum_j SCORES[j] * V[j*HD + col] ---
+\ Stride over head_dim columns: col = t, t+BDIM-X, t+2*BDIM-X, ...
+DUP BEGIN DUP HEAD_DIM < WHILE
+ 0.0
+ SEQ_LEN 0 DO
+ I CELLS SCORES + SF@
+ I HEAD_DIM * 3 PICK + CELLS V + F@
+ F* F+
+ LOOP
+ OVER 4 PICK HEAD_DIM * + CELLS O + F!
+ BDIM-X +
+REPEAT
+DROP DROP DROP

gpu_test/conftest.py

@@ -118,7 +118,10 @@ def _cleanup_orphans(self) -> None:
 
  def _launch(self) -> None:
  """Find the cheapest suitable offer and launch an instance."""
- query = f"num_gpus=1 rentable=True rented=False compute_cap>=700 dph<={MAX_COST_PER_HOUR}"
+ query = (
+ f"num_gpus=1 rentable=True rented=False compute_cap>=700"
+ f" reliability2>=0.95 inet_up>=100 dph<={MAX_COST_PER_HOUR}"
+ )
  offers = self.sdk.search_offers(query=query, order="dph", limit=5)
 
  if not offers:

gpu_test/test_kernels.py

@@ -4,6 +4,7 @@
 
 from typing import TYPE_CHECKING
 
+import numpy as np
 import pytest
 
 if TYPE_CHECKING:
@@ -583,3 +584,157 @@ def test_float_to_int_conversion(kernel_runner: KernelRunner) -> None:
  forth_source=("\\! kernel main\n\\! param DATA i64[256]\n7.9 F>S\n0 CELLS DATA + !"),
  )
  assert result[0] == 7
+
+
+# --- Attention ---
+
+_ATTENTION_KERNEL = """\
+\\! kernel attention
+\\! param Q f64[{n}]
+\\! param K f64[{n}]
+\\! param V f64[{n}]
+\\! param O f64[{n}]
+\\! param SEQ_LEN i64
+\\! param HEAD_DIM i64
+\\! shared SCORES f64[{seq_len}]
+\\! shared SCRATCH f64[{seq_len}]
+BID-X
+TID-X
+0.0
+HEAD_DIM 0 DO
+ 2 PICK HEAD_DIM * I + CELLS Q + F@
+ 2 PICK HEAD_DIM * I + CELLS K + F@
+ F* F+
+LOOP
+HEAD_DIM S>F FSQRT F/
+OVER 3 PICK >
+IF DROP -1.0e30 THEN
+OVER CELLS SCORES + SF!
+BARRIER
+TID-X 0= IF
+ 0 CELLS SCORES + SF@
+ SEQ_LEN 1 DO I CELLS SCORES + SF@ FMAX LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F- FEXP
+OVER CELLS SCORES + SF!
+BARRIER
+TID-X 0= IF
+ 0.0
+ SEQ_LEN 0 DO I CELLS SCORES + SF@ F+ LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F/
+OVER CELLS SCORES + SF!
+BARRIER
+DUP BEGIN DUP HEAD_DIM < WHILE
+ 0.0
+ SEQ_LEN 0 DO
+ I CELLS SCORES + SF@
+ I HEAD_DIM * 3 PICK + CELLS V + F@
+ F* F+
+ LOOP
+ OVER 4 PICK HEAD_DIM * + CELLS O + F!
+ BDIM-X +
+REPEAT
+DROP DROP DROP
+"""
+
+
+def _attention_reference(q: np.ndarray, k: np.ndarray, v: np.ndarray, seq_len: int) -> list[float]:
+ """Compute scaled dot-product attention with causal mask (NumPy reference)."""
+ head_dim = q.shape[1]
+ scores = q @ k.T / np.sqrt(head_dim)
+ causal_mask = np.triu(np.ones((seq_len, seq_len), dtype=bool), k=1)
+ scores[causal_mask] = -1e30
+ exp_scores = np.exp(scores - scores.max(axis=1, keepdims=True))
+ attn = exp_scores / exp_scores.sum(axis=1, keepdims=True)
+ return (attn @ v).flatten().tolist()
+
+
+def test_naive_attention_f64(kernel_runner: KernelRunner) -> None:
+ """Naive scaled dot-product attention with causal mask.
+
+ O = softmax(Q @ K^T / sqrt(d_k)) @ V, seq_len=4, head_dim=4.
+ One block per query row, one thread per key position.
+ """
+ seq_len, head_dim = 4, 4
+
+ q = np.array(
+ [
+ [1.0, 0.0, 1.0, 0.0],
+ [0.0, 1.0, 0.0, 1.0],
+ [1.0, 1.0, 0.0, 0.0],
+ [0.0, 0.0, 1.0, 1.0],
+ ]
+ )
+ k = np.array(
+ [
+ [1.0, 0.0, 0.0, 1.0],
+ [0.0, 1.0, 1.0, 0.0],
+ [1.0, 1.0, 0.0, 0.0],
+ [0.0, 0.0, 1.0, 1.0],
+ ]
+ )
+ v = np.array(
+ [
+ [1.0, 2.0, 3.0, 4.0],
+ [5.0, 6.0, 7.0, 8.0],
+ [9.0, 10.0, 11.0, 12.0],
+ [13.0, 14.0, 15.0, 16.0],
+ ]
+ )
+
+ expected = _attention_reference(q, k, v, seq_len)
+ n = seq_len * head_dim
+
+ result = kernel_runner.run(
+ forth_source=_ATTENTION_KERNEL.format(n=n, seq_len=seq_len),
+ params={
+ "Q": q.flatten().tolist(),
+ "K": k.flatten().tolist(),
+ "V": v.flatten().tolist(),
+ "SEQ_LEN": seq_len,
+ "HEAD_DIM": head_dim,
+ },
+ grid=(seq_len, 1, 1),
+ block=(seq_len, 1, 1),
+ output_param=3,
+ output_count=n,
+ )
+ assert result == [pytest.approx(v) for v in expected]
+
+
+def test_naive_attention_f64_16x64(kernel_runner: KernelRunner) -> None:
+ """Naive scaled dot-product attention, seq_len=16, head_dim=64."""
+ seq_len, head_dim = 16, 64
+
+ rng = np.random.default_rng(42)
+ q = rng.standard_normal((seq_len, head_dim))
+ k = rng.standard_normal((seq_len, head_dim))
+ v = rng.standard_normal((seq_len, head_dim))
+
+ expected = _attention_reference(q, k, v, seq_len)
+ n = seq_len * head_dim
+
+ result = kernel_runner.run(
+ forth_source=_ATTENTION_KERNEL.format(n=n, seq_len=seq_len),
+ params={
+ "Q": q.flatten().tolist(),
+ "K": k.flatten().tolist(),
+ "V": v.flatten().tolist(),
+ "SEQ_LEN": seq_len,
+ "HEAD_DIM": head_dim,
+ },
+ grid=(seq_len, 1, 1),
+ block=(seq_len, 1, 1),
+ output_param=3,
+ output_count=n,
+ )
+ assert result == [pytest.approx(v) for v in expected]

lib/Bitcode/CMakeLists.txt

@@ -0,0 +1,5 @@
+set(WARPFORTH_LIBDEVICE_PATH
+ "${CMAKE_CURRENT_BINARY_DIR}/libdevice.10.bc" PARENT_SCOPE)
+
+configure_file(libdevice.10.bc
+ "${CMAKE_CURRENT_BINARY_DIR}/libdevice.10.bc" COPYONLY)

lib/CMakeLists.txt

@@ -1,3 +1,4 @@
 add_subdirectory(Dialect)
+add_subdirectory(Bitcode)
 add_subdirectory(Conversion)
 add_subdirectory(Translation)

lib/Conversion/CMakeLists.txt

@@ -16,5 +16,9 @@ add_mlir_library(MLIRConversionPasses
  MLIRTransforms
 )
 
+target_compile_definitions(obj.MLIRConversionPasses PRIVATE
+ WARPFORTH_LIBDEVICE_PATH="${WARPFORTH_LIBDEVICE_PATH}"
+)
+
 add_subdirectory(ForthToMemRef)
 add_subdirectory(ForthToGPU)

lib/Conversion/Passes.cpp

@@ -33,7 +33,10 @@ void buildWarpForthPipeline(OpPassManager &pm) {
  pm.addPass(createCanonicalizerPass());
 
  // Stage 4: Attach NVVM target to GPU modules (sm_70 = Volta architecture)
- pm.addPass(createGpuNVVMAttachTarget());
+ GpuNVVMAttachTargetOptions nvvmOptions;
+ nvvmOptions.chip = "sm_70";
+ nvvmOptions.linkLibs.push_back(WARPFORTH_LIBDEVICE_PATH);
+ pm.addPass(createGpuNVVMAttachTarget(nvvmOptions));
 
  // Stage 5: Lower GPU to NVVM with bare pointers
  ConvertGpuOpsToNVVMOpsOptions gpuToNVVMOptions;

pyproject.toml

@@ -4,6 +4,7 @@ version = "0.1.0"
 requires-python = ">=3.11"
 dependencies = [
  "lit>=18.1.0",
+ "numpy",
  "pytest",
  "vastai-sdk",
 ]

test/Pipeline/attention.forth

@@ -0,0 +1,80 @@
+\ RUN: %warpforth-translate --forth-to-mlir %s | %warpforth-opt --warpforth-pipeline | %FileCheck %s
+
+\ Verify that a naive attention kernel with shared memory and float intrinsics
+\ survives the full pipeline to gpu.binary.
+\ CHECK: gpu.binary @warpforth_module
+
+\! kernel attention
+\! param Q f64[16]
+\! param K f64[16]
+\! param V f64[16]
+\! param O f64[16]
+\! param SEQ_LEN i64
+\! param HEAD_DIM i64
+\! shared SCORES f64[4]
+\! shared SCRATCH f64[4]
+
+\ row = BID-X, t = TID-X
+BID-X
+TID-X
+
+\ --- Dot product: Q[row,:] . K[t,:] ---
+0.0
+HEAD_DIM 0 DO
+ 2 PICK HEAD_DIM * I + CELLS Q + F@
+ 2 PICK HEAD_DIM * I + CELLS K + F@
+ F* F+
+LOOP
+HEAD_DIM S>F FSQRT F/
+
+\ --- Causal mask: if t > row, score = -inf ---
+OVER 3 PICK >
+IF DROP -1.0e30 THEN
+
+\ --- Store score to shared memory ---
+OVER CELLS SCORES + SF!
+BARRIER
+
+\ --- Softmax: max reduction (thread 0) ---
+TID-X 0= IF
+ 0 CELLS SCORES + SF@
+ SEQ_LEN 1 DO I CELLS SCORES + SF@ FMAX LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+
+\ --- Softmax: exp(score - max) ---
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F- FEXP
+OVER CELLS SCORES + SF!
+BARRIER
+
+\ --- Softmax: sum reduction (thread 0) ---
+TID-X 0= IF
+ 0.0
+ SEQ_LEN 0 DO I CELLS SCORES + SF@ F+ LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+
+\ --- Softmax: normalize ---
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F/
+OVER CELLS SCORES + SF!
+BARRIER
+
+\ --- V accumulation: O[row,col] = sum_j SCORES[j] * V[j*HD + col] ---
+\ Stride over head_dim columns: col = t, t+BDIM-X, t+2*BDIM-X, ...
+DUP BEGIN DUP HEAD_DIM < WHILE
+ 0.0
+ SEQ_LEN 0 DO
+ I CELLS SCORES + SF@
+ I HEAD_DIM * 3 PICK + CELLS V + F@
+ F* F+
+ LOOP
+ OVER 4 PICK HEAD_DIM * + CELLS O + F!
+ BDIM-X +
+REPEAT
+DROP DROP DROP

gpu_test/conftest.py

@@ -118,7 +118,10 @@ def _cleanup_orphans(self) -> None:
 
  def _launch(self) -> None:
  """Find the cheapest suitable offer and launch an instance."""
- query = f"num_gpus=1 rentable=True rented=False compute_cap>=700 dph<={MAX_COST_PER_HOUR}"
+ query = (
+ f"num_gpus=1 rentable=True rented=False compute_cap>=700"
+ f" reliability2>=0.95 inet_up>=100 dph<={MAX_COST_PER_HOUR}"
+ )
  offers = self.sdk.search_offers(query=query, order="dph", limit=5)
 
  if not offers:

gpu_test/test_kernels.py

@@ -4,6 +4,7 @@
 
 from typing import TYPE_CHECKING
 
+import numpy as np
 import pytest
 
 if TYPE_CHECKING:
@@ -583,3 +584,157 @@ def test_float_to_int_conversion(kernel_runner: KernelRunner) -> None:
  forth_source=("\\! kernel main\n\\! param DATA i64[256]\n7.9 F>S\n0 CELLS DATA + !"),
  )
  assert result[0] == 7
+
+
+# --- Attention ---
+
+_ATTENTION_KERNEL = """\
+\\! kernel attention
+\\! param Q f64[{n}]
+\\! param K f64[{n}]
+\\! param V f64[{n}]
+\\! param O f64[{n}]
+\\! param SEQ_LEN i64
+\\! param HEAD_DIM i64
+\\! shared SCORES f64[{seq_len}]
+\\! shared SCRATCH f64[{seq_len}]
+BID-X
+TID-X
+0.0
+HEAD_DIM 0 DO
+ 2 PICK HEAD_DIM * I + CELLS Q + F@
+ 2 PICK HEAD_DIM * I + CELLS K + F@
+ F* F+
+LOOP
+HEAD_DIM S>F FSQRT F/
+OVER 3 PICK >
+IF DROP -1.0e30 THEN
+OVER CELLS SCORES + SF!
+BARRIER
+TID-X 0= IF
+ 0 CELLS SCORES + SF@
+ SEQ_LEN 1 DO I CELLS SCORES + SF@ FMAX LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F- FEXP
+OVER CELLS SCORES + SF!
+BARRIER
+TID-X 0= IF
+ 0.0
+ SEQ_LEN 0 DO I CELLS SCORES + SF@ F+ LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F/
+OVER CELLS SCORES + SF!
+BARRIER
+DUP BEGIN DUP HEAD_DIM < WHILE
+ 0.0
+ SEQ_LEN 0 DO
+ I CELLS SCORES + SF@
+ I HEAD_DIM * 3 PICK + CELLS V + F@
+ F* F+
+ LOOP
+ OVER 4 PICK HEAD_DIM * + CELLS O + F!
+ BDIM-X +
+REPEAT
+DROP DROP DROP
+"""
+
+
+def _attention_reference(q: np.ndarray, k: np.ndarray, v: np.ndarray, seq_len: int) -> list[float]:
+ """Compute scaled dot-product attention with causal mask (NumPy reference)."""
+ head_dim = q.shape[1]
+ scores = q @ k.T / np.sqrt(head_dim)
+ causal_mask = np.triu(np.ones((seq_len, seq_len), dtype=bool), k=1)
+ scores[causal_mask] = -1e30
+ exp_scores = np.exp(scores - scores.max(axis=1, keepdims=True))
+ attn = exp_scores / exp_scores.sum(axis=1, keepdims=True)
+ return (attn @ v).flatten().tolist()
+
+
+def test_naive_attention_f64(kernel_runner: KernelRunner) -> None:
+ """Naive scaled dot-product attention with causal mask.
+
+ O = softmax(Q @ K^T / sqrt(d_k)) @ V, seq_len=4, head_dim=4.
+ One block per query row, one thread per key position.
+ """
+ seq_len, head_dim = 4, 4
+
+ q = np.array(
+ [
+ [1.0, 0.0, 1.0, 0.0],
+ [0.0, 1.0, 0.0, 1.0],
+ [1.0, 1.0, 0.0, 0.0],
+ [0.0, 0.0, 1.0, 1.0],
+ ]
+ )
+ k = np.array(
+ [
+ [1.0, 0.0, 0.0, 1.0],
+ [0.0, 1.0, 1.0, 0.0],
+ [1.0, 1.0, 0.0, 0.0],
+ [0.0, 0.0, 1.0, 1.0],
+ ]
+ )
+ v = np.array(
+ [
+ [1.0, 2.0, 3.0, 4.0],
+ [5.0, 6.0, 7.0, 8.0],
+ [9.0, 10.0, 11.0, 12.0],
+ [13.0, 14.0, 15.0, 16.0],
+ ]
+ )
+
+ expected = _attention_reference(q, k, v, seq_len)
+ n = seq_len * head_dim
+
+ result = kernel_runner.run(
+ forth_source=_ATTENTION_KERNEL.format(n=n, seq_len=seq_len),
+ params={
+ "Q": q.flatten().tolist(),
+ "K": k.flatten().tolist(),
+ "V": v.flatten().tolist(),
+ "SEQ_LEN": seq_len,
+ "HEAD_DIM": head_dim,
+ },
+ grid=(seq_len, 1, 1),
+ block=(seq_len, 1, 1),
+ output_param=3,
+ output_count=n,
+ )
+ assert result == [pytest.approx(v) for v in expected]
+
+
+def test_naive_attention_f64_16x64(kernel_runner: KernelRunner) -> None:
+ """Naive scaled dot-product attention, seq_len=16, head_dim=64."""
+ seq_len, head_dim = 16, 64
+
+ rng = np.random.default_rng(42)
+ q = rng.standard_normal((seq_len, head_dim))
+ k = rng.standard_normal((seq_len, head_dim))
+ v = rng.standard_normal((seq_len, head_dim))
+
+ expected = _attention_reference(q, k, v, seq_len)
+ n = seq_len * head_dim
+
+ result = kernel_runner.run(
+ forth_source=_ATTENTION_KERNEL.format(n=n, seq_len=seq_len),
+ params={
+ "Q": q.flatten().tolist(),
+ "K": k.flatten().tolist(),
+ "V": v.flatten().tolist(),
+ "SEQ_LEN": seq_len,
+ "HEAD_DIM": head_dim,
+ },
+ grid=(seq_len, 1, 1),
+ block=(seq_len, 1, 1),
+ output_param=3,
+ output_count=n,
+ )
+ assert result == [pytest.approx(v) for v in expected]

lib/Bitcode/CMakeLists.txt

@@ -0,0 +1,5 @@
+set(WARPFORTH_LIBDEVICE_PATH
+ "${CMAKE_CURRENT_BINARY_DIR}/libdevice.10.bc" PARENT_SCOPE)
+
+configure_file(libdevice.10.bc
+ "${CMAKE_CURRENT_BINARY_DIR}/libdevice.10.bc" COPYONLY)

lib/CMakeLists.txt

@@ -1,3 +1,4 @@
 add_subdirectory(Dialect)
+add_subdirectory(Bitcode)
 add_subdirectory(Conversion)
 add_subdirectory(Translation)

lib/Conversion/CMakeLists.txt

@@ -16,5 +16,9 @@ add_mlir_library(MLIRConversionPasses
  MLIRTransforms
 )
 
+target_compile_definitions(obj.MLIRConversionPasses PRIVATE
+ WARPFORTH_LIBDEVICE_PATH="${WARPFORTH_LIBDEVICE_PATH}"
+)
+
 add_subdirectory(ForthToMemRef)
 add_subdirectory(ForthToGPU)

lib/Conversion/Passes.cpp

@@ -33,7 +33,10 @@ void buildWarpForthPipeline(OpPassManager &pm) {
  pm.addPass(createCanonicalizerPass());
 
  // Stage 4: Attach NVVM target to GPU modules (sm_70 = Volta architecture)
- pm.addPass(createGpuNVVMAttachTarget());
+ GpuNVVMAttachTargetOptions nvvmOptions;
+ nvvmOptions.chip = "sm_70";
+ nvvmOptions.linkLibs.push_back(WARPFORTH_LIBDEVICE_PATH);
+ pm.addPass(createGpuNVVMAttachTarget(nvvmOptions));
 
  // Stage 5: Lower GPU to NVVM with bare pointers
  ConvertGpuOpsToNVVMOpsOptions gpuToNVVMOptions;

pyproject.toml

@@ -4,6 +4,7 @@ version = "0.1.0"
 requires-python = ">=3.11"
 dependencies = [
  "lit>=18.1.0",
+ "numpy",
  "pytest",
  "vastai-sdk",
 ]

test/Pipeline/attention.forth

@@ -0,0 +1,80 @@
+\ RUN: %warpforth-translate --forth-to-mlir %s | %warpforth-opt --warpforth-pipeline | %FileCheck %s
+
+\ Verify that a naive attention kernel with shared memory and float intrinsics
+\ survives the full pipeline to gpu.binary.
+\ CHECK: gpu.binary @warpforth_module
+
+\! kernel attention
+\! param Q f64[16]
+\! param K f64[16]
+\! param V f64[16]
+\! param O f64[16]
+\! param SEQ_LEN i64
+\! param HEAD_DIM i64
+\! shared SCORES f64[4]
+\! shared SCRATCH f64[4]
+
+\ row = BID-X, t = TID-X
+BID-X
+TID-X
+
+\ --- Dot product: Q[row,:] . K[t,:] ---
+0.0
+HEAD_DIM 0 DO
+ 2 PICK HEAD_DIM * I + CELLS Q + F@
+ 2 PICK HEAD_DIM * I + CELLS K + F@
+ F* F+
+LOOP
+HEAD_DIM S>F FSQRT F/
+
+\ --- Causal mask: if t > row, score = -inf ---
+OVER 3 PICK >
+IF DROP -1.0e30 THEN
+
+\ --- Store score to shared memory ---
+OVER CELLS SCORES + SF!
+BARRIER
+
+\ --- Softmax: max reduction (thread 0) ---
+TID-X 0= IF
+ 0 CELLS SCORES + SF@
+ SEQ_LEN 1 DO I CELLS SCORES + SF@ FMAX LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+
+\ --- Softmax: exp(score - max) ---
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F- FEXP
+OVER CELLS SCORES + SF!
+BARRIER
+
+\ --- Softmax: sum reduction (thread 0) ---
+TID-X 0= IF
+ 0.0
+ SEQ_LEN 0 DO I CELLS SCORES + SF@ F+ LOOP
+ 0 CELLS SCRATCH + SF!
+THEN
+BARRIER
+
+\ --- Softmax: normalize ---
+DUP CELLS SCORES + SF@
+0 CELLS SCRATCH + SF@
+F/
+OVER CELLS SCORES + SF!
+BARRIER
+
+\ --- V accumulation: O[row,col] = sum_j SCORES[j] * V[j*HD + col] ---
+\ Stride over head_dim columns: col = t, t+BDIM-X, t+2*BDIM-X, ...
+DUP BEGIN DUP HEAD_DIM < WHILE
+ 0.0
+ SEQ_LEN 0 DO
+ I CELLS SCORES + SF@
+ I HEAD_DIM * 3 PICK + CELLS V + F@
+ F* F+
+ LOOP
+ OVER 4 PICK HEAD_DIM * + CELLS O + F!
+ BDIM-X +
+REPEAT
+DROP DROP DROP